Update bot.py

bot.py

@@ -11,6 +11,14 @@ def get_admins():
 def is_admin(event):
     return event.sender_id in get_admins()
 
+def get_main_buttons():
+    return [
+        [Button.inline("▶️ Start Scan", b"start_scan"), Button.inline("⏸ Pause", b"pause_scan")],
+        [Button.inline("⏹ Stop", b"stop_scan"), Button.inline("📊 Status", b"show_status")],
+        [Button.inline("⚡ Set Speed", b"set_speed"), Button.inline("📥 Export", b"export_files")],
+        [Button.inline("🔄 Load Words", b"load_words"), Button.inline("💣 Reset", b"reset_confirm")]
+    ]
+
 async def generate_status_msg():
     state = await db.get_state()
     counts = await db.get_counts()
@@ -55,13 +63,7 @@ def setup_handlers(client):
             await event.respond("⛔ Unauthorized.")
             return
         
-        buttons = [
-            [Button.inline("▶️ Start Scan", b"start_scan"), Button.inline("⏸ Pause", b"pause_scan")],
-            [Button.inline("⏹ Stop", b"stop_scan"), Button.inline("📊 Status", b"show_status")],
-            [Button.inline("⚡ Set Speed", b"set_speed"), Button.inline("📥 Export", b"export_files")],
-            [Button.inline("🔄 Load Words", b"load_words"), Button.inline("💣 Reset", b"reset_confirm")]
-        ]
-        await event.respond("👋 Fragment Scanner Bot", buttons=buttons)
+        await event.respond("👋 **Fragment Scanner Bot**\nSelect an action below:", buttons=get_main_buttons())
 
     @client.on(events.NewMessage(pattern='/upload'))
     async def upload_cmd(event):
@@ -80,28 +82,32 @@ def setup_handlers(client):
             return
 
         msg = await event.respond("📥 Downloading massive file to server disk...")
-        
         await reply_msg.download_media(file="words.txt")
-        
-        await msg.edit("✅ **File uploaded to server successfully!**\n\nClick **🔄 Load Words** to bite off the first 200,000 words into the database.")
+        await msg.edit("✅ **File uploaded to server successfully!**\n\nClick **🔄 Load Words** to bite off the first 200,000 words into the database.", buttons=get_main_buttons())
 
     @client.on(events.NewMessage(pattern='/load'))
     async def load_cmd(event):
         if not is_admin(event):
             await event.respond("⛔ Unauthorized.")
             return
-        await _load_words(event.respond)
+        await _load_words(event, is_callback=False)
+
+    async def _load_words(event, is_callback=True):
+        async def update_msg(text):
+            if is_callback:
+                await event.edit(text, buttons=get_main_buttons())
+            else:
+                await event.respond(text, buttons=get_main_buttons())
 
-    async def _load_words(reply_func):
         try:
             with open("words.txt", "r", encoding="utf-8") as f:
                 all_lines = [line.strip().lower() for line in f if line.strip()]
         except Exception:
-            await reply_func("❌ words.txt not found. Please upload your master list using `/upload`.")
+            await update_msg("❌ `words.txt` not found. Please upload your master list using `/upload`.")
             return
 
         if not all_lines:
-            await reply_func("⚠️ The `words.txt` file on the server is empty. You've processed everything! Upload a new file.")
+            await update_msg("⚠️ The `words.txt` file on the server is empty. You've processed everything! Upload a new file.")
             return
 
         # --- SERVER DISK BUFFER LOGIC ---
@@ -126,12 +132,11 @@ def setup_handlers(client):
             except Exception as e:
                 err_str = str(e).lower()
                 if "oom" in err_str or "maxmemory" in err_str:
-                    await reply_func("❌ **Error:** Redis ran out of memory! Click 💣 Reset, then try loading again.")
+                    await update_msg("❌ **Error:** Redis ran out of memory! Click **💣 Reset**, then try loading again.")
                 else:
-                    await reply_func(f"❌ **Database Error:** {str(e)}")
-                return # Abort before modifying the text file so no data is lost
+                    await update_msg(f"❌ **Database Error:** {str(e)}")
+                return
         
-        # If successfully pushed to Redis, overwrite the file with the remaining lines
         with open("words.txt", "w", encoding="utf-8") as f:
             f.write("\n".join(remaining_lines))
             
@@ -144,104 +149,103 @@ def setup_handlers(client):
         else:
             msg += "✅ **No words remaining on server disk.**"
             
-        await reply_func(msg)
+        await update_msg(msg)
 
     @client.on(events.CallbackQuery())
     async def button_handler(event):
         if not is_admin(event):
-            await event.answer("⛔ Unauthorized.", alert=True)
+            try:
+                await event.answer("⛔ Unauthorized.", alert=True)
+            except Exception:
+                pass
             return
 
+        # --- CRITICAL FIX: Answer immediately to prevent QueryIdInvalidError timeouts ---
+        try:
+            await event.answer()
+        except Exception:
+            pass
+
         data = event.data.decode('utf-8')
 
-        if data == "load_words":
-            await event.answer()
-            await _load_words(event.reply)
+        if data == "back_main":
+            await event.edit("👋 **Fragment Scanner Bot**\nSelect an action below:", buttons=get_main_buttons())
+
+        elif data == "load_words":
+            await event.edit("⏳ Reading and loading words... Please wait.")
+            await _load_words(event, is_callback=True)
             
         elif data == "start_scan":
             await db.set_state(running="1", paused="0", start_time=str(time.time()))
             qlen = await db.queue_size()
-            await event.answer()
-            await event.reply(f"🚀 Scan started! {qlen} words in queue.")
+            await event.edit(f"🚀 **Scan started!**\n`{qlen:,}` words currently in the processing queue.", buttons=get_main_buttons())
             
         elif data == "pause_scan":
             state = await db.get_state()
             if state.get("paused") == "1":
                 await db.set_state(paused="0")
-                await event.answer()
-                await event.reply("▶️ Scan resumed.")
+                await event.edit("▶️ **Scan resumed.**", buttons=get_main_buttons())
             else:
                 await db.set_state(paused="1")
-                await event.answer()
-                await event.reply("⏸ Scan paused.")
+                await event.edit("⏸ **Scan paused.**", buttons=get_main_buttons())
                 
         elif data == "stop_scan":
             await db.set_state(running="0", paused="0")
-            await event.answer()
-            await event.reply("⏹ Scan stopped. Progress saved.")
+            await event.edit("⏹ **Scan stopped.** Progress has been saved.", buttons=get_main_buttons())
             
-        elif data == "show_status":
-            msg = await generate_status_msg()
-            await event.answer()
-            await event.reply(msg, buttons=[[Button.inline("🔄 Refresh Live", b"refresh_status")]])
-            
-        elif data == "refresh_status":
+        elif data in ("show_status", "refresh_status"):
             msg = await generate_status_msg()
+            buttons = [
+                [Button.inline("🔄 Refresh Live", b"refresh_status")],
+                [Button.inline("🔙 Back to Menu", b"back_main")]
+            ]
             try:
-                await event.edit(msg, buttons=[[Button.inline("🔄 Refresh Live", b"refresh_status")]])
-                await event.answer("🔄 Status Updated!")
+                await event.edit(msg, buttons=buttons)
             except Exception as e:
-                if "not modified" in str(e).lower():
-                    await event.answer("⚠️ Already up to date.", alert=False)
-                else:
-                    await event.answer("Error updating status.", alert=True)
+                pass # Ignores "Message is not modified" errors on spam click
             
         elif data == "set_speed":
             buttons = [
                 [Button.inline("🐢 10", b"spd_10"), Button.inline("🚶 20", b"spd_20"), Button.inline("🚗 30", b"spd_30")],
-                [Button.inline("🚀 50", b"spd_50"), Button.inline("⚡ 75", b"spd_75"), Button.inline("🔥 100", b"spd_100")]
+                [Button.inline("🚀 50", b"spd_50"), Button.inline("⚡ 75", b"spd_75"), Button.inline("🔥 100", b"spd_100")],
+                [Button.inline("🔙 Cancel", b"back_main")]
             ]
-            await event.answer()
-            await event.reply("Select Speed:", buttons=buttons)
+            await event.edit("⚡ **Select Scanner Speed (Concurrent Workers):**", buttons=buttons)
             
         elif data.startswith("spd_"):
             try:
                 speed_limit = int(data.split("_")[1])
                 await db.set_concurrency(speed_limit)
-                await event.answer()
-                await event.reply(f"⚡ Speed set to {speed_limit} concurrent workers.")
-            except (IndexError, ValueError):
-                await event.answer("❌ Invalid speed value.", alert=True)
-            except Exception as e:
-                await event.answer("❌ Error setting speed.", alert=True)
+                await event.edit(f"⚡ **Speed updated successfully!**\nNow running with `{speed_limit}` concurrent workers.", buttons=get_main_buttons())
+            except Exception:
+                await event.edit("❌ Error setting speed.", buttons=get_main_buttons())
             
         elif data == "export_files":
-            await event.answer()
+            await event.edit("⏳ **Exporting data...**\nPlease wait while files are generated.", buttons=get_main_buttons())
+            
             taken = sorted(await db.get_all_taken())
             unavail = sorted(await db.get_all_unavailable())
             
-            f_taken = io.BytesIO("\n".join(taken).encode('utf-8'))
-            f_taken.name = "taken.txt"
-            await client.send_file(event.chat_id, file=f_taken)
+            if taken:
+                f_taken = io.BytesIO("\n".join(taken).encode('utf-8'))
+                f_taken.name = "taken.txt"
+                await client.send_file(event.chat_id, file=f_taken)
             
-            f_unavail = io.BytesIO("\n".join(unavail).encode('utf-8'))
-            f_unavail.name = "unavailable.txt"
-            await client.send_file(event.chat_id, file=f_unavail)
+            if unavail:
+                f_unavail = io.BytesIO("\n".join(unavail).encode('utf-8'))
+                f_unavail.name = "unavailable.txt"
+                await client.send_file(event.chat_id, file=f_unavail)
             
-            await event.reply(f"✅ Exported {len(taken):,} taken, {len(unavail):,} unavailable.")
+            await event.edit(f"✅ **Export Complete!**\nSent `{len(taken):,}` taken and `{len(unavail):,}` unavailable usernames.", buttons=get_main_buttons())
             
         elif data == "reset_confirm":
             buttons = [
-                [Button.inline("✅ Yes, WIPE ALL DATA", b"reset_do"), Button.inline("❌ Cancel", b"reset_cancel")]
+                [Button.inline("✅ Yes, WIPE ALL DATA", b"reset_do")],
+                [Button.inline("❌ Cancel", b"back_main")]
             ]
-            await event.answer()
-            await event.reply("⚠️ Are you sure you want to completely wipe the database?", buttons=buttons)
+            await event.edit("⚠️ **DANGER ZONE** ⚠️\n\nAre you sure you want to completely wipe the Redis database? All progress will be lost.", buttons=buttons)
             
         elif data == "reset_do":
+            await event.edit("⏳ **Wiping database...**")
             await db.flush_all()
-            await event.answer()
-            await event.reply("💣 Database wiped. Click **🔄 Load Words** to grab the next batch from the server disk!")
-            
-        elif data == "reset_cancel":
-            await event.answer()
-            await event.reply("Cancelled.")
+            await event.edit("💣 **Database wiped.**\n\nRedis is completely clear. Click **🔄 Load Words** to grab the next batch from the server disk!", buttons=get_main_buttons())