backend/worker/telegram_bot.py

import asyncio
from datetime import datetime, timedelta
from telethon import TelegramClient, events, Button
from bson import ObjectId

from backend.core.config import settings
from backend.core.database import get_database
from backend.core.security import encrypt_hf_token, decrypt_hf_token
from backend.services.huggingface import validate_hf_token, get_user_spaces

LOG_GROUP_ID = -1001744991128
ADMIN_ID = 948247711

# Initialize the bot client
monitor_bot = TelegramClient('monitor_session', settings.TELEGRAM_API_ID, settings.TELEGRAM_API_HASH)

# In-memory state tracking
user_states = {}

async def send_sys_log(message: str):
    """Global logging function."""
    try:
        await monitor_bot.send_message(LOG_GROUP_ID, f"🔧 **SYS_LOG**\n\n{message}")
    except Exception as e:
        print(f"Failed to send log to group: {e}")

async def get_or_create_user(telegram_id: int, username: str):
    """Ensures user exists in DB on first interaction."""
    db = get_database()
    str_id = str(telegram_id)
    user = await db["users"].find_one({"telegram_id": str_id})
    if not user:
        user = {
            "telegram_id": str_id,
            "telegram_username": username or str_id,
            "is_admin": telegram_id == ADMIN_ID,
            "created_at": datetime.utcnow()
        }
        await db["users"].insert_one(user)
        await send_sys_log(f"👤 **New Bot User:** `{str_id}` just started the bot.")
    return user

def get_main_menu(is_admin=False):
    """Generates the main dashboard inline keyboard. Adds Admin Panel if user is Root."""
    buttons = [
        [Button.inline("🚀 Manage Spaces", b"menu_spaces"), Button.inline("⚙️ Custom API Jobs", b"menu_custom")],
        [Button.inline("📊 My Stats", b"menu_stats")]
    ]
    if is_admin:
        buttons.append([Button.inline("👑 Root Admin Panel", b"admin_panel")])
    return buttons

# ==========================================
# TEXT & COMMAND HANDLERS
# ==========================================

@monitor_bot.on(events.NewMessage)
async def message_handler(event):
    if event.is_group:
        return  
        
    sender = await event.get_sender()
    user_id = sender.id
    text = event.raw_text.strip()
    
    user = await get_or_create_user(user_id, sender.username)

    # 👑 SUPER ADMIN TEXT COMMANDS
    if text.startswith("/promote"):
        if user_id != ADMIN_ID:
            return await event.reply("⛔ **Access Denied:** You are not the Root Admin.")
        
        parts = text.split(" ")
        if len(parts) != 2:
            return await event.reply("⚠️ Usage: `/promote <telegram_id>`")
            
        target_id = parts[1]
        db = get_database()
        res = await db["users"].update_one({"telegram_id": target_id}, {"$set": {"is_admin": True}})
        
        if res.modified_count > 0:
            await event.reply(f"✅ **Success:** User `{target_id}` promoted to Admin.")
            await send_sys_log(f"👑 **Admin Promotion:** User `{target_id}` was promoted by Root.")
        else:
            await event.reply(f"❌ User `{target_id}` not found in database.")
        return

    # 🤖 BASIC COMMANDS
    if text == "/start":
        if not user.get("hf_token_encrypted"):
            msg = (
                "👋 **Welcome to HF Cron Manager Pro!**\n\n"
                "To manage your background jobs, you must first connect your Hugging Face account."
            )
            await event.reply(msg, buttons=[[Button.inline("🔗 Link HF Account", b"set_hf_token")]])
        else:
            await event.reply("🎛 **Main Control Panel**\nSelect an operation:", buttons=get_main_menu(user.get("is_admin")))
        return
        
    if text == "/cancel":
        if user_id in user_states:
            del user_states[user_id]
        await event.reply("❌ **Action cancelled.**", buttons=[[Button.inline("🔙 Main Menu", b"main_menu")]])
        return

    # 🔄 STATE MACHINE
    state_info = user_states.get(user_id)
    if not state_info:
        return

    state = state_info.get("state")
    db = get_database()

    if state == "waiting_hf_token":
        token = text
        await event.reply("⏳ Validating token...")
        hf_info = validate_hf_token(token)
        if not hf_info:
            return await event.reply("❌ **Invalid Token.** Please check permissions and try again, or type /cancel.")
            
        enc_token = encrypt_hf_token(token)
        await db["users"].update_one(
            {"telegram_id": str(user_id)}, 
            {"$set": {"hf_token_encrypted": enc_token, "hf_username": hf_info["name"]}}
        )
        del user_states[user_id]
        await event.reply(f"✅ **Account Linked:** `{hf_info['name']}`", buttons=[[Button.inline("➡️ Go to Main Menu", b"main_menu")]])

    elif state == "custom_api_name":
        user_states[user_id]["data"] = {"name": text}
        user_states[user_id]["state"] = "custom_api_url"
        await event.reply("🌐 **Enter the Target URL** (e.g., https://api.example.com/ping):")

    elif state == "custom_api_url":
        if not text.startswith("http"):
            return await event.reply("❌ Invalid URL. Must start with http or https. Try again or /cancel.")
        user_states[user_id]["data"]["url"] = text
        user_states[user_id]["state"] = "custom_api_method"
        await event.reply("🛠 **Enter the HTTP Method** (GET, POST, PUT, DELETE):")

    elif state == "custom_api_method":
        method = text.upper()
        if method not in ["GET", "POST", "PUT", "DELETE"]:
            return await event.reply("❌ Must be GET, POST, PUT, or DELETE. Try again.")
        user_states[user_id]["data"]["method"] = method
        user_states[user_id]["state"] = "custom_api_interval"
        await event.reply("⏱ **Enter the Interval in seconds** (e.g., 3600 for 1 hour):")

    elif state == "custom_api_interval":
        if not text.isdigit() or int(text) < 60:
            return await event.reply("❌ Interval must be a number (minimum 60). Try again.")
        
        data = user_states[user_id]["data"]
        new_job = {
            "user_id": str(user_id),
            "name": data["name"],
            "type": "custom_api",
            "target_url": data["url"],
            "method": data["method"],
            "interval_seconds": int(text),
            "status": "active",
            "next_run": datetime.utcnow() + timedelta(seconds=int(text)),
            "success_count": 0,
            "failure_count": 0,
            "created_at": datetime.utcnow()
        }
        await db["jobs"].insert_one(new_job)
        del user_states[user_id]
        await event.reply("✅ **Custom API Job Created Successfully!**", buttons=[[Button.inline("🔙 Main Menu", b"main_menu")]])

# ==========================================
# INLINE BUTTON CALLBACK HANDLERS
# ==========================================

@monitor_bot.on(events.CallbackQuery)
async def callback_handler(event):
    sender = await event.get_sender()
    user_id = sender.id
    data = event.data.decode('utf-8')
    
    db = get_database()
    user = await db["users"].find_one({"telegram_id": str(user_id)})

    if data == "set_hf_token":
        user_states[user_id] = {"state": "waiting_hf_token"}
        await event.edit(
            "🔑 **Please paste your Hugging Face Access Token:**\n"
            "*(Ensure it has Read/Write permissions)*\n\n"
            "Type /cancel to abort."
        )
        return

    if not user.get("hf_token_encrypted"):
        return await event.answer("⚠️ You must link your HF account first!", alert=True)

    if data == "main_menu":
        if user_id in user_states: del user_states[user_id]
        await event.edit("🎛 **Main Control Panel**", buttons=get_main_menu(user.get("is_admin")))

    # --- ADMIN PANEL ---
    elif data == "admin_panel":
        if not user.get("is_admin"):
            return await event.answer("⛔ Access Denied.", alert=True)
        msg = "👑 **Root Admin Panel**\nMonitor the global system."
        btns = [
            [Button.inline("📊 Global Stats", b"admin_stats"), Button.inline("👥 Recent Users", b"admin_users")],
            [Button.inline("🔙 Main Menu", b"main_menu")]
        ]
        await event.edit(msg, buttons=btns)

    elif data == "admin_stats":
        if not user.get("is_admin"): return
        users_count = await db["users"].count_documents({})
        jobs_count = await db["jobs"].count_documents({})
        active = await db["jobs"].count_documents({"status": "active"})
        logs = await db["job_logs"].count_documents({})
        
        msg = (
            "📊 **GLOBAL SYSTEM STATS**\n"
            "━━━━━━━━━━━━━━━━━━━━\n"
            f"👥 **Total Users:** `{users_count}`\n"
            f"⚙️ **Total Jobs Configured:** `{jobs_count}`\n"
            f"⚡ **Active Queue:** `{active}`\n"
            f"📡 **Total Pings Fired:** `{logs}`"
        )
        await event.edit(msg, buttons=[[Button.inline("🔙 Admin Menu", b"admin_panel")]])

    elif data == "admin_users":
        if not user.get("is_admin"): return
        cursor = db["users"].find().sort("created_at", -1).limit(10)
        users = await cursor.to_list(length=10)
        
        msg = "👥 **LAST 10 REGISTERED USERS**\n━━━━━━━━━━━━━━━━━━━━\n"
        for u in users:
            hf = u.get("hf_username", "Unlinked")
            role = "[ROOT]" if u.get("is_admin") else "[USER]"
            msg += f"ID: `{u['telegram_id']}` | HF: `{hf}` {role}\n"
            
        await event.edit(msg, buttons=[[Button.inline("🔙 Admin Menu", b"admin_panel")]])

    # --- REGULAR FEATURES ---
    elif data == "menu_stats":
        jobs_count = await db["jobs"].count_documents({"user_id": str(user_id)})
        active_jobs = await db["jobs"].count_documents({"user_id": str(user_id), "status": "active"})
        msg = (
            f"📊 **Your Personal Stats**\n"
            f"━━━━━━━━━━━━━━━━━━━━\n"
            f"👤 **HF Account:** `{user.get('hf_username')}`\n"
            f"⚙️ **Total Jobs:** `{jobs_count}`\n"
            f"⚡ **Active Jobs:** `{active_jobs}`"
        )
        await event.edit(msg, buttons=[[Button.inline("🔙 Back", b"main_menu")]])

    elif data == "menu_spaces":
        await event.edit("⏳ Scanning Hugging Face Spaces... Please wait.")
        token = decrypt_hf_token(user["hf_token_encrypted"])
        spaces = await get_user_spaces(token, user["hf_username"])
        
        if not spaces:
            return await event.edit("❌ No spaces found.", buttons=[[Button.inline("🔙 Back", b"main_menu")]])
            
        user_states[user_id] = {"spaces_cache": spaces}
        
        # Legend Explanation Added Here
        msg = (
            "🚀 **Your Spaces Grid** (Click to Manage):\n\n"
            "**Status Legend:**\n"
            "🟢 `Active / Building`\n"
            "🟡 `Sleeping / Paused`\n"
            "🔴 `Error / Offline`"
        )
        
        buttons = []
        row = []
        for idx, space in enumerate(spaces):
            raw_stage = str(space.get("runtime", {}).get("stage") or "UNKNOWN").upper()
            status_icon = "🟢" if raw_stage in ["RUNNING", "BUILDING", "INITIALIZING", "CONFIGURING"] else "🟡" if raw_stage in ["PAUSED", "SLEEPING", "STOPPED"] else "🔴"
            
            # FIX: Properly extract the space name from "username/space-name"
            space_id = space.get("id", "")
            name = space_id.split("/")[-1][:12] if space_id else "Unknown"
            
            row.append(Button.inline(f"{status_icon} {name}", f"sp_view_{idx}".encode('utf-8')))
            
            if len(row) == 3: 
                buttons.append(row)
                row = []
                
        if row: buttons.append(row)
        buttons.append([Button.inline("🔙 Main Menu", b"main_menu")])
        
        await event.edit(msg, buttons=buttons)

    elif data.startswith("sp_view_"):
        idx = int(data.split("_")[2])
        cache = user_states.get(user_id, {}).get("spaces_cache", [])
        if idx >= len(cache):
            return await event.answer("⚠️ Session expired. Please reload spaces.", alert=True)
            
        space = cache[idx]
        space_id = space["id"]
        raw_stage = str(space.get("runtime", {}).get("stage") or "UNKNOWN").upper()
        
        job = await db["jobs"].find_one({"user_id": str(user_id), "space_id": space_id, "type": "space_ping"})
        
        msg = (
            f"📦 **Space Overview**\n"
            f"━━━━━━━━━━━━━━━━━━━━\n"
            f"🏷 **Name:** `{space_id.split('/')[-1]}`\n"
            f"🔗 **URL:** [Direct Link](https://huggingface.co/spaces/{space_id})\n"
            f"🔋 **Status:** `{raw_stage}`\n"
        )
        
        btn_row = []
        if job:
            msg += f"🟢 **Auto-Wake Cron:** `ACTIVE` (Int: {job['interval_seconds']}s)"
            btn_row.append(Button.inline("🛑 Disable Cron", f"cron_del_{str(job['_id'])}".encode('utf-8')))
        else:
            msg += f"🔴 **Auto-Wake Cron:** `DISABLED`"
            btn_row.append(Button.inline("⚡ Enable Cron", f"cron_add_{idx}".encode('utf-8')))
            
        await event.edit(msg, buttons=[btn_row, [Button.inline("🔙 Back to Grid", b"menu_spaces")]], link_preview=False)

    elif data.startswith("cron_add_"):
        idx = int(data.split("_")[2])
        space = user_states.get(user_id, {}).get("spaces_cache", [])[idx]
        space_id = space["id"]
        
        new_job = {
            "user_id": str(user_id),
            "type": "space_ping",
            "space_id": space_id,
            "target_url": f"https://{space_id.replace('/', '-')}.hf.space/",
            "method": "GET",
            "interval_seconds": 900, 
            "status": "active",
            "next_run": datetime.utcnow(),
            "created_at": datetime.utcnow()
        }
        await db["jobs"].insert_one(new_job)
        await event.answer("✅ Cron Enabled!", alert=True)
        await callback_handler(event._set_data(f"sp_view_{idx}".encode('utf-8')))

    elif data.startswith("cron_del_"):
        job_id = data.split("_")[2]
        await db["jobs"].delete_one({"_id": ObjectId(job_id), "user_id": str(user_id)})
        await event.answer("🛑 Job Deleted!", alert=True)
        await callback_handler(event._set_data(b"menu_spaces"))

    elif data == "menu_custom":
        msg = "⚙️ **Custom API Jobs**\nManage your advanced webhooks and API pings."
        buttons = [
            [Button.inline("➕ Create New API Job", b"custom_create")],
            [Button.inline("📋 View Active API Jobs", b"custom_list")],
            [Button.inline("🔙 Main Menu", b"main_menu")]
        ]
        await event.edit(msg, buttons=buttons)

    elif data == "custom_create":
        user_states[user_id] = {"state": "custom_api_name"}
        await event.edit("📝 **Creating Custom Job:**\nPlease type a Name for this job (e.g., Database Backup):\n\nType /cancel to abort.")

    elif data == "custom_list":
        cursor = db["jobs"].find({"user_id": str(user_id), "type": "custom_api"})
        jobs = await cursor.to_list(length=10)
        
        if not jobs:
            return await event.edit("❌ No custom API jobs found.", buttons=[[Button.inline("🔙 Back", b"menu_custom")]])
            
        msg = "📋 **Your Custom API Jobs**\n━━━━━━━━━━━━━━━━━━━━\n"
        buttons = []
        for j in jobs:
            msg += f"🔹 `{j['name']}` | Method: `{j['method']}`\n"
            buttons.append([Button.inline(f"🗑 Delete '{j['name'][:10]}'", f"cron_del_{str(j['_id'])}".encode('utf-8'))])
            
        buttons.append([Button.inline("🔙 Back", b"menu_custom")])
        await event.edit(msg, buttons=buttons)

# ==========================================
# BOT BOOTSTRAP
# ==========================================

async def run_bot():
    """Starts the bot in the background without blocking FastAPI."""
    print("Initializing Full-Flash Telethon Bot Engine...")
    await monitor_bot.start(bot_token=settings.TELEGRAM_BOT_TOKEN)
    await send_sys_log("🚀 **System Boot Sequence Complete.** PWA and Full-Flash Telegram Bot are now ONLINE.")
    await monitor_bot.run_until_disconnected()