bot.py

# bot.py

import asyncio
import logging
import os
import re
import uuid
from typing import Dict

from aiogram import Bot, Dispatcher, F
from aiogram.filters import CommandStart
from aiogram.types import Message, FSInputFile
from aiogram.enums import ParseMode
from aiogram.exceptions import TelegramBadRequest
from aiogram.client.default import DefaultBotProperties

import config
import terabox_utils as terabox
import db_utils

# --- Setup ---
os.makedirs("logs", exist_ok=True)
os.makedirs("downloads", exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("logs/bot.log", mode='a', encoding='utf-8')
    ]
)
logger = logging.getLogger("bot")

bot = Bot(
    token=config.BOT_TOKEN,
    default=DefaultBotProperties(
        parse_mode=ParseMode.HTML
    )
)

dp = Dispatcher()
TASK_QUEUE = asyncio.Queue()
BATCH_JOBS: Dict[str, Dict] = {}

# --- Worker Logic ---
async def link_processor_worker(worker_id: int):
    logger.info(f"Link worker #{worker_id} started.")
    while True:
        task = await TASK_QUEUE.get()
        batch_id, original_link = task["batch_id"], task["original_link"]
        batch_info = BATCH_JOBS.get(batch_id)

        if not batch_info:
            TASK_QUEUE.task_done()
            continue

        file_info = None
        error_msg = f"An unknown error occurred for link: {original_link}"

        try:
            short_id = await terabox.extract_terabox_short_id(original_link)
            if not short_id:
                raise ValueError("Invalid Terabox link format.")

            download_url, raw_filename, error = await terabox.get_final_url_and_filename(original_link)
            if error:
                raise ValueError(error)

            local_filepath, _, download_error = await terabox.download_terabox_file(
                bot,
                batch_info["source_chat_id"],
                batch_info["status_message"].message_id,
                download_url,
                raw_filename
            )
            if download_error:
                raise ValueError(download_error)

            file_info = {
                'path': local_filepath,
                'name': raw_filename
            }

        except Exception as e:
            error_msg = str(e)

        async with batch_info["lock"]:
            batch_info["processed_links"] += 1
            if file_info:
                batch_info["successful_downloads"].append(file_info)
            else:
                batch_info["failed_links"].append({"link": original_link, "error": error_msg})

            processed, total = batch_info['processed_links'], batch_info['total_links']
            try:
                await bot.edit_message_text(
                    chat_id=batch_info["source_chat_id"],
                    message_id=batch_info["status_message"].message_id,
                    text=f"⚙️ Batch `{batch_id[:6]}` in progress... Processed {processed}/{total} links."
                )
            except TelegramBadRequest:
                pass

            if processed == total:
                logger.info(f"Batch {batch_id[:6]} complete. Triggering final processing.")
                await handle_batch_completion(batch_id)

        TASK_QUEUE.task_done()

async def send_file(chat_id: int, file_info: dict, caption: str) -> Message:
    file_path = FSInputFile(file_info['path'], filename=file_info['name'])
    sent_message = await bot.send_document(chat_id, file_path, caption=caption)
    return sent_message

async def handle_batch_completion(batch_id: str):
    batch = BATCH_JOBS.get(batch_id)
    if not batch:
        return

    status_msg = batch["status_message"]
    successful_downloads = batch["successful_downloads"]

    try:
        if not successful_downloads:
            failed_links_text = "\n".join(
                [f"- {x['link']} → {x['error']}" for x in batch['failed_links']]
            )
            await bot.edit_message_text(
                chat_id=batch["source_chat_id"],
                message_id=status_msg.message_id,
                text=f"❌ Batch `{batch_id[:6]}` failed. No files could be processed.\nDetails:\n{failed_links_text}"
            )
            return

        await bot.edit_message_text(
            chat_id=batch["source_chat_id"],
            message_id=status_msg.message_id,
            text=f"✅ Batch `{batch_id[:6]}` downloaded. Preparing to send {len(successful_downloads)} files..."
        )

        for item in successful_downloads:
            caption = f"`{item.get('name')}`"
            await send_file(batch["source_chat_id"], item, caption)

        summary = f"✅ Batch `{batch_id[:6]}` complete: {len(successful_downloads)} files sent."
        if batch["failed_links"]:
            summary += f"\n❌ {len(batch['failed_links'])} links failed."

        await bot.edit_message_text(
            chat_id=batch["source_chat_id"],
            message_id=status_msg.message_id,
            text=summary
        )

    except Exception as e:
        logger.error(f"Error during batch completion for {batch_id}: {e}", exc_info=True)
        await bot.edit_message_text(
            chat_id=batch["source_chat_id"],
            message_id=status_msg.message_id,
            text=f"A critical error occurred while sending files for batch `{batch_id[:6]}`."
        )
    finally:
        for item in successful_downloads:
            if os.path.exists(item['path']):
                os.remove(item['path'])

        del BATCH_JOBS[batch_id]

# --- Handlers ---
@dp.message(CommandStart())
async def start_handler(message: Message):
    text = (
        "👋 <b>Welcome to the Terabox Downloader Bot!</b>\n\n"
        "📥 Send me any valid Terabox link and I will fetch the file and send it to you.\n\n"
        f"📢 Please make sure you are a member of: @{config.FORCE_SUB_CHANNEL_USERNAME}\n\n"
        "✅ <i>Just send your link below ⬇️</i>"
    )
    await message.reply(text)

@dp.message(F.text | F.caption)
async def message_handler(message: Message):
    if message.text and message.text.startswith('/'):
        return

    await db_utils.add_or_update_user_db(
        message.from_user.id,
        message.from_user.username,
        message.from_user.first_name
    )

    links = list(set(re.findall(r'https?://[^\s<>"\']+', message.text or message.caption or "")))
    terabox_links = [link for link in links if any(domain in link for domain in [
        "terabox.com", "teraboxapp.com", "terasharelink.com", "1024tera.com",
        "freeterabox.com", "4funbox.com", "box-links.com"
    ])]

    if not terabox_links:
        return

    batch_id = str(uuid.uuid4())
    status_msg = await bot.send_message(
        chat_id=message.chat.id,
        text=f"✅ Found {len(terabox_links)} links. Queued as batch `{batch_id[:6]}`."
    )

    BATCH_JOBS[batch_id] = {
        "total_links": len(terabox_links),
        "processed_links": 0,
        "successful_downloads": [],
        "failed_links": [],
        "source_chat_id": message.chat.id,
        "source_user_id": message.from_user.id,
        "source_message_id": message.message_id,
        "status_message": status_msg,
        "lock": asyncio.Lock()
    }

    for link in terabox_links:
        await TASK_QUEUE.put({"batch_id": batch_id, "original_link": link})

# --- Startup ---
async def on_startup(dispatcher: Dispatcher):
    await db_utils.initialize_database()
    for i in range(config.CONCURRENT_WORKERS):
        asyncio.create_task(link_processor_worker(i + 1))
    logger.info(f"Bot started with {config.CONCURRENT_WORKERS} concurrent workers.")

def start_bot():
    dp.startup.register(on_startup)
    dp.message.register(message_handler)
    dp.message.register(start_handler, CommandStart())
    return dp, bot