handlers/admin.py

# filename: handlers/admin.py

import logging
from telethon import events
from datetime import datetime, timedelta
import re

from core.bot import bot, BATCH_JOBS, PREMIUM_QUEUE, FREE_QUEUE
from database import manager as db_manager
import templates
import config

logger = logging.getLogger(__name__)

# This decorator ensures only the owner can use these commands
@bot.on(events.NewMessage(from_users=config.OWNER_ID))
async def admin_handler(event):
    # This function acts as a router for all admin commands
    text = event.text.strip()
    command = text.split(' ')[0]

    # Simple commands
    if command == '/stats':
        await stats_command(event)
    # Commands with arguments
    elif command == '/grantpremium':
        await grant_premium_command(event)
    elif command == '/ban':
        await ban_command(event, ban_status=True)
    elif command == '/unban':
        await ban_command(event, ban_status=False)
    elif command == '/authgroup':
        await auth_group_command(event, status=True)
    elif command == '/deauthgroup':
        await auth_group_command(event, status=False)

async def stats_command(event):
    active_batches = len(BATCH_JOBS)
    premium_q_size = PREMIUM_QUEUE.qsize()
    free_q_size = FREE_QUEUE.qsize()
    
    stats_text = (
        "**🤖 Bot Status & Health**\n\n"
        f"- **Active Batches:** `{active_batches}`\n"
        f"- **Premium Queue:** `{premium_q_size}` tasks pending\n"
        f"- **Free Queue:** `{free_q_size}` tasks pending\n"
        f"- **Premium Workers:** `{config.PREMIUM_WORKERS}`\n"
        f"- **Free Workers:** `{config.FREE_WORKERS}`"
    )
    await event.reply(stats_text)

async def grant_premium_command(event):
    # Expected format: /grantpremium 1d 123456789
    parts = event.text.split()
    if len(parts) != 3:
        await event.reply("Usage: `/grantpremium <duration> <user_id>` (e.g., `1d`, `7m`, `12h`)")
        return

    duration_str = parts[1]
    user_id_str = parts[2]

    try:
        user_id = int(user_id_str)
        
        match = re.match(r"(\d+)([dhm])", duration_str.lower())
        if not match:
            raise ValueError("Invalid duration format.")
        
        value, unit = int(match.group(1)), match.group(2)
        
        if unit == 'd':
            delta = timedelta(days=value)
        elif unit == 'm':
            delta = timedelta(days=value * 30) # Approx. for months
        elif unit == 'h':
            delta = timedelta(hours=value)
        
        success = await db_manager.grant_premium(user_id, delta)
        
        if success:
            await event.reply(templates.BotResponses.ADMIN_PREMIUM_GRANTED.format(user_id=user_id, duration=duration_str))
        else:
            await event.reply(f"Error: Could not find user with ID `{user_id}`.")

    except ValueError as e:
        await event.reply(f"Invalid input: {e}. Please check the User ID and duration format.")

async def ban_command(event, ban_status: bool):
    # Expected format: /ban 123456789
    parts = event.text.split()
    if len(parts) != 2:
        await event.reply(f"Usage: `/{'ban' if ban_status else 'unban'} <user_id>`")
        return
        
    try:
        user_id = int(parts[1])
        success = await db_manager.ban_user(user_id, ban_status=ban_status)
        if success:
            template = templates.BotResponses.ADMIN_USER_BANNED if ban_status else templates.BotResponses.ADMIN_USER_UNBANNED
            await event.reply(template.format(user_id=user_id))
        else:
            await event.reply(f"Error: Could not find user with ID `{user_id}`.")
    except ValueError:
        await event.reply("Invalid User ID format.")

async def auth_group_command(event, status: bool):
    # Expected format: /authgroup -10012345678
    parts = event.text.split()
    if len(parts) != 2:
        await event.reply(f"Usage: `/{'authgroup' if status else 'deauthgroup'} <group_id>`")
        return
        
    try:
        group_id = int(parts[1])
        await db_manager.authorize_group(group_id, status=status)
        template = templates.BotResponses.ADMIN_GROUP_AUTH if status else templates.BotResponses.ADMIN_GROUP_DEAUTH
        await event.reply(template.format(group_id=group_id))
    except ValueError:
        await event.reply("Invalid Group ID format. Make sure it's a number (and usually negative).")