"""MuseForge HeartMuLa ZeroGPU Test — standalone real model check.

Hugging Face ZeroGPU -> HeartMuLa 3B -> HeartCodec -> real Hindi MP3.

Flow:
  1. Space starts on CPU: checks disk, downloads the 3 public model repos
     only when their marker files are missing, then builds the heartlib
     pipeline with lazy_load=True (tokenizer + config only). NO GPU
     allocation and NO model weights in VRAM at startup.
  2. Clicking "Generate MP3" calls the ZeroGPU-decorated function -> ZeroGPU
     allocates a GPU -> heartlib lazily loads HeartMuLa weights onto it,
     generates up to the requested audio length, decodes with HeartCodec
     and saves an MP3. heartlib's own _unload() (del + gc + empty_cache)
     frees VRAM after the forward pass and after codec decoding, before
     the function returns.

Model repos (all public, Apache-2.0):
  - HeartMuLa/HeartMuLaGen          -> ckpt/gen_config.json + ckpt/tokenizer.json
  - HeartMuLa/HeartMuLa-oss-3B      -> ckpt/HeartMuLa-oss-3B/
  - HeartMuLa/HeartCodec-oss-20260123 -> ckpt/HeartCodec-oss/
    (the bare "HeartCodec-oss" repo is gated/not downloadable — 401 on the
     Hub API; the official heartlib README recommends the -20260123 codec,
     which is public. Swap HEARTMULA_REPO to the happy-new-year variant by
     changing one constant if desired.)
"""

import gc
import os
import shutil
import threading
import time
import traceback
from pathlib import Path

import spaces
import gradio as gr
import torch

try:
    from heartlib.pipelines.music_generation import HeartMuLaGenPipeline
except ModuleNotFoundError:
    # heartlib pins transformers==4.57.0 (needs huggingface-hub<1.0), which
    # conflicts with the Space runtime's gradio 6.24.0 (needs huggingface-hub
    # >=1.16). Install heartlib with --no-deps so pip does not enforce its
    # incompatible pins; all real runtime deps are pinned in requirements.txt.
    import subprocess
    import sys

    print("[startup] installing heartlib (--no-deps, pinned commit)...", flush=True)
    subprocess.check_call(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--no-deps",
            "heartlib @ git+https://github.com/HeartMuLa/heartlib.git@3783bdb8441f2c298b1e64c8651173aac200361c",
        ]
    )
    from heartlib.pipelines.music_generation import HeartMuLaGenPipeline

APP_DIR = Path("/home/user/app")
CKPT_DIR = APP_DIR / "ckpt"
VERIFY_DIR = APP_DIR / "verify"

HEARTMULA_REPO = "HeartMuLa/HeartMuLa-oss-3B"
CODEC_REPO = "HeartMuLa/HeartCodec-oss-20260123"
GEN_CONFIG_REPO = "HeartMuLa/HeartMuLaGen"

MODEL_VERSION = "3B"

DEFAULT_LYRICS = """[verse]
रात के सन्नाटे में
एक नई सुबह की आहट है
गिरकर भी मैं उठता हूँ
दिल में अब भी चाहत है

[chorus]
मैं फिर उठूँगा
मैं फिर चलूँगा
हज़ार बार गिरकर
फिर संभलूँगा

[verse]
रास्ते चाहे मुश्किल हों
आँधियाँ चाहे तेज़ चलें
मेरे कदम नहीं रुकेंगे
जब तक सपने साथ चलें

[chorus]
मैं फिर उठूँगा
मैं फिर चलूँगा
हज़ार बार गिरकर
फिर संभलूँगा"""

DEFAULT_TAGS = "hindi, cinematic, emotional, powerful male vocal, warm vocal, piano, strings, atmospheric, melodic"

GEN_LOCK = threading.Lock()
PIPE = None


# ---------------------------------------------------------------- startup (CPU)
def _download_if_missing(repo_id: str, local_dir: Path, marker: Path) -> None:
    if marker.exists():
        print(f"[startup] {repo_id}: already present ({marker.name})", flush=True)
        return
    from huggingface_hub import snapshot_download

    print(f"[startup] downloading {repo_id} -> {local_dir} ...", flush=True)
    snapshot_download(repo_id=repo_id, local_dir=str(local_dir))
    if not marker.exists():
        raise RuntimeError(f"{repo_id} download finished but {marker} is missing")


def _ensure_models() -> None:
    total, used, free = shutil.disk_usage(APP_DIR)
    print(f"[startup] disk: {free / 1024**3:.1f} GB free (models need ~22.4 GB)", flush=True)
    if free < 24 * 1024**3:
        raise RuntimeError(
            f"Insufficient disk: {free / 1024**3:.1f} GB free, need >= 24 GB"
        )

    _download_if_missing(GEN_CONFIG_REPO, CKPT_DIR, CKPT_DIR / "gen_config.json")
    _download_if_missing(
        HEARTMULA_REPO,
        CKPT_DIR / "HeartMuLa-oss-3B",
        CKPT_DIR / "HeartMuLa-oss-3B" / "model.safetensors.index.json",
    )
    _download_if_missing(
        CODEC_REPO,
        CKPT_DIR / "HeartCodec-oss",
        CKPT_DIR / "HeartCodec-oss" / "model.safetensors.index.json",
    )


print(f"[startup] PyTorch {torch.__version__} (CUDA emulation outside the ZeroGPU decorator)", flush=True)
_ensure_models()

print("[startup] creating heartlib pipeline with lazy_load=True (no GPU used yet)...", flush=True)
PIPE = HeartMuLaGenPipeline.from_pretrained(
    str(CKPT_DIR),
    device=torch.device("cuda"),
    dtype={"mula": torch.bfloat16, "codec": torch.float32},
    version=MODEL_VERSION,
    lazy_load=True,
)
print("[startup] pipeline ready. HeartMuLa/HeartCodec weights load on first generation.", flush=True)


# ---------------------------------------------------------------- GPU function
def _allocation_duration(lyrics: str, tags: str, duration_sec: float) -> int:
    """Declared ZeroGPU runtime ceiling (quota is billed on actual runtime).

    The spaces SDK multiplies this value by its duration_factor (1.5 on
    zero-a10g -> RTX PRO 6000 Blackwell) before requesting ZeroGPU, so cap
    at 196 so the SDK-equivalent request is 196 * 1.5 = 294s, below the
    300s free quota (300s requests are rejected).
    """
    return min(196, 120 + int(duration_sec) * 2)


@spaces.GPU(duration=_allocation_duration)
def generate_music(lyrics: str, tags: str, duration_sec: float):
    """Runs inside a real ZeroGPU allocation. Returns (mp3_path, info_text)."""
    started = time.time()
    current_stage = "unknown"

    def stage(name: str) -> None:
        print(f"[generation] {name} (+{time.time() - started:.1f}s)", flush=True)

    try:
        props = torch.cuda.get_device_properties(0)
        info = [
            f"GPU: {props.name}",
            f"VRAM total: {props.total_memory / 1024**3:.2f} GB",
            f"CUDA: {torch.version.cuda}",
            f"PyTorch: {torch.__version__}",
        ]

        VERIFY_DIR.mkdir(parents=True, exist_ok=True)
        save_path = str(VERIFY_DIR / "hindi_heartmula_test.mp3")
        if os.path.exists(save_path):
            os.remove(save_path)
        max_ms = int(max(10, min(int(duration_sec), 240)) * 1000)

        current_stage = "loading model"
        stage("Loading model (HeartMuLa weights -> VRAM)")
        t0 = time.time()
        _ = PIPE.mula  # heartlib lazy weight load
        info.append(f"model load: {time.time() - t0:.1f}s")

        current_stage = "generating"
        stage("Generating (HeartMuLa forward)")
        t0 = time.time()
        PIPE(
            {"lyrics": lyrics, "tags": tags},
            max_audio_length_ms=max_ms,
            topk=50,
            temperature=1.0,
            cfg_scale=1.5,
            save_path=save_path,
        )
        info.append(f"generation + codec decode: {time.time() - t0:.1f}s")

        current_stage = "processing audio"
        stage("Processing audio (MP3 saved)")
        if not os.path.isfile(save_path) or os.path.getsize(save_path) < 10_000:
            raise RuntimeError("MP3 output missing or too small (silence/corrupt audio?)")

        info.append(f"peak VRAM used: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
        info.append(f"output: {save_path}")
        info.append(f"total time: {time.time() - started:.1f}s")
        info.append("STATUS: SUCCESS")
        stage("Completed")
        return save_path, "\n".join(info)
    except Exception as exc:
        traceback.print_exc()
        try:
            mem = (
                f"\nCUDA allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB, "
                f"reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB"
            )
        except Exception:
            mem = ""
        return "", (
            "STATUS: FAILED\n"
            f"stage: {current_stage}\n"
            f"error type: {type(exc).__name__}\n"
            f"error: {exc}{mem}"
        )
    finally:
        gc.collect()
        torch.cuda.empty_cache()


def generate_clicked(lyrics: str, tags: str, duration_sec: float):
    with GEN_LOCK:  # one generation at a time, ever
        return generate_music(lyrics or DEFAULT_LYRICS, tags or DEFAULT_TAGS, int(duration_sec or 40))


# ---------------------------------------------------------------- Gradio UI
with gr.Blocks(title="MuseForge HeartMuLa ZeroGPU Test") as demo:
    gr.Markdown("# MuseForge HeartMuLa ZeroGPU Test")
    lyrics = gr.Textbox(label="Lyrics", value=DEFAULT_LYRICS, lines=14)
    tags = gr.Textbox(label="Music tags", value=DEFAULT_TAGS, lines=2)
    duration = gr.Slider(label="Duration (seconds)", minimum=10, maximum=120, value=40, step=5)
    generate_btn = gr.Button("Generate MP3", variant="primary")
    audio = gr.Audio(label="Generated MP3", type="filepath")
    info = gr.Textbox(label="GPU status / result", lines=14)
    generate_btn.click(generate_clicked, inputs=[lyrics, tags, duration], outputs=[audio, info])

demo.launch()