bot.py

# bot.py
import asyncio
import logging
import os
import re
import uuid
import time
from typing import Dict

from aiogram import Bot, Dispatcher, types, 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(__name__)

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] = {}
START_TIME = time.time()

# --- Utils ---
def get_uptime():
    return round(time.time() - START_TIME)

# --- Worker ---
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"Unknown error 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}.")
                file_info = {'cached': True, **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 = batch_info['processed_links']
            total = batch_info['total_links']
            eta = estimate_eta(batch_info)

            try:
                await batch_info["status_message"].edit_text(
                    f"⚙️ Batch `{batch_id[:6]}` processing... {processed}/{total} links.\nETA: {eta}s"
                )
            except TelegramBadRequest:
                pass

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

        TASK_QUEUE.task_done()

# --- ETA ---
def estimate_eta(batch_info):
    elapsed = time.time() - batch_info["start_time"]
    processed = batch_info["processed_links"]
    total = batch_info["total_links"]

    if processed == 0:
        return "calculating..."

    avg_time_per_link = elapsed / processed
    remaining_links = total - processed
    eta = int(avg_time_per_link * remaining_links)
    return eta

# --- Send file ---
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'])
    filename = file_info.get('name') or file_info.get('filename')
    sent_message = None

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

    if 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 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

# --- Batch complete ---
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 `{batch_id[:6]}` failed. No files processed.\nDetails:\n{failed_links_text}"
            )
            return

        await status_msg.edit_text(
            f"✅ Batch `{batch_id[:6]}` downloaded. Sending {len(successful_downloads)} 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 `{batch_id[:6]}` complete: {len(successful_downloads)} 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 in batch {batch_id}: {e}", exc_info=True)
        await status_msg.edit_text(
            f"A critical error occurred while sending files for batch `{batch_id[:6]}`."
        )
    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]

# --- Message ---
@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"✅ 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(),
        "start_time": time.time()
    }

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

@dp.message(CommandStart())
async def start_handler(message: Message):
    text = (
        "👋 Welcome to <b>Terabox Downloader Bot</b>!\n\n"
        "✅ Send me <b>Terabox links</b>, I will download and send you the files.\n"
        "✅ Files also forwarded to our channel.\n"
        "✅ Progress + ETA shown.\n\n"
        f"Bot uptime: {get_uptime()} sec 🚀"
    )
    await message.reply(text)

async def run_bot():
    dp.startup.register(on_startup)
    dp.message.register(message_handler)
    dp.message.register(start_handler, CommandStart())
    await dp.start_polling(bot)