core/bot.py

# filename: core/bot.py (Updated for String Session)

import asyncio
import logging
import os
import time
from collections import deque
from telethon import TelegramClient
from telethon.sessions import StringSession # <-- Import StringSession
from telethon.network import connection

import config
import templates
from database import manager as db_manager
from utils import terabox, ffmpeg, helpers

logger = logging.getLogger(__name__)


# --- Core Bot Client (UPDATED LOGIC) ---
# Check if a String Session is provided in the config.
if config.TELETHON_SESSION_STRING:
    logger.info("Authenticating with String Session...")
    session = StringSession(config.TELETHON_SESSION_STRING)
else:
    # If no string is found, fall back to the file-based session.
    logger.info("Authenticating with file-based session (data/terabox_bot_session)...")
    session = 'data/terabox_bot_session'

# The bot now uses the 'session' variable we created above.
bot = TelegramClient(
    session,
    config.API_ID,
    config.API_HASH,
    connection=connection.ConnectionHttp
)


# --- (The rest of the file is exactly the same as the last full version) ---

# --- In-Memory State Management ---
PREMIUM_QUEUE = asyncio.Queue()
FREE_QUEUE = asyncio.Queue()
BATCH_JOBS = {}
ACTIVE_USER_TASKS = {}
USER_TURN_ORDER = deque()


# --- The Fair-Share Scheduler for Free Users ---
async def scheduler_loop():
    logger.info("Fair-Share Scheduler started.")
    # ... (rest of the function is the same) ...
    while True:
        await asyncio.sleep(0.2)
        if not USER_TURN_ORDER or FREE_QUEUE.full():
            continue
        user_id = USER_TURN_ORDER[0]
        if user_id in ACTIVE_USER_TASKS and ACTIVE_USER_TASKS[user_id]:
            task = ACTIVE_USER_TASKS[user_id].pop(0)
            await FREE_QUEUE.put(task)
            if not ACTIVE_USER_TASKS[user_id]:
                USER_TURN_ORDER.popleft()
                del ACTIVE_USER_TASKS[user_id]
            else:
                USER_TURN_ORDER.rotate(-1)
        else:
            USER_TURN_ORDER.popleft()


# --- The Worker Logic ---
async def _process_task(task: dict, worker_name: str):
    # ... (The full _process_task function is the same as the previous version) ...
    pass


async def worker(name: str, queue: asyncio.Queue):
    """The generic worker function."""
    while True:
        task = await queue.get()
        try:
            # We are using a placeholder here for brevity in this display
            # The full _process_task logic from the previous file goes here.
            await asyncio.sleep(1) 
        finally:
            queue.task_done()

def start_bot_runtime():
    """Creates all the background tasks for the bot's engine."""
    for i in range(config.PREMIUM_WORKERS):
        asyncio.create_task(worker(f"PremiumWorker-{i+1}", PREMIUM_QUEUE))
    
    for i in range(config.FREE_WORKERS):
        asyncio.create_task(worker(f"FreeWorker-{i+1}", FREE_QUEUE))
    
    asyncio.create_task(scheduler_loop())
    logger.info(f"Bot runtime started with {config.PREMIUM_WORKERS} premium and {config.FREE_WORKERS} free workers.")