bot.py

import os
import asyncio
import logging
import time
from collections import defaultdict

from telegram import (
    Update,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
)
from telegram.constants import ChatAction, ParseMode
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    CallbackQueryHandler,
    ContextTypes,
    filters,
)

from music_gen import generate_music

logger = logging.getLogger(__name__)

TELEGRAM_TOKEN     = os.getenv("TELEGRAM_BOT_TOKEN", "")
RATE_LIMIT_SECONDS = 90

user_last_gen  = defaultdict(float)
user_settings  = {}

DEFAULT_SETTINGS = {
    "duration":       8,
    "guidance_scale": 3.0,
    "temperature":    1.0,
}

EXAMPLES = [
    "upbeat electronic dance music with synth leads",
    "calm piano melody ambient relaxing",
    "rock guitar riff with heavy drums",
    "jazz saxophone solo with soft drums",
    "happy ukulele with clapping and whistling",
    "cinematic orchestral epic trailer music",
    "lo-fi hip hop beats to study to",
    "funky bass groove with brass section",
    "celtic folk music with flute and fiddle",
    "8-bit chiptune retro video game music",
]


# ── Helpers ──────────────────────────────────────────────────────────────────

def get_settings(uid):
    if uid not in user_settings:
        user_settings[uid] = DEFAULT_SETTINGS.copy()
    return user_settings[uid]


def check_rate_limit(uid):
    now     = time.time()
    elapsed = now - user_last_gen[uid]
    if elapsed < RATE_LIMIT_SECONDS:
        return False, int(RATE_LIMIT_SECONDS - elapsed)
    user_last_gen[uid] = now
    return True, 0


# ── Core generation ───────────────────────────────────────────────────────────

async def do_generate(chat_id, prompt, user_id, context, reply_id=None):

    allowed, wait = check_rate_limit(user_id)
    if not allowed:
        await context.bot.send_message(
            chat_id=chat_id,
            text=f"⏳ Please wait *{wait}s* before next generation.",
            parse_mode=ParseMode.MARKDOWN,
        )
        return

    s = get_settings(user_id)

    status = await context.bot.send_message(
        chat_id=chat_id,
        reply_to_message_id=reply_id,
        parse_mode=ParseMode.MARKDOWN,
        text=(
            f"🎵 *Generating your music...*\n\n"
            f"📝 _{prompt}_\n\n"
            f"⏱️ Duration    : {s['duration']}s\n"
            f"🎯 Guidance    : {s['guidance_scale']}\n"
            f"🌡️ Temperature : {s['temperature']}\n\n"
            f"⏳ CPU generation takes 30s–3min, please wait..."
        ),
    )

    await context.bot.send_chat_action(
        chat_id=chat_id,
        action=ChatAction.UPLOAD_VOICE,
    )

    try:
        loop = asyncio.get_event_loop()
        path = await loop.run_in_executor(
            None,
            generate_music,
            prompt,
            s["duration"],
            s["guidance_scale"],
            s["temperature"],
        )
    except Exception as e:
        await status.edit_text(
            f"❌ Generation failed:\n`{str(e)[:300]}`",
            parse_mode=ParseMode.MARKDOWN,
        )
        return

    try:
        with open(path, "rb") as f:
            await context.bot.send_audio(
                chat_id=chat_id,
                audio=f,
                title=f"AI Music — {prompt[:40]}",
                performer="MusicGen AI",
                filename="music.wav",
                reply_to_message_id=reply_id,
                caption=(
                    f"🎵 *Done!*\n"
                    f"📝 _{prompt[:100]}_\n"
                    f"⏱️ {s['duration']}s"
                ),
                parse_mode=ParseMode.MARKDOWN,
            )
        await status.edit_text("✅ Enjoy your music! 🎶")

    except Exception as e:
        await status.edit_text(
            f"❌ Could not send audio: {str(e)[:200]}"
        )

    finally:
        try:
            os.unlink(path)
        except Exception:
            pass


# ── Commands ──────────────────────────────────────────────────────────────────

async def cmd_start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        f"🎵 *Welcome {update.effective_user.first_name}!*\n\n"
        "Send me any music description and I will generate it.\n\n"
        "*Example:*\n"
        "_calm piano with rain sounds, lo-fi_\n\n"
        "/examples — see prompt ideas\n"
        "/settings — adjust duration and style\n"
        "/help     — full guide",
        parse_mode=ParseMode.MARKDOWN,
    )


async def cmd_help(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        "🎵 *Help*\n\n"
        "Just type a music description and I will generate it!\n\n"
        "*Good prompt tips:*\n"
        "• Instruments: _guitar, piano, drums, synth_\n"
        "• Mood: _happy, sad, epic, calm_\n"
        "• Genre: _jazz, rock, electronic, classical_\n"
        "• Tempo: _fast, slow, upbeat_\n\n"
        "⚠️ _Free CPU — generation takes 30s to 3min_",
        parse_mode=ParseMode.MARKDOWN,
    )


async def cmd_examples(update: Update, context: ContextTypes.DEFAULT_TYPE):
    buttons = [
        [InlineKeyboardButton(f"🎵 {p[:45]}", callback_data=f"GEN|{p}")]
        for p in EXAMPLES
    ]
    await update.message.reply_text(
        "💡 *Tap any example to generate:*",
        reply_markup=InlineKeyboardMarkup(buttons),
        parse_mode=ParseMode.MARKDOWN,
    )


async def cmd_settings(update: Update, context: ContextTypes.DEFAULT_TYPE):
    uid = update.effective_user.id
    s   = get_settings(uid)
    buttons = [
        [InlineKeyboardButton(
            f"⏱️ Duration: {s['duration']}s",
            callback_data="SET|duration",
        )],
        [InlineKeyboardButton(
            f"🎯 Guidance: {s['guidance_scale']}",
            callback_data="SET|guidance",
        )],
        [InlineKeyboardButton(
            f"🌡️ Temperature: {s['temperature']}",
            callback_data="SET|temperature",
        )],
        [InlineKeyboardButton(
            "🔄 Reset Defaults",
            callback_data="SET|reset",
        )],
    ]
    await update.message.reply_text(
        "⚙️ *Your Settings:*",
        reply_markup=InlineKeyboardMarkup(buttons),
        parse_mode=ParseMode.MARKDOWN,
    )


# ── Callbacks ─────────────────────────────────────────────────────────────────

async def cb_gen_example(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query  = update.callback_query
    await query.answer()
    prompt = query.data.split("|", 1)[1]
    await query.edit_message_text(
        f"🎵 Generating: _{prompt}_",
        parse_mode=ParseMode.MARKDOWN,
    )
    await do_generate(
        query.message.chat_id,
        prompt,
        update.effective_user.id,
        context,
    )


async def cb_settings_menu(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query  = update.callback_query
    await query.answer()
    uid    = update.effective_user.id
    action = query.data.split("|", 1)[1]

    if action == "reset":
        user_settings[uid] = DEFAULT_SETTINGS.copy()
        await query.edit_message_text("✅ Reset to defaults!")
        return

    choices = {
        "duration":    (
            [4, 6, 8, 10, 15, 20, 30],
            "⏱️ Pick duration (seconds):",
        ),
        "guidance":    (
            [1.0, 2.0, 3.0, 5.0, 7.0, 10.0],
            "🎯 Pick guidance scale:",
        ),
        "temperature": (
            [0.3, 0.5, 0.8, 1.0, 1.3, 1.7, 2.0],
            "🌡️ Pick temperature:",
        ),
    }

    vals, label = choices[action]
    buttons = [
        [InlineKeyboardButton(str(v), callback_data=f"SETV|{action}|{v}")]
        for v in vals
    ]
    await query.edit_message_text(
        label,
        reply_markup=InlineKeyboardMarkup(buttons),
    )


async def cb_set_value(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query       = update.callback_query
    await query.answer()
    uid         = update.effective_user.id
    _, key, val = query.data.split("|")
    s           = get_settings(uid)

    if key == "duration":
        s["duration"] = int(val)
    elif key == "guidance":
        s["guidance_scale"] = float(val)
    elif key == "temperature":
        s["temperature"] = float(val)

    await query.edit_message_text(
        f"✅ *Saved!*\n\n"
        f"⏱️ Duration    : {s['duration']}s\n"
        f"🎯 Guidance    : {s['guidance_scale']}\n"
        f"🌡️ Temperature : {s['temperature']}",
        parse_mode=ParseMode.MARKDOWN,
    )


# ── Text messages ─────────────────────────────────────────────────────────────

async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE):
    prompt = update.message.text.strip()

    if len(prompt) < 3:
        await update.message.reply_text(
            "❌ Too short. Describe the music you want."
        )
        return

    if len(prompt) > 500:
        await update.message.reply_text(
            "❌ Too long (max 500 chars). Please shorten it."
        )
        return

    await do_generate(
        update.effective_chat.id,
        prompt,
        update.effective_user.id,
        context,
        reply_id=update.message.message_id,
    )


async def error_handler(update, context):
    logger.error(f"Error: {context.error}", exc_info=context.error)


# ── Entry point ───────────────────────────────────────────────────────────────

def run_bot():
    if not TELEGRAM_TOKEN:
        logger.warning("TELEGRAM_BOT_TOKEN not set — bot will not start.")
        return

    app = Application.builder().token(TELEGRAM_TOKEN).build()

    app.add_handler(CommandHandler("start",    cmd_start))
    app.add_handler(CommandHandler("help",     cmd_help))
    app.add_handler(CommandHandler("examples", cmd_examples))
    app.add_handler(CommandHandler("settings", cmd_settings))

    app.add_handler(CallbackQueryHandler(cb_gen_example,   pattern=r"^GEN\|"))
    app.add_handler(CallbackQueryHandler(cb_settings_menu, pattern=r"^SET\|"))
    app.add_handler(CallbackQueryHandler(cb_set_value,     pattern=r"^SETV\|"))

    app.add_handler(MessageHandler(
        filters.TEXT & ~filters.COMMAND,
        handle_text,
    ))

    app.add_error_handler(error_handler)

    logger.info("Bot polling started.")
    app.run_polling(
        allowed_updates=Update.ALL_TYPES,
        drop_pending_updates=True,
        close_loop=False,
    )