main.py

import asyncio
import os
import httpx
from fastapi import FastAPI
from contextlib import asynccontextmanager
from pyrogram import Client, filters
import config
import database

# Helper function to log user activity and update metadata
def log_and_update(message):
    if message.from_user:
        database.update_user_metadata(
            message.from_user.id,
            message.from_user.username,
            message.from_user.first_name,
            message.from_user.last_name
        )

# Initialize Pyrogram Bot Client
# Use in_memory=True to avoid creating local session db files on Hugging Face Docker Space disk
bot = Client(
    "pinger_bot",
    api_id=config.API_ID,
    api_hash=config.API_HASH,
    bot_token=config.BOT_TOKEN,
    in_memory=True
)

# Reference to background loop task
ping_task = None

async def ping_url(url: str):
    headers = {
        "User-Agent": "HF-Space-Pinger-Bot/1.0"
    }
    try:
        # Perform HTTP GET request to wake/keep the container alive
        async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
            response = await client.get(url, headers=headers)
            print(f"Pinger: Success {url} -> Status {response.status_code}")
            return response.status_code
    except Exception as e:
        print(f"Pinger: Error pinging {url} -> {e}")
        return None

def get_own_space_url():
    space_host = os.environ.get("SPACE_HOST")
    if space_host:
        return f"https://{space_host}"
    
    space_id = os.environ.get("SPACE_ID")
    if space_id:
        try:
            return database.normalize_hf_url(space_id)
        except Exception:
            pass
    return None

async def ping_loop():
    print("Background ping loop started.")
    # Initial sleep of 10s to let the startup finalize
    await asyncio.sleep(10)
    
    while True:
        try:
            print("Pinger: Starting a new ping cycle...")
            # Retrieve all URLs
            urls = database.get_all_urls()
            
            # Make sure we ping our own space URL
            own_url = get_own_space_url()
            if own_url:
                if own_url not in urls:
                    urls = [own_url] + urls
                    print(f"Pinger: Appended own Space URL: {own_url}")
            
            print(f"Pinger: Total spaces to ping: {len(urls)}")
            
            # Ping one by one sequentially
            for url in urls:
                await ping_url(url)
                # Wait 2 seconds between spaces as requested
                await asyncio.sleep(2)
                
            print("Pinger: Ping cycle complete. Waiting 5 minutes.")
        except asyncio.CancelledError:
            print("Pinger: Background task cancelled.")
            break
        except Exception as e:
            print(f"Pinger: Unexpected error in loop: {e}")
            
        # Wait 5 minutes (300 seconds)
        await asyncio.sleep(300)

@asynccontextmanager
async def lifespan(app: FastAPI):
    global ping_task
    
    # 1. Initialize database and load sync from HF dataset
    print("Lifespan: Initializing database state...")
    database.init_db()
    
    # 2. Start the Pyrogram client
    print("Lifespan: Starting Telegram bot...")
    await bot.start()
    print("Lifespan: Telegram bot started successfully.")
    
    # 3. Spawn background pinger loop
    ping_task = asyncio.create_task(ping_loop())
    
    yield
    
    # Shutdown sequence
    print("Lifespan: Initiating shutdown sequence...")
    if ping_task:
        ping_task.cancel()
        try:
            await ping_task
        except asyncio.CancelledError:
            pass
    
    await bot.stop()
    print("Lifespan: Shutdown complete.")

app = FastAPI(lifespan=lifespan)

# Health Check Endpoint
@app.get("/")
async def health_check():
    return {"status": "ok", "message": "HF Space Pinger Bot is running successfully."}

# ================== PYROGRAM BOT HANDLERS ==================

@bot.on_message(filters.command("start"))
async def start_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_authorized(user_id):
        await message.reply_text(
            "👋 **Hello! Welcome to the Hugging Face Space Pinger Bot.** 🚀\n\n"
            "You are currently unauthorized to use this bot.\n"
            f"🔑 Your ID: `{user_id}`\n\n"
            "Please contact an admin to allow you access."
        )
        return

    # Welcomes the user
    await message.reply_text(
        "👋 **Hello! Welcome to the Hugging Face Space Pinger Bot.** 🚀\n\n"
        "I am designed to keep Hugging Face spaces alive by pinging them at a 5-minute interval."
    )

@bot.on_message(filters.command("help"))
async def help_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    # Explains the bot
    await message.reply_text(
        "🤖 **What this bot does:**\n\n"
        "Hugging Face free tier spaces automatically go to sleep after some inactivity.\n"
        "This bot pings registered Space URLs (in the format of `username-spacename.hf.space`) "
        "sequentially every 5 minutes to keep them awake, including its own running space container.\n\n"
        "Send `/commands` to see what instructions you can run, or simply send a Hugging Face Space URL to register it automatically!"
    )

@bot.on_message(filters.command("commands"))
async def commands_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    is_auth = database.is_authorized(user_id)
    is_adm = database.is_admin(user_id)
    is_sa = database.is_super_admin(user_id)
    
    cmds = "📋 **Available Commands:**\n\n"
    
    # Public command
    cmds += "💬 **Public Commands:**\n"
    cmds += "- `/id` - Show your Telegram User ID (useful for getting authorized)\n"
    cmds += "- `/start` - Send welcome message\n"
    cmds += "- `/help` - Tell what the bot does\n"
    cmds += "- `/commands` - List available commands\n\n"
    
    if is_auth:
        cmds += "👤 **User Commands:**\n"
        cmds += "- `/add <hf_url>` - Register a Hugging Face Space URL\n"
        cmds += "- `/remove <hf_url>` - Remove a registered Space URL\n"
        cmds += "- `/list_url` - List your registered URLs\n\n"
        
    if is_adm:
        cmds += "👮 **Admin Commands:**\n"
        cmds += "- `/add_user <user_id>` - Authorize a user\n"
        cmds += "- `/remove_user <user_id>` - Deauthorize a user\n"
        cmds += "- `/add_admin <user_id>` - Promote a user to Admin\n"
        cmds += "- `/list` - List admins and authorized users (excluding super admin)\n"
        cmds += "- `/list_url` - List ALL registered space URLs in the system\n\n"
        
    if is_sa:
        cmds += "👑 **Super Admin Commands:**\n"
        cmds += "- `/remove_admin <user_id>` - Demote an Admin\n"
        cmds += "- `/list` - List all registered admins, users, and the super admin\n\n"
        
    if not is_auth:
        cmds += "⚠️ _You are not currently authorized. Contact the administrator to register._"
        
    await message.reply_text(cmds)

@bot.on_message(filters.command("id"))
async def id_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    # Shows user ID
    await message.reply_text(f"🔑 Your Telegram User ID is: `{message.from_user.id}`")

@bot.on_message(filters.command("add"))
async def add_url_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_authorized(user_id):
        await message.reply_text("❌ You are not authorized. Contact an admin.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/add <hf_space_url1> <hf_space_url2> ...`\n\nExample:\n`/add https://huggingface.co/spaces/iozxv/C1`")
        return
        
    urls_input = message.command[1:]
    results = []
    
    for url_input in urls_input:
        try:
            success, normalized = database.add_url(url_input, user_id)
            if success:
                results.append(f"✅ Added: `{normalized}`")
            else:
                results.append(f"⚠️ Already exists: `{normalized}`")
        except ValueError as e:
            results.append(f"❌ Invalid: {url_input} - {str(e)}")
        except Exception as e:
            results.append(f"❌ Error: {url_input} - {str(e)}")
            
    await message.reply_text("\n".join(results))

@bot.on_message(filters.command("remove"))
async def remove_url_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_authorized(user_id):
        await message.reply_text("❌ You are not authorized. Contact an admin.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/remove <hf_space_url1> <hf_space_url2> ...`")
        return
        
    urls_input = message.command[1:]
    results = []
    
    for url_input in urls_input:
        try:
            success, normalized = database.remove_url(url_input, user_id)
            if success:
                results.append(f"✅ Removed: `{normalized}`")
            else:
                results.append(f"❌ Not found: `{normalized}`")
        except PermissionError as e:
            results.append(f"❌ Permission Denied: {url_input} - {str(e)}")
        except Exception as e:
            results.append(f"❌ Error: {url_input} - {str(e)}")
            
    await message.reply_text("\n".join(results))

@bot.on_message(filters.command("list_url"))
async def list_url_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_authorized(user_id):
        await message.reply_text("❌ You are not authorized. Contact an admin.")
        return
        
    urls_data = database.get_urls(user_id)
    
    if not urls_data:
        await message.reply_text("ℹ️ No URLs found in your list.")
        return
        
    # If admin or super admin, show who added it. Otherwise, just list URLs
    if database.is_admin(user_id):
        response = f"🔗 **All Monitored Space URLs ({len(urls_data)}):**\n\n"
        
        # Group by user
        grouped = {}
        for item in urls_data:
            added_by = item["added_by"]
            if added_by not in grouped:
                grouped[added_by] = []
            grouped[added_by].append(item["url"])
            
        for user, urls in grouped.items():
            user_display = database.get_user_display(user)
            response += f"👤 **Added by: {user_display}**\n"
            for u in urls:
                response += f"- `{u}`\n"
            response += "\n"
    else:
        response = f"🔗 **Your Monitored Space URLs ({len(urls_data)}):**\n\n"
        for item in urls_data:
            response += f"- `{item['url']}`\n"
            
    await message.reply_text(response)

@bot.on_message(filters.command("list"))
async def list_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_admin(user_id):
        await message.reply_text("❌ You must be an Admin to use this command.")
        return
        
    admins = database.list_admins()
    users = database.list_users()
    
    response = "👥 **Registered Bot Accounts**\n\n"
    
    if database.is_super_admin(user_id):
        # Super Admin sees everything including Super Admin itself
        response += f"👑 **Super Admin:**\n- {database.get_user_display(config.SUPER_ADMIN_ID)}\n\n"
        
        response += f"👮 **Admins ({len(admins)}):**\n"
        if admins:
            for admin_id in admins:
                response += f"- {database.get_user_display(admin_id)}\n"
        else:
            response += "- No admins added yet.\n"
            
        response += f"\n👤 **Authorized Users ({len(users)}):**\n"
        if users:
            for u_id in users:
                response += f"- {database.get_user_display(u_id)}\n"
        else:
            response += "- No authorized users added yet.\n"
    else:
        # Regular Admin sees only admins and users (not super admin, as requested)
        response += f"👮 **Admins ({len(admins)}):**\n"
        if admins:
            for admin_id in admins:
                response += f"- {database.get_user_display(admin_id)}\n"
        else:
            response += "- No admins added yet.\n"
            
        response += f"\n👤 **Authorized Users ({len(users)}):**\n"
        if users:
            for u_id in users:
                response += f"- {database.get_user_display(u_id)}\n"
        else:
            response += "- No authorized users added yet.\n"
            
    await message.reply_text(response)

@bot.on_message(filters.command("add_user"))
async def add_user_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_admin(user_id):
        await message.reply_text("❌ You must be an Admin to use this command.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/add_user <telegram_user_id>`")
        return
        
    target_id_str = message.command[1]
    try:
        target_id = int(target_id_str)
        success = database.add_user(target_id)
        if success:
            await message.reply_text(f"✅ User `{target_id}` has been authorized successfully.")
        else:
            await message.reply_text(f"ℹ️ User `{target_id}` is already authorized.")
    except ValueError:
        await message.reply_text("❌ **Error:** Please provide a valid integer Telegram User ID.")

@bot.on_message(filters.command("remove_user"))
async def remove_user_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_admin(user_id):
        await message.reply_text("❌ You must be an Admin to use this command.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/remove_user <telegram_user_id>`")
        return
        
    target_id_str = message.command[1]
    try:
        target_id = int(target_id_str)
        success = database.remove_user(target_id)
        if success:
            await message.reply_text(f"✅ User `{target_id}` has been deauthorized.")
        else:
            await message.reply_text(f"❌ User `{target_id}` is not in the authorized list.")
    except ValueError:
        await message.reply_text("❌ **Error:** Please provide a valid integer Telegram User ID.")

@bot.on_message(filters.command("add_admin"))
async def add_admin_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_admin(user_id):
        await message.reply_text("❌ You must be an Admin to use this command.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/add_admin <telegram_user_id>`")
        return
        
    target_id_str = message.command[1]
    try:
        target_id = int(target_id_str)
        success = database.add_admin(target_id)
        if success:
            await message.reply_text(f"✅ User `{target_id}` has been promoted to Admin.")
        else:
            await message.reply_text(f"ℹ️ User `{target_id}` is already an Admin.")
    except ValueError:
        await message.reply_text("❌ **Error:** Please provide a valid integer Telegram User ID.")

@bot.on_message(filters.command("remove_admin"))
async def remove_admin_command(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    user_id = message.from_user.id
    if not database.is_super_admin(user_id):
        await message.reply_text("❌ Only the Super Admin can demote other Admins.")
        return
        
    if len(message.command) < 2:
        await message.reply_text("Usage: `/remove_admin <telegram_user_id>`")
        return
        
    target_id_str = message.command[1]
    try:
        target_id = int(target_id_str)
        if target_id == config.SUPER_ADMIN_ID:
            await message.reply_text("❌ You cannot demote yourself.")
            return
        success = database.remove_admin(target_id)
        if success:
            await message.reply_text(f"✅ Admin `{target_id}` has been demoted to User.")
        else:
            await message.reply_text(f"❌ Admin `{target_id}` is not in the Admin list.")
    except ValueError:
        await message.reply_text("❌ **Error:** Please provide a valid integer Telegram User ID.")

# If someone sends a direct Hugging Face URL in the chat, capture it and add it
@bot.on_message(filters.text)
async def handle_text_message(client, message):
    if not message.from_user:
        return
    log_and_update(message)
    
    text = message.text.strip()
    # Ignore command messages
    if text.startswith("/"):
        return
    # Check if the text matches a HF Space URL or direct app endpoint
    is_hf_url = "huggingface.co/spaces" in text or "hf.space" in text
    
    if is_hf_url:
        user_id = message.from_user.id
        if not database.is_authorized(user_id):
            await message.reply_text("❌ You are not authorized to add URLs. Ask an Admin to authorize your ID.")
            return
            
        # Split text into possible multiple URLs
        words = text.split()
        potential_urls = [w for w in words if "huggingface.co/spaces" in w or "hf.space" in w]
        
        results = []
        for url_input in potential_urls:
            try:
                success, normalized = database.add_url(url_input, user_id)
                if success:
                    results.append(f"✅ Added: `{normalized}`")
                else:
                    results.append(f"⚠️ Already exists: `{normalized}`")
            except ValueError as e:
                results.append(f"❌ Invalid: {url_input} - {str(e)}")
            except Exception as e:
                results.append(f"❌ Error: {url_input} - {str(e)}")
                
        if results:
            await message.reply_text("\n".join(results))

# ==================== MAIN EXECUTION ====================

if __name__ == "__main__":
    import uvicorn
    # Bind to 0.0.0.0 and port 7860 as required by Hugging Face Spaces
    port = int(os.environ.get("PORT", 7860))
    print(f"Starting FastAPI webserver and bot client on port {port}...")
    uvicorn.run("main:app", host="0.0.0.0", port=port, log_level="info")