utils/clients.py

import asyncio, config
import traceback
from pathlib import Path
from pyrogram import Client
from pyrogram.errors import AuthKeyDuplicated
from utils.directoryHandler import backup_drive_data, loadDriveData
from utils.logger import Logger

logger = Logger(__name__)

multi_clients = {}
premium_clients = {}
work_loads = {}
premium_work_loads = {}
main_bot = None


async def initialize_clients():
    global multi_clients, work_loads, premium_clients, premium_work_loads
    logger.info("Initializing Clients")

    session_cache_path = Path(f"./cache")
    session_cache_path.parent.mkdir(parents=True, exist_ok=True)

    # Log configuration for debugging
    logger.info(f"BOT_SESSIONS configured: {len(config.BOT_SESSIONS)} session(s)")
    logger.info(f"BOT_TOKENS configured: {len(config.BOT_TOKENS)} token(s)")
    if config.BOT_SESSIONS:
        logger.info(f"BOT_SESSIONS: {config.BOT_SESSIONS}")
    if config.BOT_TOKENS:
        logger.info(f"BOT_TOKENS: {len(config.BOT_TOKENS[0]) if config.BOT_TOKENS else 0} chars (hidden)")

    # Use BOT_SESSIONS as Client names (if provided), otherwise use client_id
    # Match BOT_SESSIONS with BOT_TOKENS
    all_bot_sessions = dict((i, s) for i, s in enumerate(config.BOT_SESSIONS, start=1))
    all_tokens = dict((i, t) for i, t in enumerate(config.BOT_TOKENS, start=1))
    
    # Check if BOT_TOKENS is empty
    if not all_tokens:
        logger.error("BOT_TOKENS is empty! You must provide BOT_TOKENS in environment variables.")
        logger.error("BOT_SESSIONS alone cannot work - Pyrogram requires bot_token for bot clients.")
        return False
    
    # Match sessions with tokens
    if all_bot_sessions:
        if len(all_bot_sessions) != len(all_tokens):
            logger.warning(f"BOT_SESSIONS count ({len(all_bot_sessions)}) doesn't match BOT_TOKENS count ({len(all_tokens)})")
            min_len = min(len(all_bot_sessions), len(all_tokens))
            all_bot_sessions = dict(list(all_bot_sessions.items())[:min_len])
            all_tokens = dict(list(all_tokens.items())[:min_len])
        # Zip sessions and tokens
        bot_configs = dict((i, (session, token)) for i, (session, token) in 
                          enumerate(zip(all_bot_sessions.values(), all_tokens.values()), start=1))
    else:
        # No BOT_SESSIONS, use client_id as name
        bot_configs = dict((i, (None, token)) for i, token in all_tokens.items())
    
    all_string_sessions = dict(
        (i, s) for i, s in enumerate(config.STRING_SESSIONS, start=len(bot_configs) + 1)
    )

    async def start_client(client_id, session_name_and_token, type):
        try:
            logger.info(f"Starting - {type.title()} Client {client_id}")

            if type == "bot":
                # Get session name and token
                session_name, token = session_name_and_token
                
                # Use session name as Client name (remove .session extension if present)
                # If no session name, use client_id
                if session_name:
                    client_name = session_name.replace(".session", "")
                else:
                    client_name = str(client_id)
                
                # Pyrogram will use existing session file if it exists in cache/, or create new one
                client = Client(
                    name=client_name,
                    api_id=config.API_ID,
                    api_hash=config.API_HASH,
                    bot_token=token,
                    workdir=session_cache_path,
                )
                client.loop = asyncio.get_running_loop()
                try:
                    await client.start()
                    await client.send_message(
                        config.STORAGE_CHANNEL,
                        f"Started - Bot Client {client_id}",
                    )
                    multi_clients[client_id] = client
                    work_loads[client_id] = 0
                except AuthKeyDuplicated as e:
                    # Delete the invalidated session file and mark it as invalidated
                    session_file = None
                    try:
                        session_file = session_cache_path / f"{client_name}.session"
                        if session_file.exists():
                            session_file.unlink()
                            logger.info(f"Deleted invalidated session file: {session_file}")
                            # Create a marker file so we don't restore this session on next restart
                            invalidated_marker = session_cache_path / f"{client_name}.session.invalidated"
                            invalidated_marker.touch()
                            logger.info(f"Marked session file as invalidated: {invalidated_marker}")
                    except Exception as cleanup_error:
                        logger.error(f"Failed to delete session file: {cleanup_error}")
                    
                    error_msg = (
                        f"\n{'='*60}\n"
                        f"AUTH_KEY_DUPLICATED for Bot Client {client_id}\n"
                        f"{'='*60}\n"
                        f"CRITICAL: The same bot token is being used in MULTIPLE places simultaneously.\n"
                        f"Telegram has invalidated the session file for security reasons.\n\n"
                        f"✅ Session file has been deleted: {session_file if session_file else 'N/A'}\n\n"
                        f"⚠️  TO FIX THIS:\n"
                        f"   1. STOP using this bot token in ALL other locations (local servers, other Spaces, etc.)\n"
                        f"   2. Wait a few seconds for the other instance to disconnect\n"
                        f"   3. Restart this application (the session file will be recreated automatically)\n\n"
                        f"📌 REMEMBER: Each bot token can only be used in ONE place at a time!\n"
                        f"{'='*60}"
                    )
                    logger.error(error_msg)
                    logger.error(f"Error details: {e}")
                    
                    # Don't re-raise - let the client fail gracefully and app continue in offline mode
                    return
            elif type == "string_session":
                # Use string session (user account)
                session_string = session_name_and_token  # For string sessions, it's just the string
                client = await Client(
                    name=str(client_id),
                    api_id=config.API_ID,
                    api_hash=config.API_HASH,
                    session_string=session_string,
                    sleep_threshold=config.SLEEP_THRESHOLD,
                    workdir=session_cache_path,
                    no_updates=True,
                ).start()
                await client.send_message(
                    config.STORAGE_CHANNEL,
                    f"Started - String Session Client {client_id}",
                )
                premium_clients[client_id] = client
                premium_work_loads[client_id] = 0

            logger.info(f"Started - {type.title()} Client {client_id}")
        except Exception as e:
            error_msg = str(e) if str(e) else repr(e)
            error_trace = traceback.format_exc()
            logger.error(
                f"Failed To Start {type.title()} Client - {client_id} Error: {error_msg}"
            )
            logger.error(f"Traceback: {error_trace}")

    # Start bot clients and string session clients
    await asyncio.gather(
        *(
            [
                start_client(client_id, (session_name, token), "bot")
                for client_id, (session_name, token) in bot_configs.items()
            ]
            + [
                start_client(client_id, session, "string_session")
                for client_id, session in all_string_sessions.items()
            ]
        )
    )
    
    clients_initialized = len(multi_clients) > 0
    
    if not clients_initialized:
        logger.warning("No Clients Were Initialized - Website will run in offline mode")
        logger.warning("Please configure BOT_TOKENS (and optionally BOT_SESSIONS) in your environment variables")
        return False

    logger.info("Clients Initialized")

    # Load the drive data
    await loadDriveData()

    # Start the backup drive data task
    asyncio.create_task(backup_drive_data())
    
    return True


def has_clients():
    """Check if any clients are available"""
    global multi_clients
    return len(multi_clients) > 0


def get_client(premium_required=False) -> Client:
    global multi_clients, work_loads, premium_clients, premium_work_loads

    if premium_required:
        if not premium_work_loads:
            raise RuntimeError("No Premium Clients Available")
        index = min(premium_work_loads, key=premium_work_loads.get)
        premium_work_loads[index] += 1
        return premium_clients[index]

    if not work_loads:
        raise RuntimeError("No Clients Available - Check BOT_TOKENS/BOT_SESSIONS configuration")
    index = min(work_loads, key=work_loads.get)
    work_loads[index] += 1
    return multi_clients[index]