"""MiniMax-Music3-Jam: describe a song in plain English, an LLM writes the lyrics and the
structured caption, MiniMax Music 3 sings it, and the song can be shared to a community feed
backed by an HF bucket mounted at /data."""

import os

# Cache dirs must be set before any HF/torch import (ZeroGPU: ~/.cache is not writable).
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
for _v in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"):
    os.environ.pop(_v, None)

import spaces  # noqa: E402  must precede torch

import base64  # noqa: E402
import io  # noqa: E402
import json  # noqa: E402
import random  # noqa: E402
import time  # noqa: E402
import traceback  # noqa: E402
import uuid  # noqa: E402
from datetime import datetime, timezone  # noqa: E402

import numpy as np  # noqa: E402
import scipy.io.wavfile  # noqa: E402
import torch  # noqa: E402
from fastapi.responses import FileResponse, HTMLResponse  # noqa: E402
import gradio as gr  # noqa: E402
from gradio import Server  # noqa: E402

import composer  # noqa: E402
import engine  # noqa: E402  (loads MiniMax Music 3 + AoTI kernels onto the GPU)

_HERE = os.path.dirname(os.path.abspath(__file__))
BUCKET_ID = os.environ.get("COMMUNITY_BUCKET", "victor/minimax-music3-community")
BUCKET_URL = f"https://huggingface.co/buckets/{BUCKET_ID}/resolve"
DATA_DIR = os.environ.get("COMMUNITY_DIR", "/data")
SONGS_DIR = os.path.join(DATA_DIR, "songs")
MAX_SEED = int(np.iinfo(np.int32).max)
MAX_DURATION = 180.0
STEPS, GUIDANCE = 30, 1.7
FEED_LIMIT = 60

print(f"[startup] model ready: sr={engine.SAMPLE_RATE} frame_rate={engine.FRAME_RATE}", flush=True)
print(f"[startup] community bucket: {BUCKET_ID} mounted at {DATA_DIR} (exists={os.path.isdir(DATA_DIR)})", flush=True)

# ── Cover art (Z-Image-Turbo), optional ──────────────────────────────────────
try:
    from diffusers import FlowMatchEulerDiscreteScheduler, ZImagePipeline

    _zimage = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16)
    _zimage.to("cuda")
    print("[startup] Z-Image-Turbo loaded for cover art", flush=True)
except Exception as e:  # noqa: BLE001
    _zimage = None
    print(f"[startup] Z-Image-Turbo unavailable, songs will ship without cover art: {e}", flush=True)


def _render_cover(word: str) -> bytes | None:
    if _zimage is None or not word:
        return None
    try:
        t0 = time.time()
        _zimage.scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
        image = _zimage(
            prompt=f"{word} studio photography close-up black background",
            height=768, width=768,
            guidance_scale=0.0,
            num_inference_steps=9,
            generator=torch.Generator("cuda").manual_seed(random.randint(1, 1_000_000)),
            max_sequence_length=512,
        ).images[0]
        buf = io.BytesIO()
        image.save(buf, format="PNG", optimize=True)
        print(f"[cover] '{word}' in {time.time() - t0:.1f}s ({len(buf.getvalue()) // 1024}KB)", flush=True)
        return buf.getvalue()
    except Exception as e:  # noqa: BLE001
        print(f"[cover] failed: {e}", flush=True)
        return None


# ── GPU job: song (+ optional cover) in a single ZeroGPU allocation ───────────
def _gpu_seconds(caption, lyrics, duration, seed, cover_word, *args, **kwargs):
    return engine.estimate_gpu_seconds(duration, STEPS) + (12 if cover_word else 0)


@spaces.GPU(duration=_gpu_seconds, size="xlarge")
def _gpu_job(caption, lyrics, duration, seed, cover_word):
    # Errors are returned as data: exceptions raised in the ZeroGPU worker reach the caller
    # stripped down to their class name.
    try:
        wav, sr, audio_seconds, wall = engine.generate_wav(caption, lyrics, duration, seed, STEPS, GUIDANCE)
    except Exception as e:  # noqa: BLE001
        print(f"[gpu ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
        return {"error": f"{type(e).__name__}: {e}"}
    print(f"[gpu] {audio_seconds:.1f}s of audio in {wall:.1f}s (seed={seed})", flush=True)
    cover = _render_cover(cover_word) if cover_word else None
    return {"wav": wav, "sr": sr, "seconds": audio_seconds, "cover": cover}


class GenerationError(RuntimeError):
    pass


def _wav_bytes(wav: np.ndarray, sr: int) -> bytes:
    buf = io.BytesIO()
    scipy.io.wavfile.write(buf, sr, wav)
    return buf.getvalue()


def _friendly_error(err: Exception) -> str:
    msg = (str(err) or "").lower()
    if any(h in msg for h in ("gpu limit", "quota", "no gpu", "could not allocate", "gpu is busy", "too many", "concurrent")):
        return "The shared GPU is at capacity right now. Please wait a minute and retry."
    if "out of memory" in msg or "oom" in msg:
        return "Generation ran out of GPU memory. Try a shorter duration."
    if "zero audio frames" in msg:
        return "The model produced no audio for these inputs; try different lyrics or a longer duration."
    if isinstance(err, GenerationError):
        return "Generation failed on the GPU. Please try again (a shorter duration helps if it repeats)."
    if isinstance(err, (ValueError, RuntimeError)) and str(err):
        return str(err)
    return f"Generation failed ({type(err).__name__}). Please try again."


# ── Community feed (bucket-backed, cached in memory) ─────────────────────────
_feed: list[dict] = []


def _public_meta(meta: dict) -> dict:
    keys = ("id", "title", "tags", "description", "lyrics", "caption", "duration", "audio_url", "thumb_url", "created_at", "seed")
    return {k: meta.get(k) for k in keys}


def _load_feed():
    if not os.path.isdir(SONGS_DIR):
        print("[feed] no songs dir yet, starting empty", flush=True)
        return
    t0 = time.time()
    for song_id in os.listdir(SONGS_DIR):
        meta_path = os.path.join(SONGS_DIR, song_id, "meta.json")
        if not os.path.isfile(meta_path):
            continue
        try:
            with open(meta_path) as f:
                meta = json.load(f)
            meta["audio_url"] = f"{BUCKET_URL}/songs/{song_id}/{song_id}.wav"
            if os.path.isfile(os.path.join(SONGS_DIR, song_id, "thumb.png")):
                meta["thumb_url"] = f"{BUCKET_URL}/songs/{song_id}/thumb.png"
            else:
                meta["thumb_url"] = None
            _feed.append(_public_meta(meta))
        except Exception as e:  # noqa: BLE001
            print(f"[feed] skipping {song_id}: {e}", flush=True)
    _feed.sort(key=lambda s: s.get("created_at") or "", reverse=True)
    print(f"[feed] loaded {len(_feed)} songs in {time.time() - t0:.1f}s", flush=True)


_load_feed()


def _share(wav_bytes: bytes, cover: bytes | None, meta: dict) -> dict:
    """Persist a song to the bucket mount and prepend it to the in-memory feed."""
    song_id = uuid.uuid4().hex[:12]
    song_dir = os.path.join(SONGS_DIR, song_id)
    os.makedirs(song_dir, exist_ok=True)
    with open(os.path.join(song_dir, f"{song_id}.wav"), "wb") as f:
        f.write(wav_bytes)
    thumb_url = None
    if cover:
        with open(os.path.join(song_dir, "thumb.png"), "wb") as f:
            f.write(cover)
        thumb_url = f"{BUCKET_URL}/songs/{song_id}/thumb.png"
    meta = {
        **meta,
        "id": song_id,
        "audio_url": f"{BUCKET_URL}/songs/{song_id}/{song_id}.wav",
        "thumb_url": thumb_url,
        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }
    with open(os.path.join(song_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)
    public = _public_meta(meta)
    _feed.insert(0, public)
    print(f"[share] {song_id} '{meta.get('title')}' -> {meta['audio_url']}", flush=True)
    return public


# ── gr.Server app ────────────────────────────────────────────────────────────
app = Server(title="MiniMax-Music3-Jam")


def _clamp_duration(d) -> float:
    try:
        d = float(d)
    except (TypeError, ValueError):
        d = 60.0
    return max(10.0, min(MAX_DURATION, d))


def _resolve_seed(seed) -> int:
    try:
        seed = int(seed)
    except (TypeError, ValueError):
        seed = -1
    return random.randint(0, MAX_SEED) if seed < 0 else min(seed, MAX_SEED)


def _run_song(*, description, title, tags, caption, lyrics, duration, seed, community, cover_word):
    """Shared tail of /create and /generate: GPU job, encode, optional share. Returns result dict."""
    engine.validate(caption, lyrics)
    out = _gpu_job(caption, lyrics, duration, seed, cover_word)
    if out.get("error"):
        msg = out["error"]
        if "out of memory" in msg.lower():
            raise GenerationError("out of memory")
        raise GenerationError(msg)
    wav, sr, audio_seconds, cover = out["wav"], out["sr"], out["seconds"], out["cover"]
    wav_bytes = _wav_bytes(wav, sr)
    result = {
        "title": title,
        "tags": tags,
        "lyrics": lyrics,
        "caption": caption,
        "seed": seed,
        "duration": round(audio_seconds, 1),
        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "audio": "data:audio/wav;base64," + base64.b64encode(wav_bytes).decode(),
    }
    if cover:
        result["thumbnail"] = "data:image/png;base64," + base64.b64encode(cover).decode()
    if community:
        try:
            shared = _share(wav_bytes, cover, {
                "title": title, "tags": tags, "description": description, "lyrics": lyrics,
                "caption": caption, "duration": round(audio_seconds, 1), "seed": seed,
            })
            result["community_url"] = shared["audio_url"]
            result["id"] = shared["id"]
            result["created_at"] = shared["created_at"]
        except Exception as e:  # noqa: BLE001
            print(f"[share] failed: {e}", flush=True)
            result["share_error"] = "Song generated, but sharing to the community feed failed."
    return result


@app.api(name="create", concurrency_limit=None, time_limit=900)
def create(
    description: str,
    audio_duration: float = 60.0,
    seed: int = -1,
    community: bool = False,
    instrumental: bool = False,
) -> str:
    """One-box: describe a song -> LLM writes title/lyrics/structured caption -> MiniMax Music 3 generates it.
    Streams JSON status messages ({"status": ...}) and finally the result
    ({audio, title, tags, lyrics, caption, seed, duration, created_at, thumbnail?, community_url?})."""
    try:
        description = (description or "").strip()
        if not description:
            raise ValueError("Describe the song you want first.")
        duration = _clamp_duration(audio_duration)
        seed = _resolve_seed(seed)

        yield json.dumps({"status": "composing", "message": "Writing lyrics & structured caption…"})
        t0 = time.time()
        song = composer.compose(description, duration, instrumental=bool(instrumental))
        print(f"[create] composed '{song['title']}' in {time.time() - t0:.1f}s | tags={song['tags']}", flush=True)
        cover_word = composer.visual_word(song["title"], song["tags"], song["lyrics"], description) if _zimage else ""

        yield json.dumps({
            "status": "generating",
            "message": f"MiniMax Music 3 is recording “{song['title']}”…",
            "title": song["title"], "tags": song["tags"], "lyrics": song["lyrics"],
            "eta": engine.estimate_gpu_seconds(duration, STEPS),
        })
        result = _run_song(
            description=description, title=song["title"], tags=song["tags"], caption=song["caption"],
            lyrics=song["lyrics"], duration=duration, seed=seed, community=bool(community), cover_word=cover_word,
        )
        yield json.dumps(result)
    except Exception as e:  # noqa: BLE001
        print(f"[create ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
        raise gr.Error(_friendly_error(e), print_exception=False) from e


@app.api(name="generate", concurrency_limit=None, time_limit=900)
def generate(
    caption: str,
    lyrics: str,
    audio_duration: float = 60.0,
    seed: int = -1,
    title: str = "",
    tags: str = "",
    community: bool = False,
    cover_word: str = "",
) -> str:
    """Advanced: generate from an explicit structured caption + tagged lyrics (no LLM).
    cover_word: optional visual noun for the Z-Image cover art. Returns the same JSON result shape as /create."""
    try:
        duration = _clamp_duration(audio_duration)
        seed = _resolve_seed(seed)
        lyrics = composer.normalize_lyrics(lyrics)
        title = (title or "").strip()[:80] or "Untitled"
        tags = (tags or "").strip()[:120]
        result = _run_song(
            description="", title=title, tags=tags, caption=(caption or "").strip(), lyrics=lyrics,
            duration=duration, seed=seed, community=bool(community),
            cover_word=(cover_word or "").strip()[:40] if _zimage else "",
        )
        return json.dumps(result)
    except Exception as e:  # noqa: BLE001
        print(f"[generate ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
        raise gr.Error(_friendly_error(e), print_exception=False) from e


@app.api(name="community", concurrency_limit=8)
def community() -> str:
    """Newest community songs (JSON list), served from memory. Each entry carries created_at (UTC ISO)."""
    return json.dumps(_feed[:FEED_LIMIT], ensure_ascii=False)


@app.get("/", response_class=HTMLResponse)
async def homepage():
    with open(os.path.join(_HERE, "index.html"), encoding="utf-8") as f:
        return f.read()


@app.get("/logo.png", include_in_schema=False)
async def logo():
    return FileResponse(os.path.join(_HERE, "logo_dark.png"), media_type="image/png")


demo = app

if __name__ == "__main__":
    # mcp_server=True exposes /create, /generate, /community as MCP tools (schema comes
    # straight from each function's type hints + docstring, already written for the REST API).
    demo.launch(show_error=True, mcp_server=True)