from __future__ import annotations

import os
import uuid
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import gradio as gr
import numpy as np
import soundfile as sf
import spaces
import torch
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from magenta_rt import paths
from magenta_rt.torch import MagentaRT2
from magenta_rt.torch.musiccoca import MusicCoCa


MODEL_REPO = "google/magenta-realtime-2"
MODEL_NAME = "mrt2_small"
CHECKPOINT = f"{MODEL_NAME}.safetensors"
AOTI_REPO = "magenta-torch/magenta-rt-aoti-small"
SAMPLE_RATE = 48_000
FRAMES_PER_SECOND = 25
OUTPUT_DIR = Path("/tmp/magenta_rt_outputs")

model_root = Path("/data" if Path("/data").is_dir() else "/tmp/magenta")
magenta_home = model_root / "magenta-rt-v2"
magenta_home.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

hf_hub_download(
    repo_id=MODEL_REPO,
    filename=f"checkpoints/{CHECKPOINT}",
    local_dir=magenta_home,
)
paths.set_magenta_home(magenta_home)

style_model = MusicCoCa(device="cpu")
model = MagentaRT2(
    size=MODEL_NAME,
    device="cuda",
    dtype=torch.bfloat16,
    style_model=style_model,
)

try:
    model.load_compiled(repo_id=AOTI_REPO)
except Exception as exc:
    print(f"AOTI loading failed; using eager inference: {exc}", flush=True)


model_card = ModelCard(
    name="Magenta RealTime 2",
    description=(
        "Generate short instrumental music clips from text prompts using "
        "Google's open-weights Magenta RealTime 2 small model."
    ),
    author="Google DeepMind",
    tags=[
        "music-generation",
        "text-to-music",
        "instrument-synthesis",
        "real-time-music",
    ],
)


@spaces.GPU(duration=45)
def process_fn(
    prompt: str,
    duration: str,
    temperature: float,
    top_k: float,
    seed: float,
) -> str:
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please describe the music you want to generate.")
    if len(prompt) > 300:
        raise gr.Error("The prompt must be 300 characters or fewer.")

    duration_seconds = int(duration)
    expected_samples = duration_seconds * SAMPLE_RATE
    frames = duration_seconds * FRAMES_PER_SECOND + 1

    try:
        if style_model.device != "cuda":
            style_model.to("cuda")
        style_tokens = style_model.embed_tokens(prompt)
        audio, _ = model.generate(
            style=style_tokens,
            temperature=float(temperature),
            top_k=int(top_k),
            frames=frames,
            seed=int(seed),
            flush=True,
        )
    except Exception as exc:
        raise gr.Error(f"Magenta RealTime 2 inference failed: {exc}") from exc

    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim != 2 or audio.shape[1] != 2:
        raise gr.Error("The model returned an unexpected audio shape.")
    if len(audio) < expected_samples:
        raise gr.Error("The model returned less audio than requested.")

    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}.wav"
    sf.write(
        output_path,
        audio[:expected_samples],
        SAMPLE_RATE,
        subtype="PCM_16",
    )
    return str(output_path)


with gr.Blocks(title="Magenta RealTime 2") as demo:
    input_components = [
        gr.Textbox(
            value="warm analog synthesizer with a gentle rhythmic pulse",
            label="Music Prompt",
            info="Describe the instruments, texture, style, or mood.",
            lines=2,
            max_lines=4,
        ),
        gr.Dropdown(
            choices=["2", "4", "8"],
            value="4",
            label="Duration (seconds)",
            info="Length of the generated clip.",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            step=0.1,
            value=1.1,
            label="Temperature",
            info="Higher values produce more variation.",
        ),
        gr.Slider(
            minimum=10,
            maximum=100,
            step=5,
            value=50,
            label="Top-k",
            info="Limits each sampling step to the most likely tokens.",
        ),
        gr.Number(
            value=0,
            minimum=0,
            maximum=2_147_483_647,
            precision=0,
            label="Seed",
            info="Use the same seed and controls to reproduce a result.",
        ),
    ]
    output_components = [
        gr.Audio(
            type="filepath",
            label="Generated Music",
        ).set_info("A 48 kHz stereo WAV file."),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True,
        pwa=True,
    )