src/bot.py

#!/usr/bin/env python3
"""
AIGC Telegram Bot — AI Cover Generation Pipeline for Telegram

Features:
    - Upload audio or provide YouTube URL to generate AI voice covers
    - Choose from multiple RVC voice models
    - Customize pitch, index rate, f0 method, reverb, and more
    - Download generated covers directly in Telegram
    - Admin controls for managing models and users

Usage:
    python -m src.bot
"""

import asyncio
import html
import json
import logging
import os
import signal
import sys
import time
import traceback
from pathlib import Path
from typing import Optional

from telegram import (
    Update,
    BotCommand,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
)
from telegram.ext import (
    Application,
    ApplicationBuilder,
    CallbackQueryHandler,
    CommandHandler,
    ContextTypes,
    ConversationHandler,
    MessageHandler,
    filters,
)

from .config import load_config, load_pipeline_config, get_path_config, BotConfig, PipelineConfig, PathConfig
from .state import state_manager, UserState, InputType, UserSession
from .pipeline import PipelineRunner

# ============================================================
# Setup logging
# ============================================================

logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("aigc_bot")


def setup_logging(config: BotConfig):
    """Configure logging based on bot config."""
    root_logger = logging.getLogger()
    root_logger.setLevel(getattr(logging, config.log_level.upper(), logging.INFO))

    if config.log_file:
        file_handler = logging.FileHandler(config.log_file)
        file_handler.setFormatter(
            logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
        )
        root_logger.addHandler(file_handler)


# ============================================================
# Constants
# ============================================================

CONVO_STATES = {
    "INPUT_TYPE": 1,
    "SONG_INPUT": 2,
    "VOICE_MODEL": 3,
    "CONFIRM": 4,
    "SETTINGS_MENU": 5,
    "PITCH_INPUT": 6,
    "MODE_SELECT": 7,
}

# Callback data prefixes
CB_PREFIX = {
    "INPUT_YT": "input_yt",
    "INPUT_FILE": "input_file",
    "MODEL_SELECT": "mdl_",
    "CONFIRM_START": "confirm_start",
    "CANCEL": "cancel",
    "SETTINGS": "settings",
    "SET_PITCH": "set_pitch",
    "SET_MODE": "set_mode",
    "SET_FORMAT": "set_fmt_",
    "SET_F0": "set_f0_",
    "SET_DENOISE": "set_denoise",
    "SET_KEEP": "set_keep",
    "SET_REVERB": "set_reverb",
    "BACK_MENU": "back_menu",
}


# ============================================================
# AIGC Telegram Bot Class
# ============================================================

class AIGCTelegramBot:
    """
    Main Telegram Bot application for AI Cover Generation.
    """

    def __init__(self):
        self.config: BotConfig = load_config()
        self.pipeline_config: PipelineConfig = load_pipeline_config()
        self.paths: PathConfig = get_path_config()
        self.pipeline: Optional[PipelineRunner] = None
        self.application: Optional[Application] = None
        self._progress_messages: dict = {}  # user_id -> last progress message

        setup_logging(self.config)

        # Validate config
        if not self.config.bot_token:
            logger.error("BOT_TOKEN is not set! Please configure it in .env file.")
            sys.exit(1)

    async def initialize(self):
        """Initialize the bot and pipeline."""
        logger.info("Initializing AIGC Telegram Bot...")

        # Initialize pipeline runner
        self.pipeline = PipelineRunner(
            models_dir=str(self.paths.models_dir),
            output_dir=str(self.paths.output_dir),
            temp_dir=str(self.paths.temp_dir),
            mdx_models_dir=str(self.paths.mdx_models_dir),
            max_workers=self.pipeline_config.max_concurrent_jobs,
        )
        logger.info(f"Pipeline initialized with {self.pipeline_config.max_concurrent_jobs} workers")

        # Build application
        self.application = (
            ApplicationBuilder()
            .token(self.config.bot_token)
            .concurrent_updates(True)
            .build()
        )

        # Register handlers
        self._register_handlers()

        # Register bot commands
        await self._register_commands()

        logger.info(f"Bot '{self.config.bot_name}' initialized successfully")

    def _register_handlers(self):
        """Register all Telegram bot handlers."""
        app = self.application

        # Start handler
        app.add_handler(CommandHandler("start", self.cmd_start))

        # Help handler
        app.add_handler(CommandHandler("help", self.cmd_help))

        # Cancel handler
        app.add_handler(CommandHandler("cancel", self.cmd_cancel))

        # Status handler
        app.add_handler(CommandHandler("status", self.cmd_status))

        # Models list handler
        app.add_handler(CommandHandler("models", self.cmd_models))

        # Settings handler
        app.add_handler(CommandHandler("settings", self.cmd_settings))

        # Admin commands
        app.add_handler(CommandHandler("admin", self.cmd_admin))

        # Cover generation conversation handler
        cover_handler = ConversationHandler(
            entry_points=[
                CommandHandler("cover", self.cmd_cover),
                CommandHandler("newcover", self.cmd_cover),
            ],
            states={
                CONVO_STATES["INPUT_TYPE"]: [
                    CallbackQueryHandler(self.on_input_type_select, pattern=f"^({'|'.join([CB_PREFIX['INPUT_YT'], CB_PREFIX['INPUT_FILE'], CB_PREFIX['CANCEL']])})$"),
                ],
                CONVO_STATES["SONG_INPUT"]: [
                    MessageHandler(filters.TEXT & ~filters.COMMAND, self.on_song_url_input),
                    MessageHandler(filters.AUDIO | filters.VOICE, self.on_song_file_input),
                    CommandHandler("cancel", self.cmd_cancel),
                ],
                CONVO_STATES["VOICE_MODEL"]: [
                    CallbackQueryHandler(self.on_model_select, pattern=f"^{CB_PREFIX['MODEL_SELECT']}"),
                    CallbackQueryHandler(self.on_cancel, pattern=f"^{CB_PREFIX['CANCEL']}$"),
                ],
                CONVO_STATES["CONFIRM"]: [
                    CallbackQueryHandler(self.on_confirm_start, pattern=f"^{CB_PREFIX['CONFIRM_START']}$"),
                    CallbackQueryHandler(self.on_settings_from_confirm, pattern=f"^{CB_PREFIX['SETTINGS']}$"),
                    CallbackQueryHandler(self.on_cancel, pattern=f"^{CB_PREFIX['CANCEL']}$"),
                ],
            },
            fallbacks=[
                CommandHandler("cancel", self.cmd_cancel),
                CommandHandler("start", self.cmd_start),
            ],
            allow_reentry=True,
        )
        app.add_handler(cover_handler)

        # Settings conversation handler
        settings_handler = ConversationHandler(
            entry_points=[
                CallbackQueryHandler(self.on_settings_callback, pattern=f"^{CB_PREFIX['SET_PITCH']}|{CB_PREFIX['SET_MODE']}|{CB_PREFIX['SET_FORMAT']}|{CB_PREFIX['SET_F0']}|{CB_PREFIX['SET_DENOISE']}|{CB_PREFIX['SET_KEEP']}|{CB_PREFIX['SET_REVERB']}|{CB_PREFIX['BACK_MENU']}$"),
            ],
            states={
                CONVO_STATES["PITCH_INPUT"]: [
                    MessageHandler(filters.TEXT & ~filters.COMMAND, self.on_pitch_input),
                ],
                CONVO_STATES["MODE_SELECT"]: [
                    CallbackQueryHandler(self.on_mode_select),
                ],
            },
            fallbacks=[
                CommandHandler("cancel", self.cmd_cancel),
            ],
            allow_reentry=True,
        )
        app.add_handler(settings_handler)

        # Error handler
        app.add_error_handler(self.error_handler)

        # Log all messages for debugging
        if self.config.log_level == "DEBUG":
            app.add_handler(MessageHandler(filters.ALL, self._log_message), group=-1)

    async def _register_commands(self):
        """Register bot commands with Telegram."""
        commands = [
            BotCommand("start", "Start the bot and see welcome message"),
            BotCommand("cover", "Generate a new AI cover"),
            BotCommand("newcover", "Start a new cover generation"),
            BotCommand("models", "List available voice models"),
            BotCommand("settings", "Adjust generation settings"),
            BotCommand("status", "Check bot status and queue"),
            BotCommand("cancel", "Cancel current operation"),
            BotCommand("help", "Show help and usage guide"),
        ]
        await self.application.bot.set_my_commands(commands)

    # ============================================================
    # Auth & Access Control
    # ============================================================

    def _is_authorized(self, user_id: int) -> bool:
        """Check if a user is authorized to use the bot."""
        if self.config.allowed_user_ids is None:
            return True  # Public bot
        return user_id in self.config.allowed_user_ids

    def _is_admin(self, user_id: int) -> bool:
        """Check if a user is an admin."""
        return user_id in self.config.admin_ids

    # ============================================================
    # Command Handlers
    # ============================================================

    async def cmd_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /start command."""
        user = update.effective_user

        if not self._is_authorized(user.id):
            await update.message.reply_text(
                "Sorry, this bot is private and you don't have access.\n"
                "Contact the bot administrator for access."
            )
            return ConversationHandler.END

        state_manager.clear_session(user.id)

        welcome_text = (
            f"Hi <b>{html.escape(user.first_name)}</b>! Welcome to <b>{html.escape(self.config.bot_name)}</b>\n\n"
            "I can generate AI voice covers of your favorite songs using RVC voice models.\n\n"
            "<b>How to use:</b>\n"
            "1. Send /cover to start\n"
            "2. Choose to upload audio or paste a YouTube link\n"
            "3. Select a voice model\n"
            "4. Adjust settings (optional)\n"
            "5. Confirm and wait for your cover!\n\n"
            "<b>Commands:</b>\n"
            "/cover - Start cover generation\n"
            "/models - List available voice models\n"
            "/settings - Adjust generation settings\n"
            "/status - Check bot status\n"
            "/help - Show detailed help"
        )

        keyboard = InlineKeyboardMarkup([
            [InlineKeyboardButton("Generate Cover", callback_data="start_cover")],
            [InlineKeyboardButton("View Models", callback_data="view_models")],
            [InlineKeyboardButton("Settings", callback_data="settings")],
        ])

        await update.message.reply_text(
            welcome_text,
            parse_mode="HTML",
            reply_markup=keyboard,
        )

        # Handle inline button presses for start menu
        # These are captured by the general callback handler below
        return ConversationHandler.END

    async def cmd_help(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /help command with detailed usage guide."""
        help_text = (
            "<b>AIGC Cover Bot - Help Guide</b>\n\n"
            "<b>Quick Start:</b>\n"
            "1. Send <code>/cover</code> to begin\n"
            "2. Choose input method (YouTube URL or audio file upload)\n"
            "3. For YouTube: paste the video link\n"
            "   For file: upload an MP3, WAV, or other audio file\n"
            "4. Select a voice model from the list\n"
            "5. Review settings and tap 'Generate Cover'\n\n"
            "<b>Voice Models:</b>\n"
            "Voice models determine how the AI will sound. Each model is trained "
            "on a different voice. Use /models to see all available models.\n\n"
            "<b>Settings:</b>\n"
            "/settings lets you customize:\n"
            "- <b>Pitch Change:</b> Shift the AI voice pitch in octaves. "
            "Use +1 for male-to-female, -1 for female-to-male.\n"
            "- <b>Pitch Change All:</b> Shift pitch of all audio (vocals + instruments).\n"
            "- <b>Index Rate:</b> Timbre similarity (0-1). Higher = more like training data.\n"
            "- <b>F0 Method:</b> Pitch detection algorithm. rmvpe = best clarity.\n"
            "- <b>Reverb:</b> Add room reverb to the AI vocals.\n"
            "- <b>Denoise:</b> Apply extra noise reduction.\n"
            "- <b>Output Format:</b> MP3 (smaller) or WAV (best quality).\n\n"
            "<b>Inference Modes:</b>\n"
            "- <b>Full:</b> Separate vocals + RVC voice conversion + mix\n"
            "- <b>RVC Only:</b> Voice conversion only (skip vocal separation)\n"
            "- <b>MDX Only:</b> Vocal separation only (skip voice conversion)\n\n"
            "<b>Tips:</b>\n"
            "- YouTube links must be direct video URLs (not playlists)\n"
            "- Audio files up to 20MB are supported\n"
            "- Processing time depends on song length (usually 1-3 minutes)\n"
            "- Use /cancel to abort at any time\n\n"
            "<b>Admin Commands:</b>\n"
            "/admin stats - View bot usage statistics\n"
            "/admin cleanup - Clean up old temp files"
        )

        await update.message.reply_text(help_text, parse_mode="HTML")

    async def cmd_cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /cancel command."""
        user = update.effective_user
        state_manager.clear_session(user.id)
        state_manager.set_state(user.id, UserState.IDLE)

        await update.message.reply_text(
            "Operation cancelled. Send /cover to start a new one."
        )
        return ConversationHandler.END

    async def cmd_models(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /models command - list available voice models."""
        if not self.pipeline:
            await update.message.reply_text("Pipeline is not initialized yet.")
            return

        models = self.pipeline.get_available_models()

        if not models:
            await update.message.reply_text(
                "No voice models available. Please add models to the "
                f"<code>{self.paths.models_dir}</code> directory.\n\n"
                "Each model should be in its own folder containing a <code>.pth</code> file "
                "(and optionally a <code>.index</code> file).",
                parse_mode="HTML",
            )
            return

        models_text = (
            f"<b>Available Voice Models</b> ({len(models)}):\n\n"
        )
        for i, model_name in enumerate(models, 1):
            models_text += f"  {i}. <code>{html.escape(model_name)}</code>\n"

        models_text += (
            "\nSend /cover to use one of these models."
        )

        await update.message.reply_text(models_text, parse_mode="HTML")

    async def cmd_status(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /status command."""
        active_count = state_manager.get_active_count()
        models_count = len(self.pipeline.get_available_models()) if self.pipeline else 0

        status_text = (
            f"<b>{html.escape(self.config.bot_name)} - Status</b>\n\n"
            f"Active jobs: {active_count}/{self.pipeline_config.max_concurrent_jobs}\n"
            f"Available models: {models_count}\n"
            f"Output format: {self.pipeline_config.default_output_format}\n"
            f"Inference mode: {self.pipeline_config.default_inference_mode}\n"
        )

        session = state_manager.get_session(update.effective_user.id)
        status_text += f"\n<b>Your state:</b> {session.state.value}"

        await update.message.reply_text(status_text, parse_mode="HTML")

    async def cmd_settings(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /settings command."""
        session = state_manager.get_session(update.effective_user.id)
        await self._send_settings_menu(update.effective_chat.id, session, context)
        return ConversationHandler.END

    async def cmd_admin(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /admin command."""
        user = update.effective_user

        if not self._is_admin(user.id):
            await update.message.reply_text("You don't have admin access.")
            return

        if not context.args:
            await update.message.reply_text(
                "<b>Admin Commands:</b>\n"
                "/admin stats - View statistics\n"
                "/admin cleanup - Clean temp files\n"
                "/admin broadcast <msg> - Send message to all users",
                parse_mode="HTML",
            )
            return

        action = context.args[0].lower()

        if action == "stats":
            models = self.pipeline.get_available_models() if self.pipeline else []
            await update.message.reply_text(
                f"<b>Bot Stats</b>\n"
                f"Models: {len(models)}\n"
                f"Active jobs: {state_manager.get_active_count()}\n"
                f"Max concurrent: {self.pipeline_config.max_concurrent_jobs}",
                parse_mode="HTML",
            )

        elif action == "cleanup":
            if self.pipeline:
                self.pipeline.cleanup_temp_files()
                await update.message.reply_text("Temp files cleaned up.")
            else:
                await update.message.reply_text("Pipeline not initialized.")

        elif action == "broadcast" and len(context.args) > 1:
            msg = " ".join(context.args[1:])
            await update.message.reply_text("Broadcast feature not yet implemented.")

        else:
            await update.message.reply_text("Unknown admin command.")

    # ============================================================
    # Cover Generation Flow
    # ============================================================

    async def cmd_cover(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle /cover command - start cover generation flow."""
        user = update.effective_user

        if not self._is_authorized(user.id):
            await update.message.reply_text("You don't have access to this bot.")
            return ConversationHandler.END

        if state_manager.is_processing(user.id):
            await update.message.reply_text(
                "You already have a cover being generated. "
                "Please wait or use /cancel to abort."
            )
            return ConversationHandler.END

        state_manager.clear_session(user.id)
        session = state_manager.get_session(user.id)

        # Load defaults from config
        session.pitch_change = self.pipeline_config.default_pitch_change
        session.index_rate = self.pipeline_config.default_index_rate
        session.filter_radius = self.pipeline_config.default_filter_radius
        session.volume_envelope = self.pipeline_config.default_volume_envelope
        session.f0_method = self.pipeline_config.default_f0_method
        session.hop_length = self.pipeline_config.default_hop_length
        session.protect = self.pipeline_config.default_protect
        session.reverb_size = self.pipeline_config.default_reverb_size
        session.reverb_wet = self.pipeline_config.default_reverb_wet
        session.reverb_dry = self.pipeline_config.default_reverb_dry
        session.reverb_damping = self.pipeline_config.default_reverb_damping
        session.main_gain = self.pipeline_config.default_main_gain
        session.backup_gain = self.pipeline_config.default_backup_gain
        session.inst_gain = self.pipeline_config.default_inst_gain
        session.output_format = self.pipeline_config.default_output_format
        session.inference_mode = self.pipeline_config.default_inference_mode

        # Step 1: Ask for input type
        keyboard = InlineKeyboardMarkup([
            [
                InlineKeyboardButton(
                    "YouTube URL",
                    callback_data=CB_PREFIX["INPUT_YT"]
                ),
                InlineKeyboardButton(
                    "Upload Audio File",
                    callback_data=CB_PREFIX["INPUT_FILE"]
                ),
            ],
            [InlineKeyboardButton("Cancel", callback_data=CB_PREFIX["CANCEL"])],
        ])

        await update.message.reply_text(
            "<b>Step 1/3: Choose Input Method</b>\n\n"
            "How would you like to provide the song?",
            parse_mode="HTML",
            reply_markup=keyboard,
        )

        return CONVO_STATES["INPUT_TYPE"]

    async def on_input_type_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle input type selection (YouTube or File)."""
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        if query.data == CB_PREFIX["INPUT_YT"]:
            session.input_type = InputType.YOUTUBE
            state_manager.set_state(user_id, UserState.AWAITING_YOUTUBE_URL)

            await query.edit_message_text(
                "<b>Step 2/3: Provide Song</b>\n\n"
                "Send me a YouTube video URL.\n"
                "Example: <code>https://www.youtube.com/watch?v=dQw4w9WgXcQ</code>\n\n"
                "<b>Note:</b> Only direct video URLs are supported (no playlists).",
                parse_mode="HTML",
            )
            return CONVO_STATES["SONG_INPUT"]

        elif query.data == CB_PREFIX["INPUT_FILE"]:
            session.input_type = InputType.FILE
            state_manager.set_state(user_id, UserState.AWAITING_SONG_INPUT)

            await query.edit_message_text(
                "<b>Step 2/3: Provide Song</b>\n\n"
                "Upload an audio file (MP3, WAV, OGG, FLAC).\n"
                f"Maximum file size: {self.pipeline_config.max_file_size_mb}MB\n\n"
                "Send the audio file now:",
                parse_mode="HTML",
            )
            return CONVO_STATES["SONG_INPUT"]

        elif query.data == CB_PREFIX["CANCEL"]:
            state_manager.clear_session(user_id)
            await query.edit_message_text("Cancelled. Send /cover to start again.")
            return ConversationHandler.END

    async def on_song_url_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle YouTube URL input from user."""
        user_id = update.effective_user.id
        session = state_manager.get_session(user_id)
        url = update.message.text.strip()

        # Validate YouTube URL
        if not self._is_valid_youtube_url(url):
            await update.message.reply_text(
                "That doesn't look like a valid YouTube URL.\n"
                "Please send a direct YouTube video link like:\n"
                "<code>https://www.youtube.com/watch?v=...</code>",
                parse_mode="HTML",
            )
            return CONVO_STATES["SONG_INPUT"]

        session.song_input = url
        await self._show_model_selection(update.effective_chat.id, session, context)
        return CONVO_STATES["VOICE_MODEL"]

    async def on_song_file_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle audio file upload from user."""
        user_id = update.effective_user.id
        session = state_manager.get_session(user_id)

        # Get the audio file
        audio = update.message.audio or update.message.voice
        if not audio:
            await update.message.reply_text("Please send an audio file.")
            return CONVO_STATES["SONG_INPUT"]

        # Check file size
        file_size_mb = audio.file_size / (1024 * 1024) if audio.file_size else 0
        if file_size_mb > self.pipeline_config.max_file_size_mb:
            await update.message.reply_text(
                f"File too large ({file_size_mb:.1f}MB). "
                f"Maximum is {self.pipeline_config.max_file_size_mb}MB."
            )
            return CONVO_STATES["SONG_INPUT"]

        # Download the file to temp directory
        status_msg = await update.message.reply_text("Downloading audio file...")

        try:
            file = await audio.get_file()
            temp_filename = f"input_{user_id}_{int(time.time())}.ogg"
            temp_path = str(self.paths.temp_dir / temp_filename)
            await file.download_to_drive(temp_path)
            session.song_input = temp_path
            session.song_file_id = audio.file_id

            await status_msg.edit_text("Audio file downloaded. Choose a voice model:")
        except Exception as e:
            logger.error(f"Failed to download audio: {e}")
            await status_msg.edit_text(f"Failed to download audio file: {html.escape(str(e))}")
            return ConversationHandler.END

        await self._show_model_selection(update.effective_chat.id, session, context)
        return CONVO_STATES["VOICE_MODEL"]

    async def _show_model_selection(self, chat_id: int, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
        """Display voice model selection keyboard."""
        models = self.pipeline.get_available_models() if self.pipeline else []

        if not models:
            await context.bot.send_message(
                chat_id,
                "No voice models available. Please add models to the "
                f"<code>{self.paths.models_dir}</code> directory and try again.",
                parse_mode="HTML",
            )
            return ConversationHandler.END

        # Build inline keyboard with models (paginated)
        buttons = []
        for model in models:
            buttons.append([InlineKeyboardButton(
                model,
                callback_data=f"{CB_PREFIX['MODEL_SELECT']}{model}"
            )])
        buttons.append([InlineKeyboardButton("Cancel", callback_data=CB_PREFIX["CANCEL"])])

        keyboard = InlineKeyboardMarkup(buttons)

        text = (
            f"<b>Step 3/3: Select Voice Model</b>\n\n"
            f"Choose a voice model for the cover ({len(models)} available):\n"
        )
        if session.pitch_change != 0:
            text += f"\nCurrent pitch change: <b>{session.pitch_change}</b> octaves"

        await context.bot.send_message(chat_id, text, parse_mode="HTML", reply_markup=keyboard)

    async def on_model_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle voice model selection."""
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        model_name = query.data.replace(CB_PREFIX["MODEL_SELECT"], "")
        session.voice_model = model_name

        # Show confirmation with current settings
        await self._show_confirmation(query, session, context)
        return CONVO_STATES["CONFIRM"]

    async def _show_confirmation(self, query_or_msg, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
        """Show generation confirmation with settings summary."""
        pitch_display = f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)

        mode_labels = {
            "full": "Full (Separate + Convert + Mix)",
            "mdx": "MDX Only (Vocal Separation)",
            "rvc": "RVC Only (Voice Conversion)",
        }

        input_display = (
            "YouTube URL" if session.input_type == InputType.YOUTUBE
            else "Audio File Upload"
        )

        song_display = session.song_input
        if session.input_type == InputType.YOUTUBE:
            if len(song_display) > 50:
                song_display = song_display[:50] + "..."

        text = (
            "<b>Confirm Cover Generation</b>\n\n"
            f"<b>Song:</b> <code>{html.escape(song_display)}</code>\n"
            f"<b>Input:</b> {input_display}\n"
            f"<b>Voice Model:</b> <code>{html.escape(session.voice_model or 'None')}</code>\n\n"
            "<b>Settings:</b>\n"
            f"  Mode: {mode_labels.get(session.inference_mode, session.inference_mode)}\n"
            f"  Pitch: {pitch_display} octaves"
        )
        if session.pitch_change_all != 0:
            text += f" (All: +{session.pitch_change_all})"

        text += (
            f"\n  F0 Method: {session.f0_method}\n"
            f"  Index Rate: {session.index_rate}\n"
            f"  Filter Radius: {session.filter_radius}\n"
            f"  Protect: {session.protect}\n"
            f"  Output: {session.output_format.upper()}\n"
            f"  Extra Denoise: {'Yes' if session.extra_denoise else 'No'}\n"
        )

        keyboard = InlineKeyboardMarkup([
            [
                InlineKeyboardButton(
                    "Generate Cover",
                    callback_data=CB_PREFIX["CONFIRM_START"]
                ),
            ],
            [
                InlineKeyboardButton(
                    "Adjust Settings",
                    callback_data=CB_PREFIX["SETTINGS"]
                ),
                InlineKeyboardButton(
                    "Cancel",
                    callback_data=CB_PREFIX["CANCEL"]
                ),
            ],
        ])

        if hasattr(query_or_msg, "edit_message_text"):
            await query_or_msg.edit_message_text(text, parse_mode="HTML", reply_markup=keyboard)
        else:
            await context.bot.send_message(
                query_or_msg.chat_id, text, parse_mode="HTML", reply_markup=keyboard
            )

    async def on_confirm_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle confirmation to start generation."""
        query = update.callback_query
        await query.answer("Starting generation...")

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        # Mark user as processing
        state_manager.start_task(user_id, f"cover_{int(time.time())}")

        processing_msg = await query.edit_message_text(
            "Generating your AI cover...\n\n"
            "This may take a few minutes depending on the song length. "
            "I'll send you the result when it's ready!"
        )
        self._progress_messages[user_id] = processing_msg.message_id

        # Start generation in background
        asyncio.create_task(
            self._run_generation(user_id, query.effective_chat.id, session, context)
        )

        return ConversationHandler.END

    async def on_settings_from_confirm(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle settings button from confirmation screen."""
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        await self._send_settings_menu(query.effective_chat.id, session, context)
        # Stay in CONFIRM state so user can go back
        return CONVO_STATES["CONFIRM"]

    async def on_cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle cancel button."""
        query = update.callback_query
        await query.answer("Cancelled")
        user_id = query.from_user.id
        state_manager.clear_session(user_id)
        await query.edit_message_text(
            "Cancelled. Send /cover to start again."
        )
        return ConversationHandler.END

    # ============================================================
    # Settings Menus
    # ============================================================

    async def _send_settings_menu(self, chat_id: int, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
        """Send the settings menu."""
        pitch_display = f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)

        mode_labels = {
            "full": "Full Pipeline",
            "mdx": "MDX Only",
            "rvc": "RVC Only",
        }

        text = (
            "<b>Generation Settings</b>\n\n"
            f"<b>Voice Model:</b> <code>{html.escape(session.voice_model or 'Not selected')}</code>\n"
            f"<b>Mode:</b> {mode_labels.get(session.inference_mode, session.inference_mode)}\n"
            f"<b>Pitch Change:</b> {pitch_display} octaves\n"
            f"<b>Pitch All:</b> {'+' if session.pitch_change_all >= 0 else ''}{session.pitch_change_all}\n"
            f"<b>F0 Method:</b> {session.f0_method}\n"
            f"<b>Index Rate:</b> {session.index_rate}\n"
            f"<b>Filter Radius:</b> {session.filter_radius}\n"
            f"<b>Volume Envelope:</b> {session.volume_envelope}\n"
            f"<b>Protect:</b> {session.protect}\n"
            f"<b>Output Format:</b> {session.output_format.upper()}\n"
            f"<b>Reverb Size:</b> {session.reverb_size}\n"
            f"<b>Reverb Wetness:</b> {session.reverb_wet}\n"
            f"<b>Extra Denoise:</b> {'On' if session.extra_denoise else 'Off'}\n"
        )

        keyboard = InlineKeyboardMarkup([
            [
                InlineKeyboardButton("Change Pitch", callback_data=CB_PREFIX["SET_PITCH"]),
                InlineKeyboardButton("Change Mode", callback_data=CB_PREFIX["SET_MODE"]),
            ],
            [
                InlineKeyboardButton("Format: MP3" if session.output_format == "mp3" else "Format: WAV",
                                     callback_data=f"{CB_PREFIX['SET_FORMAT']}{'wav' if session.output_format == 'mp3' else 'mp3'}"),
                InlineKeyboardButton("Denoise: " + ("ON" if session.extra_denoise else "OFF"),
                                     callback_data=CB_PREFIX["SET_DENOISE"]),
            ],
            [
                InlineKeyboardButton("F0: rmvpe" if session.f0_method == "rmvpe" else f"F0: {session.f0_method}",
                                     callback_data=f"{CB_PREFIX['SET_F0']}{'mangio-crepe' if session.f0_method == 'rmvpe' else 'rmvpe'}"),
                InlineKeyboardButton("Reverb", callback_data=CB_PREFIX["SET_REVERB"]),
            ],
            [InlineKeyboardButton("Back", callback_data=CB_PREFIX["BACK_MENU"])],
        ])

        await context.bot.send_message(chat_id, text, parse_mode="HTML", reply_markup=keyboard)

    async def on_settings_callback(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle settings button callbacks."""
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        if query.data == CB_PREFIX["SET_PITCH"]:
            await query.edit_message_text(
                "Send the pitch change value in octaves.\n\n"
                "Examples:\n"
                "  <code>1</code> = one octave up (male to female)\n"
                "  <code>-1</code> = one octave down (female to male)\n"
                "  <code>0</code> = no change\n"
                "  <code>2</code> = two octaves up\n\n"
                "Current: " + (f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)) + " octaves",
                parse_mode="HTML",
            )
            return CONVO_STATES["PITCH_INPUT"]

        elif query.data == CB_PREFIX["SET_MODE"]:
            keyboard = InlineKeyboardMarkup([
                [
                    InlineKeyboardButton("Full Pipeline", callback_data="mode_full"),
                    InlineKeyboardButton("RVC Only", callback_data="mode_rvc"),
                ],
                [InlineKeyboardButton("MDX Only", callback_data="mode_mdx")],
            ])
            await query.edit_message_text(
                "<b>Select Inference Mode</b>\n\n"
                "<b>Full Pipeline:</b> Separate vocals + RVC conversion + mix\n"
                "<b>RVC Only:</b> Voice conversion only\n"
                "<b>MDX Only:</b> Vocal separation only",
                parse_mode="HTML",
                reply_markup=keyboard,
            )
            return CONVO_STATES["MODE_SELECT"]

        elif query.data.startswith(CB_PREFIX["SET_FORMAT"]):
            new_format = query.data.replace(CB_PREFIX["SET_FORMAT"], "")
            session.output_format = new_format
            await self._send_settings_menu(query.effective_chat.id, session, context)

        elif query.data.startswith(CB_PREFIX["SET_F0"]):
            new_f0 = query.data.replace(CB_PREFIX["SET_F0"], "")
            session.f0_method = new_f0
            await self._send_settings_menu(query.effective_chat.id, session, context)

        elif query.data == CB_PREFIX["SET_DENOISE"]:
            session.extra_denoise = not session.extra_denoise
            await self._send_settings_menu(query.effective_chat.id, session, context)

        elif query.data == CB_PREFIX["SET_REVERB"]:
            # Cycle through reverb presets: None -> Light -> Medium -> Heavy -> None
            if session.reverb_wet < 0.1:
                session.reverb_wet, session.reverb_size = 0.1, 0.1
                label = "Light"
            elif session.reverb_wet < 0.3:
                session.reverb_wet, session.reverb_size = 0.3, 0.2
                label = "Medium"
            elif session.reverb_wet < 0.5:
                session.reverb_wet, session.reverb_size = 0.5, 0.4
                label = "Heavy"
            else:
                session.reverb_wet, session.reverb_size = 0.0, 0.0
                label = "None"
            await self._send_settings_menu(query.effective_chat.id, session, context)

        elif query.data == CB_PREFIX["BACK_MENU"]:
            await query.edit_message_text("Settings saved. Send /cover to generate a cover.")

        return ConversationHandler.END

    async def on_pitch_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle pitch value input."""
        user_id = update.effective_user.id
        session = state_manager.get_session(user_id)

        try:
            pitch_value = int(update.message.text.strip())
        except ValueError:
            await update.message.reply_text(
                "Invalid value. Please enter an integer (e.g., 1, -1, 0, 2)."
            )
            return CONVO_STATES["PITCH_INPUT"]

        if pitch_value < -12 or pitch_value > 12:
            await update.message.reply_text(
                "Pitch must be between -12 and +12 octaves."
            )
            return CONVO_STATES["PITCH_INPUT"]

        session.pitch_change = pitch_value
        await update.message.reply_text(
            f"Pitch set to <b>{pitch_value}</b> octaves.",
            parse_mode="HTML",
        )
        await self._send_settings_menu(update.effective_chat.id, session, context)
        return ConversationHandler.END

    async def on_mode_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle inference mode selection."""
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        session = state_manager.get_session(user_id)

        if query.data == "mode_full":
            session.inference_mode = "full"
        elif query.data == "mode_rvc":
            session.inference_mode = "rvc"
        elif query.data == "mode_mdx":
            session.inference_mode = "mdx"

        await self._send_settings_menu(query.effective_chat.id, session, context)
        return ConversationHandler.END

    # ============================================================
    # Generation Runner
    # ============================================================

    async def _run_generation(
        self,
        user_id: int,
        chat_id: int,
        session: UserSession,
        context: ContextTypes.DEFAULT_TYPE,
    ):
        """Run the cover generation pipeline and send the result."""
        progress_msg_id = self._progress_messages.get(user_id)

        async def progress_callback(percent: float, desc: str = ""):
            """Update progress message."""
            if progress_msg_id and percent % 10 < 5:  # Update every ~10%
                try:
                    bar_len = 20
                    filled = int(bar_len * percent / 100)
                    bar = "█" * filled + "░" * (bar_len - filled)
                    text = (
                        f"<b>Generating AI Cover...</b>\n\n"
                        f"[{bar}] {percent:.0f}%\n"
                        f"{html.escape(desc or 'Processing...')}"
                    )
                    await context.bot.edit_message_text(
                        text,
                        chat_id=chat_id,
                        message_id=progress_msg_id,
                        parse_mode="HTML",
                    )
                except Exception:
                    pass  # Ignore edit failures (rate limits, etc.)

        try:
            result_path = await self.pipeline.process_cover(
                song_input=session.song_input,
                voice_model=session.voice_model,
                pitch_change=session.pitch_change,
                keep_files=session.keep_files,
                main_gain=session.main_gain,
                backup_gain=session.backup_gain,
                inst_gain=session.inst_gain,
                index_rate=session.index_rate,
                filter_radius=session.filter_radius,
                volume_envelope=session.volume_envelope,
                f0_method=session.f0_method,
                hop_length=session.hop_length,
                protect=session.protect,
                pitch_change_all=session.pitch_change_all,
                reverb_rm_size=session.reverb_size,
                reverb_wet=session.reverb_wet,
                reverb_dry=session.reverb_dry,
                reverb_damping=session.reverb_damping,
                output_format=session.output_format,
                extra_denoise=session.extra_denoise,
                inference_mode=session.inference_mode,
                progress_callback=progress_callback,
            )

            # Send the result
            if result_path and os.path.exists(result_path):
                file_size = os.path.getsize(result_path) / (1024 * 1024)

                # Check if file is within Telegram's 50MB limit
                if file_size > 50:
                    # File too large for Telegram - send note
                    await context.bot.send_message(
                        chat_id,
                        f"Cover generated but file is too large for Telegram ({file_size:.1f}MB).\n"
                        f"File saved at: <code>{html.escape(result_path)}</code>",
                        parse_mode="HTML",
                    )
                else:
                    # Get the base filename
                    filename = os.path.basename(result_path)
                    caption = (
                        f"Your AI cover is ready!\n\n"
                        f"Voice: <code>{html.escape(session.voice_model or 'Unknown')}</code>\n"
                        f"Pitch: {session.pitch_change}\n"
                        f"Format: {session.output_format.upper()}\n"
                    )

                    await context.bot.send_audio(
                        chat_id=chat_id,
                        audio=open(result_path, "rb"),
                        caption=caption,
                        parse_mode="HTML",
                    )
            else:
                await context.bot.send_message(
                    chat_id,
                    "Generation completed but no output file was produced. "
                    "This might indicate an issue with the input or model."
                )

        except Exception as e:
            error_text = f"Generation failed: {html.escape(str(e))}"
            logger.error(f"Generation failed for user {user_id}: {e}")
            logger.debug(traceback.format_exc())

            if progress_msg_id:
                try:
                    await context.bot.edit_message_text(
                        f"<b>Error</b>\n\n{error_text}\n\n"
                        "Try again with /cover or check your settings.",
                        chat_id=chat_id,
                        message_id=progress_msg_id,
                        parse_mode="HTML",
                    )
                except Exception:
                    await context.bot.send_message(chat_id, error_text, parse_mode="HTML")
            else:
                await context.bot.send_message(chat_id, error_text, parse_mode="HTML")

        finally:
            state_manager.end_task(user_id)
            self._progress_messages.pop(user_id, None)

            # Cleanup temp files
            if session.song_input and os.path.exists(session.song_input) and session.input_type == InputType.FILE:
                try:
                    os.remove(session.song_input)
                except Exception:
                    pass

    # ============================================================
    # Utility Methods
    # ============================================================

    @staticmethod
    def _is_valid_youtube_url(url: str) -> bool:
        """Validate a YouTube URL."""
        from urllib.parse import urlparse

        parsed = urlparse(url)
        valid_hosts = {
            "www.youtube.com",
            "youtube.com",
            "m.youtube.com",
            "youtu.be",
            "music.youtube.com",
        }

        if parsed.hostname in valid_hosts:
            if parsed.hostname == "youtu.be" and parsed.path[1:]:
                return True
            if "v=" in parsed.query:
                return True
            if parsed.path.startswith("/watch"):
                return True

        return False

    async def _log_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Log all incoming messages (debug mode only)."""
        logger.debug(f"Message from {update.effective_user.id}: {update.message.text or '[media]'}")

    # ============================================================
    # Error Handler
    # ============================================================

    async def error_handler(self, update: object, context: ContextTypes.DEFAULT_TYPE):
        """Handle all uncaught exceptions."""
        error = context.error
        tb = "".join(traceback.format_exception(type(error), error, error.__traceback__))

        logger.error(f"Unhandled exception: {error}\n{tb}")

        if update and hasattr(update, "effective_message") and update.effective_message:
            try:
                await update.effective_message.reply_text(
                    f"An unexpected error occurred: {html.escape(str(error))}\n"
                    "Please try again or contact the admin."
                )
            except Exception:
                pass

    # ============================================================
    # Run
    # ============================================================

    def run(self):
        """Start the bot (blocking)."""
        if not self.application:
            raise RuntimeError("Bot not initialized. Call initialize() first.")

        logger.info(f"Starting {self.config.bot_name}...")

        # Setup signal handlers for graceful shutdown
        def signal_handler(sig, frame):
            logger.info("Received shutdown signal")
            if self.pipeline:
                self.pipeline.shutdown()
            if self.application:
                asyncio.create_task(self.application.shutdown())
            sys.exit(0)

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

        # Start polling
        self.application.run_polling(
            drop_pending_updates=True,
            allowed_updates=Update.ALL_TYPES,
        )

    async def run_async(self):
        """Start the bot asynchronously."""
        if not self.application:
            raise RuntimeError("Bot not initialized. Call initialize() first.")

        await self.application.initialize()
        await self.application.start()
        await self.application.updater.start_polling(
            drop_pending_updates=True,
            allowed_updates=Update.ALL_TYPES,
        )

        logger.info(f"{self.config.bot_name} is running")


# ============================================================
# Entry Point
# ============================================================

def main():
    """Main entry point."""
    bot = AIGCTelegramBot()

    async def startup():
        await bot.initialize()
        await bot.run_async()

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        loop.run_until_complete(startup())
        loop.run_forever()
    except (KeyboardInterrupt, SystemExit):
        logger.info("Shutting down...")
    finally:
        if bot.pipeline:
            bot.pipeline.shutdown()
        loop.close()


if __name__ == "__main__":
    main()