import gradio as gr
import requests
import base64
import io
import soundfile as sf

# 🔥 Your ngrok API endpoint — update this each time you restart Colab
API_URL = "https://bustled-hertha-unprojective.ngrok-free.dev/generate"

# Required to bypass ngrok browser warning page (fixes Invalid JSON error)
HEADERS = {"ngrok-skip-browser-warning": "true"}


def generate(prompt, duration, steps, cfg):
    try:
        print(f"📤 Sending request → prompt='{prompt}' duration={duration}s steps={steps}")

        res = requests.post(
            API_URL,
            json={
                "prompt":    prompt,
                "duration":  float(duration),
                "steps":     int(steps),
                "cfg_scale": float(cfg),
            },
            headers=HEADERS,
            timeout=300,  # 5 min — generation can be slow
        )

        print(f"STATUS: {res.status_code}")

        if res.status_code != 200:
            print(f"❌ Backend error: {res.text[:500]}")
            return None

        if not res.text:
            print("❌ Empty response from backend")
            return None

        try:
            data = res.json()
        except Exception:
            print(f"❌ Invalid JSON (got HTML?): {res.text[:300]}")
            return None

        if "audio" not in data:
            print(f"❌ No audio key in response: {data}")
            return None

        # Decode base64 WAV → numpy
        audio_bytes = base64.b64decode(data["audio"])
        audio, sr   = sf.read(io.BytesIO(audio_bytes))

        print(f"✅ Got audio: {len(audio)/sr:.2f}s @ {sr}Hz")
        return sr, audio

    except requests.exceptions.Timeout:
        print("❌ Request timed out — try fewer steps or shorter duration")
        return None
    except Exception as e:
        print(f"❌ Request failed: {e}")
        return None


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="AutoMix AI 🎵") as demo:
    gr.Markdown("# 🎵 AutoMix AI Beat Generator")
    gr.Markdown("Generate AI beats using a diffusion model fine-tuned on trap/rap/R&B 🚀")

    with gr.Row():
        with gr.Column():
            prompt_in = gr.Textbox(
                label="🎧 Prompt",
                placeholder="A dark trap beat at 140 BPM in C minor, featuring 808 bass and synth bells.",
                lines=3,
            )
            duration_in = gr.Slider(
                minimum=5, maximum=47, value=30, step=1,
                label="⏱ Duration (seconds)",
            )
            steps_in = gr.Slider(
                minimum=20, maximum=200, value=100, step=10,
                label="⚙️ Diffusion Steps  (more = better quality, slower)",
            )
            cfg_in = gr.Slider(
                minimum=1.0, maximum=15.0, value=7.0, step=0.5,
                label="🎛 CFG Scale",
            )
            btn = gr.Button("🚀 Generate Beat", variant="primary")

        with gr.Column():
            output = gr.Audio(label="🎵 Generated Beat", type="numpy")

    btn.click(
        fn=generate,
        inputs=[prompt_in, duration_in, steps_in, cfg_in],
        outputs=output,
    )

demo.launch()