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,
        link_preview_is_disabled=True
    )
)

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.")

            cached_file = await db_utils.get_cached_file(short_id)
            if cached_file:
                logger.info(f"Cache hit for {short_id}. Reusing file_id.")
                file_info = {'cached': True, 'short_id': short_id, **cached_file}
            else:
                download_url, raw_filename, error = await terabox.get_final_url_and_filename(original_link)
                if error:
                    raise ValueError(error)

                local_filepath, thumb_path, 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 = {
                    'cached': False,
                    'path': local_filepath,
                    'name': raw_filename,
                    'size': os.path.getsize(local_filepath),
                    'short_id': short_id,
                    'thumb': thumb_path
                }

        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 batch_info["status_message"].edit_text(
                    f"⚙️ Batch <code>{batch_id[:6]}</code> in progress...\n"
                    f"Processed <b>{processed}</b>/<b>{total}</b> 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_and_cache_file(chat_id: int, file_info: dict, caption: str) -> Message:
    file_path_or_id = file_info.get('file_id') or FSInputFile(file_info['path'], filename=file_info['name'])
    media_type = file_info.get('type')
    filename = file_info.get('name') or file_info.get('filename')

    video_exts = ('.mp4', '.mkv', '.mov', '.avi', '.webm')
    audio_exts = ('.mp3', '.flac', '.ogg', '.wav')

    sent_message = None
    if media_type == 'video' or filename.lower().endswith(video_exts):
        sent_message = await bot.send_video(chat_id, file_path_or_id, caption=caption, supports_streaming=True)
        media_type = 'video'
    elif media_type == 'audio' or filename.lower().endswith(audio_exts):
        sent_message = await bot.send_audio(chat_id, file_path_or_id, caption=caption)
        media_type = 'audio'
    else:
        sent_message = await bot.send_document(chat_id, file_path_or_id, caption=caption)
        media_type = 'document'

    if not file_info.get('cached') and sent_message:
        file_id_to_cache = getattr(sent_message, media_type).file_id
        await db_utils.add_to_cache(
            file_info['short_id'], file_id_to_cache, filename, media_type, file_info['size']
        )

    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 status_msg.edit_text(
                f"❌ Batch <code>{batch_id[:6]}</code> failed.\nNo files could be processed.\n\n<b>Details:</b>\n{failed_links_text}"
            )
            return

        await status_msg.edit_text(
            f"✅ Batch <code>{batch_id[:6]}</code> downloaded.\nPreparing to send <b>{len(successful_downloads)}</b> files..."
        )

        if config.FORWARD_CHANNEL_ID:
            await bot.forward_message(
                config.FORWARD_CHANNEL_ID,
                batch["source_chat_id"],
                batch["source_message_id"]
            )
            for item in successful_downloads:
                caption = f"`{item.get('name') or item.get('filename')}`"
                await send_and_cache_file(config.FORWARD_CHANNEL_ID, item, caption)
                await asyncio.sleep(1)

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

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

        await status_msg.edit_text(summary)

    except Exception as e:
        logger.error(f"Error during batch completion for {batch_id}: {e}", exc_info=True)
        await status_msg.edit_text(
            f"❌ Critical error occurred while sending files for batch <code>{batch_id[:6]}</code>."
        )
    finally:
        for item in successful_downloads:
            if not item.get('cached') and os.path.exists(item['path']):
                os.remove(item['path'])
            if item.get('thumb') and os.path.exists(item['thumb']):
                os.remove(item['thumb'])

        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(s) and I will fetch the files and send them to you.\n\n"
        f"📢 Make sure you are a member of @{config.FORCE_SUB_CHANNEL_USERNAME}\n\n"
        "🚀 Supports multiple links per message (batch mode).\n"
        "💾 Auto-cache enabled.\n"
        "🔄 Forwarding original message to channel enabled.\n\n"
        "✅ <i>Just send your links 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 message.reply(
        f"✅ <b>Found</b> {len(terabox_links)} links.\n<b>Queued</b> as batch <code>{batch_id[:6]}</code>."
    )

    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