import os
import threading
import logging
import gradio as gr

from music_gen import load_model, generate_music

logging.basicConfig(
    format="%(asctime)s | %(levelname)-8s | %(name)s — %(message)s",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


# ── Load model once at startup ────────────────────────────────────────────────
logger.info("Loading model at startup...")
load_model()
logger.info("Model ready!")


# ── Gradio generation wrapper ─────────────────────────────────────────────────
def ui_generate(prompt, duration, guidance, temperature):
    if not prompt.strip():
        raise gr.Error("Please enter a music description.")
    try:
        path = generate_music(
            prompt=prompt,
            duration=int(duration),
            guidance_scale=float(guidance),
            temperature=float(temperature),
        )
        return path, f"✅ Generated {duration}s for: '{prompt}'"
    except Exception as e:
        raise gr.Error(str(e))


# ── Gradio UI ─────────────────────────────────────────────────────────────────
EXAMPLES = [
    ["upbeat electronic dance music with synth leads",  10, 3.0, 1.0],
    ["calm piano melody ambient relaxing",               8, 3.0, 0.8],
    ["rock guitar riff with heavy drums",                8, 4.0, 1.0],
    ["jazz saxophone solo with soft drums",             10, 3.0, 1.0],
    ["cinematic orchestral epic trailer music",         12, 4.0, 0.9],
    ["lo-fi hip hop beats to study to",                 10, 3.0, 1.0],
    ["8-bit chiptune retro video game music",            8, 3.0, 1.1],
]

with gr.Blocks(
    title="🎵 AI Music Generator",
    theme=gr.themes.Soft(),
    css="h1{text-align:center} footer{display:none!important}",
) as demo:

    gr.Markdown("# 🎵 AI Music Generator")
    gr.Markdown(
        "<p style='text-align:center;color:#666'>"
        "Generate music from text • Powered by MusicGen Small"
        "</p>"
    )

    with gr.Row():

        with gr.Column(scale=3):
            prompt_box = gr.Textbox(
                label="🎼 Describe the music",
                placeholder="e.g. calm piano with rain sounds, lo-fi, relaxing...",
                lines=3,
            )
            duration_sl = gr.Slider(
                minimum=1,
                maximum=30,
                value=8,
                step=1,
                label="⏱️ Duration (seconds)",
                info="Longer = slower on free CPU",
            )
            with gr.Accordion("⚙️ Advanced", open=False):
                guidance_sl = gr.Slider(
                    minimum=1.0,
                    maximum=10.0,
                    value=3.0,
                    step=0.5,
                    label="🎯 Guidance Scale",
                )
                temp_sl = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="🌡️ Temperature",
                )
            gen_btn = gr.Button(
                "🎵 Generate",
                variant="primary",
                size="lg",
            )

        with gr.Column(scale=2):
            audio_out  = gr.Audio(
                label="🔊 Output",
                type="filepath",
            )
            status_out = gr.Textbox(
                label="Status",
                interactive=False,
            )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt_box, duration_sl, guidance_sl, temp_sl],
        label="💡 Examples",
    )

    gen_btn.click(
        fn=ui_generate,
        inputs=[prompt_box, duration_sl, guidance_sl, temp_sl],
        outputs=[audio_out, status_out],
    )
    prompt_box.submit(
        fn=ui_generate,
        inputs=[prompt_box, duration_sl, guidance_sl, temp_sl],
        outputs=[audio_out, status_out],
    )


# ── Start Telegram bot in background thread ───────────────────────────────────
def start_bot():
    import asyncio
    if not os.getenv("TELEGRAM_BOT_TOKEN"):
        logger.warning("No TELEGRAM_BOT_TOKEN set — skipping bot.")
        return
    try:
        from bot import run_bot
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        run_bot()
    except Exception as e:
        logger.error(f"Bot thread crashed: {e}", exc_info=True)


threading.Thread(
    target=start_bot,
    name="TelegramBot",
    daemon=True,
).start()
logger.info("Telegram bot thread started.")


# ── Launch Gradio ─────────────────────────────────────────────────────────────
demo.queue(max_size=3)
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,
)