Update bot.py

bot.py

@@ -4,7 +4,9 @@ import time
 import io
 import asyncio
 from telethon import events, Button
+from PIL import Image, ImageDraw, ImageFont
 import db
+import scraper # Required for the /check command
 
 def get_admins():
     return {int(x) for x in os.environ.get("ADMIN_IDS", "").split(",") if x}
@@ -12,14 +14,75 @@ def get_admins():
 def is_admin(event):
     return event.sender_id in get_admins()
 
-def get_main_buttons():
+# --- PROFESSIONAL UI MENUS ---
+def get_main_menu():
     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")]
+        [Button.inline("Start Scan", b"start_scan"), Button.inline("Pause", b"pause_scan"), Button.inline("Stop", b"stop_scan")],
+        [Button.inline("Live Status", b"show_status"), Button.inline("Export Data", b"export_files")],
+        [Button.inline("System Settings", b"menu_settings")]
     ]
 
+def get_settings_menu():
+    return [
+        [Button.inline("Load Database", b"load_words")],
+        [Button.inline("Set Speed Limit", b"set_speed")],
+        [Button.inline("Wipe Database", b"reset_confirm")],
+        [Button.inline("« Back to Main Menu", b"back_main")]
+    ]
+
+# --- IMAGE GENERATOR (LIQUID GLASS STYLE) ---
+def generate_status_image(username, status):
+    # Base dark theme colors
+    bg_color = (15, 23, 42)       # Slate 900
+    card_color = (30, 41, 59)     # Slate 800 (Glass card)
+    outline_color = (51, 65, 85)  # Slate 700
+    text_primary = (248, 250, 252) # Slate 50
+    
+    # Determine status color
+    if status == "TAKEN":
+        status_color = (239, 68, 68)   # Red
+    elif status == "UNAVAILABLE":
+        status_color = (156, 163, 175) # Gray
+    elif status == "FOR_SALE" or status == "ON_AUCTION":
+        status_color = (59, 130, 246)  # Blue
+    else:
+        status_color = (34, 197, 94)   # Green (Available)
+
+    # Create Image
+    img = Image.new('RGB', (800, 400), color=bg_color)
+    d = ImageDraw.Draw(img)
+    
+    # Draw "Glass" Card (Rounded Rectangle)
+    d.rounded_rectangle([(50, 50), (750, 350)], radius=20, fill=card_color, outline=outline_color, width=3)
+    
+    # Status Indicator Dot
+    d.ellipse([(90, 90), (110, 110)], fill=status_color)
+    
+    # Text Setup (Using default font scaled up via ImageFont if custom isn't available)
+    try:
+        # Tries to load a standard Linux font, falls back to default
+        font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 60)
+        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 30)
+    except IOError:
+        font_large = ImageFont.load_default()
+        font_small = ImageFont.load_default()
+
+    # Draw Text
+    d.text((130, 75), "FRAGMENT STATUS REPORT", fill=outline_color, font=font_small)
+    d.text((90, 160), f"@{username}", fill=text_primary, font=font_large)
+    
+    # Draw Status Pill
+    d.rounded_rectangle([(90, 260), (350, 310)], radius=25, fill=bg_color, outline=status_color, width=2)
+    d.text((120, 265), status, fill=status_color, font=font_small)
+
+    # Save to buffer
+    buf = io.BytesIO()
+    img.save(buf, format='PNG')
+    buf.seek(0)
+    buf.name = f"{username}_status.png"
+    return buf
+
+# --- MODERN STATUS DASHBOARD ---
 async def generate_status_msg():
     state = await db.get_state()
     counts = await db.get_counts()
@@ -29,243 +92,220 @@ async def generate_status_msg():
     total = int(state.get("total", 0))
     processed = int(state.get("processed", 0))
     pct = (processed / total * 100) if total > 0 else 0
-    run_str = "YES" if state.get("running") == "1" else "NO"
-    pause_str = "YES" if state.get("paused") == "1" else "NO"
     
-    bar_length = 15
+    run_str = "RUNNING" if state.get("running") == "1" else "STOPPED"
+    if state.get("paused") == "1": run_str = "PAUSED"
+    
+    bar_length = 20
     filled = int(bar_length * (processed / total)) if total > 0 else 0
-    bar = "█" * filled + "░" * (bar_length - filled)
+    bar = "â–ˆ" * filled + "â–‘" * (bar_length - filled)
     
     msg = (
-        f"**📊 SCANNER STATUS**\n"
-        f"━━━━━━━━━━━━━━━━━━━\n"
-        f"🟢 Running  : `{run_str:<10}`\n"
-        f"⏸ Paused   : `{pause_str:<10}`\n"
-        f"📦 Total    : `{total:<10,}`\n"
-        f"✅ Processed : `{processed:<10,}`\n"
-        f"📈 Progress : `{pct:.2f}%`\n"
+        f"**SYSTEM DASHBOARD**\n"
+        f"────────────────────────\n"
+        f"**State:** `{run_str}`\n"
+        f"**Queue:** `{qlen:,}` waiting\n"
+        f"**Speed:** `{concurrency}` workers\n"
+        f"**Progress:** `{pct:.2f}%`\n"
         f"`[{bar}]`\n"
-        f"━━━━━━━━━━━━━━━━━━━\n"
-        f"🔴 Taken    : `{counts['taken']:<10,}`\n"
-        f"🚫 Unavail  : `{counts['unavailable']:<10,}`\n"
-        f"💰 For Sale : `{counts['forsale']:<10,}`\n"
-        f"🔨 Auction  : `{counts['auction']:<10,}`\n"
-        f"🛒 Sold     : `{counts['sold']:<10,}`\n"
-        f"━━━━━━━━━━━━━━━━━━━\n"
-        f"⚡ Speed    : `{concurrency:<10}`\n"
-        f"📋 Queue    : `{qlen:<10,}`\n"
+        f"`{processed:,} / {total:,}`\n\n"
+        f"**DATABASE METRICS**\n"
+        f"────────────────────────\n"
+        f"Taken      : `{counts.get('taken', 0):,}`\n"
+        f"Unavailable: `{counts.get('unavailable', 0):,}`\n"
+        f"For Sale   : `{counts.get('forsale', 0):,}`\n"
+        f"Auction    : `{counts.get('auction', 0):,}`\n"
+        f"Sold       : `{counts.get('sold', 0):,}`\n"
     )
     return msg
 
+# --- BACKGROUND HELPERS ---
+def read_file_sync():
+    with open("words.txt", "r", encoding="utf-8") as f:
+        return [line.strip().lower() for line in f if line.strip()]
+
+def write_file_sync(lines):
+    with open("words.txt", "w", encoding="utf-8") as f:
+        f.write("\n".join(lines))
+
+def filter_words_sync(chunk):
+    return {w for w in chunk if re.match(r'^[a-z0-9_]{4,32}$', w)}
+
 def setup_handlers(client):
+    
     @client.on(events.NewMessage(pattern='/start'))
     async def start_cmd(event):
-        if not is_admin(event):
-            return
-        await event.respond("👋 **Fragment Scanner Bot**\nSelect an action below:", buttons=get_main_buttons())
+        if not is_admin(event): return
+        await event.respond("**Fragment Control System**\nSelect an operation:", buttons=get_main_menu())
+
+    # --- NEW MANUAL CHECK FEATURE ---
+    @client.on(events.NewMessage(pattern=r'/check\s+([a-zA-Z0-9_]+)'))
+    async def check_cmd(event):
+        if not is_admin(event): return
+        username = event.pattern_match.group(1).lower()
+        
+        msg = await event.respond(f"Scanning `{username}`...")
+        
+        # Scrape data
+        try:
+            status = await scraper.check_fragment(username)
+            if not status:
+                status = "UNKNOWN"
+        except Exception:
+            status = "ERROR"
+
+        # Generate sleek image
+        img_buffer = await asyncio.to_thread(generate_status_image, username, status)
+        
+        # Send result
+        await msg.delete()
+        await client.send_file(
+            event.chat_id, 
+            file=img_buffer, 
+            caption=f"**Target:** `@{username}`\n**Status:** `{status}`"
+        )
 
     @client.on(events.NewMessage(pattern='/upload'))
     async def upload_cmd(event):
-        if not is_admin(event):
-            return
+        if not is_admin(event): return
         if not event.is_reply:
-            await event.respond("⚠️ Please reply to a `.txt` file with `/upload`.")
+            await event.respond("Reply to a `.txt` file with `/upload`.")
             return
 
         reply_msg = await event.get_reply_message()
         if not reply_msg.document or reply_msg.file.ext != '.txt':
-            await event.respond("⚠️ The message you replied to must be a `.txt` document.")
+            await event.respond("Target must be a `.txt` document.")
             return
 
-        msg = await event.respond("📥 Downloading massive file to server disk...")
+        msg = await event.respond("Transferring to secure server storage...")
         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.", buttons=get_main_buttons())
-
-    @client.on(events.NewMessage(pattern='/load'))
-    async def load_cmd(event):
-        if not is_admin(event):
-            return
-        await _load_words(event, is_callback=False)
+        await msg.edit("Transfer complete. Open **System Settings** -> **Load Database** to initialize.", buttons=get_main_menu())
 
     async def _load_words(event, is_callback=True):
-        async def update_msg(text):
+        async def update_msg(text, buttons=None):
             if is_callback:
-                await event.edit(text, buttons=get_main_buttons())
+                await event.edit(text, buttons=buttons or get_settings_menu())
             else:
-                await event.respond(text, buttons=get_main_buttons())
+                await event.respond(text, buttons=buttons or get_settings_menu())
 
         try:
-            with open("words.txt", "r", encoding="utf-8") as f:
-                all_lines = [line.strip().lower() for line in f if line.strip()]
+            all_lines = await asyncio.to_thread(read_file_sync)
         except Exception:
-            await update_msg("❌ `words.txt` not found. Please upload your master list using `/upload`.")
+            await update_msg("`words.txt` missing. Please use `/upload`.")
             return
 
         if not all_lines:
-            await update_msg("⚠️ The `words.txt` file on the server is empty. You've processed everything! Upload a new file.")
+            await update_msg("Database is empty. Upload a new master list.")
             return
 
-        # Load chunks safely
         CHUNK_SIZE = 200000
         current_chunk = all_lines[:CHUNK_SIZE]
         remaining_lines = all_lines[CHUNK_SIZE:]
 
-        # OPTIMIZATION 1: Fast regex filtering
-        valid_words = {w for w in current_chunk if re.match(r'^[a-z0-9_]{4,32}$', w)}
-        
+        valid_words = await asyncio.to_thread(filter_words_sync, current_chunk)
         r = await db.get_redis()
         done = await r.smembers("frag:done")
-        
-        # OPTIMIZATION 2: Mathematical Set Subtraction (Instant calculation instead of freezing CPU)
         to_queue = list(valid_words - done)
         
         if to_queue:
             pipe = r.pipeline()
             for i in range(0, len(to_queue), 1000):
                 pipe.lpush("frag:queue", *to_queue[i:i+1000])
-                # OPTIMIZATION 3: Yield to event loop so bot buttons stay smooth
                 await asyncio.sleep(0) 
-            
             try:
                 await pipe.execute()
             except Exception as e:
-                err_str = str(e).lower()
-                if "oom" in err_str or "maxmemory" in err_str:
-                    await update_msg("❌ **Error:** Redis ran out of memory! Click **💣 Reset**, then try loading again.")
-                else:
-                    await update_msg(f"❌ **Database Error:** {str(e)}")
+                if "oom" in str(e).lower() or "maxmemory" in str(e).lower():
+                    await update_msg("Memory limit exceeded. Please Wipe Database first.")
                 return
         
-        with open("words.txt", "w", encoding="utf-8") as f:
-            f.write("\n".join(remaining_lines))
-            
+        await asyncio.to_thread(write_file_sync, remaining_lines)
         await db.set_state(total=len(valid_words))
         skipped = len(valid_words) - len(to_queue)
         
-        msg = f"📥 Loaded **{len(to_queue):,}** words into the queue. Skipped {skipped:,}.\n\n"
-        if remaining_lines:
-            msg += f"📦 **{len(remaining_lines):,} words remaining on server disk.**\n*(When this batch finishes: Export, Reset, and click Load Words again!)*"
-        else:
-            msg += "✅ **No words remaining on server disk.**"
-            
-        await update_msg(msg)
+        msg = f"**Database Loaded**\n• Added: `{len(to_queue):,}`\n• Skipped: `{skipped:,}`\n• Disk Remaining: `{len(remaining_lines):,}`"
+        await update_msg(msg, buttons=get_main_menu())
 
     @client.on(events.CallbackQuery())
     async def button_handler(event):
-        if not is_admin(event):
-            try:
-                await event.answer("⛔ Unauthorized.", alert=True)
-            except Exception:
-                pass
-            return
-
-        # Instantly tell Telegram we received the click to stop the loading clock
-        try:
-            await event.answer()
-        except Exception:
-            pass
+        if not is_admin(event): return
+        try: await event.answer() 
+        except Exception: pass
 
         data = event.data.decode('utf-8')
 
+        # Menu Navigation
         if data == "back_main":
-            await event.edit("👋 **Fragment Scanner Bot**\nSelect an action below:", buttons=get_main_buttons())
+            await event.edit("**Fragment Control System**\nSelect an operation:", buttons=get_main_menu())
+        elif data == "menu_settings":
+            await event.edit("**System Settings**\nConfigure scanner variables:", buttons=get_settings_menu())
 
+        # Logic
         elif data == "load_words":
-            # Immediate UI feedback
-            await event.edit("⏳ **Reading and loading words... Please wait.**")
+            await event.edit("Processing chunk logic in background...")
             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.edit(f"🚀 **Scan started!**\n`{qlen:,}` words currently in the processing queue.", buttons=get_main_buttons())
+            await event.edit("Scanner initialized and running.", buttons=get_main_menu())
             
         elif data == "pause_scan":
             state = await db.get_state()
             if state.get("paused") == "1":
                 await db.set_state(paused="0")
-                await event.edit("▶️ **Scan resumed.**", buttons=get_main_buttons())
+                await event.edit("Scanner resumed.", buttons=get_main_menu())
             else:
                 await db.set_state(paused="1")
-                await event.edit("⏸ **Scan paused.**", buttons=get_main_buttons())
+                await event.edit("Scanner paused.", buttons=get_main_menu())
                 
         elif data == "stop_scan":
             await db.set_state(running="0", paused="0")
-            await event.edit("⏹ **Scan stopped.** Progress has been saved.", buttons=get_main_buttons())
+            await event.edit("Scanner terminated.", buttons=get_main_menu())
             
         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=buttons)
-            except Exception:
-                pass 
+            buttons = [[Button.inline("Refresh Data", b"refresh_status")], [Button.inline("« Back", b"back_main")]]
+            try: await event.edit(msg, buttons=buttons)
+            except Exception: pass 
             
         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("🔙 Cancel", b"back_main")]
+                [Button.inline("Slow (10)", b"spd_10"), Button.inline("Normal (20)", b"spd_20")],
+                [Button.inline("Fast (30)", b"spd_30"), Button.inline("Max (50)", b"spd_50")],
+                [Button.inline("« Back", b"menu_settings")]
             ]
-            await event.edit("⚡ **Select Scanner Speed (Concurrent Workers):**", buttons=buttons)
+            await event.edit("**Thread Allocation**\nSelect maximum concurrent workers:", buttons=buttons)
             
         elif data.startswith("spd_"):
             try:
                 speed_limit = int(data.split("_")[1])
                 await db.set_concurrency(speed_limit)
-                await event.edit(f"⚡ **Speed updated successfully!**\nNow running with `{speed_limit}` concurrent workers.", buttons=get_main_buttons())
+                await event.edit(f"Thread limit updated to `{speed_limit}`.", buttons=get_settings_menu())
             except Exception:
-                await event.edit("❌ Error setting speed.", buttons=get_main_buttons())
+                pass
             
         elif data == "export_files":
-            await event.edit("⏳ **Exporting data...**\nPlease wait while files are generated.")
+            await event.edit("Generating export packages...")
             
             taken = sorted(await db.get_all_taken())
             unavail = sorted(await db.get_all_unavailable())
             sold = sorted(await db.get_all_sold())
             forsale = sorted(await db.get_all_forsale())
             
-            await asyncio.sleep(0) # Keep UI smooth
+            await asyncio.sleep(0) 
             
-            exported_count = 0
+            if taken: await client.send_file(event.chat_id, file=io.BytesIO("\n".join(taken).encode('utf-8')), caption="taken.txt")
+            if unavail: await client.send_file(event.chat_id, file=io.BytesIO("\n".join(unavail).encode('utf-8')), caption="unavailable.txt")
+            if sold: await client.send_file(event.chat_id, file=io.BytesIO("\n".join(sold).encode('utf-8')), caption="sold.txt")
+            if forsale: await client.send_file(event.chat_id, file=io.BytesIO("\n".join(forsale).encode('utf-8')), caption="forsale.txt")
             
-            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)
-                exported_count += 1
-                
-            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)
-                exported_count += 1
-                
-            if sold:
-                f_sold = io.BytesIO("\n".join(sold).encode('utf-8'))
-                f_sold.name = "sold.txt"
-                await client.send_file(event.chat_id, file=f_sold)
-                exported_count += 1
-                
-            if forsale:
-                f_forsale = io.BytesIO("\n".join(forsale).encode('utf-8'))
-                f_forsale.name = "forsale.txt"
-                await client.send_file(event.chat_id, file=f_forsale)
-                exported_count += 1
-            
-            await event.edit(f"✅ **Export Complete!**\nGenerated `{exported_count}` text files.", buttons=get_main_buttons())
+            await event.edit("Export sequence complete.", buttons=get_main_menu())
             
         elif data == "reset_confirm":
-            buttons = [
-                [Button.inline("✅ Yes, WIPE ALL DATA", b"reset_do")],
-                [Button.inline("❌ Cancel", b"back_main")]
-            ]
-            await event.edit("⚠️ **DANGER ZONE** ⚠️\n\nAre you sure you want to completely wipe the Redis database? All progress will be lost.", buttons=buttons)
+            buttons = [[Button.inline("CONFIRM WIPE", b"reset_do")], [Button.inline("Cancel", b"menu_settings")]]
+            await event.edit("**AUTHORIZATION REQUIRED**\n\nThis will purge the active Redis database. Confirm?", buttons=buttons)
             
         elif data == "reset_do":
-            await event.edit("⏳ **Wiping database...**")
+            await event.edit("Purging memory...")
             await db.flush_all()
-            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())
+            await event.edit("Database wiped. Awaiting new chunk allocation.", buttons=get_settings_menu())