"""SemBridge — unified speech & singing generation (TTS · SVS · mixed) on ZeroGPU."""

import json
import os
import tempfile
import time
from pathlib import Path

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

import spaces  # MUST come before any CUDA-touching import

from sembridge import SemBridge
from sembridge.infer import _write_pcm16

import gradio as gr

MODEL_ID = "ASLP-lab/SemBridge"
EXAMPLES_DIR = Path(__file__).parent / "examples"
MAX_PATCHES = 500  # 40 s hard cap of the released checkpoint (2 × 25 Hz frames per patch)

print("Loading SemBridge checkpoint (≈4.3 GB)…", flush=True)
_t0 = time.perf_counter()
model = SemBridge.from_pretrained(MODEL_ID, device="cuda", dtype="bfloat16")
print(f"SemBridge loaded in {time.perf_counter() - _t0:.1f}s", flush=True)

# --------------------------------------------------------------- catalog
# (demo id, request json, task) — packaged straight from the official repo.
CATALOG = [
    ("tts_zh_r001", "tts/zh_r001.json", "tts"),
    ("tts_zh_r002", "tts/zh_r002.json", "tts"),
    ("tts_en_r001", "tts/en_r001.json", "tts"),
    ("tts_en_r002", "tts/en_r002.json", "tts"),
    ("svs_zh_r001", "svs/zh_r001.json", "svs"),
    ("svs_zh_r002", "svs/zh_r002.json", "svs"),
    ("svs_en_r001", "svs/en_r001.json", "svs"),
    ("svs_en_r002", "svs/en_r002.json", "svs"),
    ("unified_zh_novel_wukong", "unified/zh_novel_wukong.json", "unified"),
    ("unified_zh_robust_caihong", "unified/zh_robust_caihong.json", "unified"),
]


def _read_json(relative: str) -> dict:
    record = json.loads((EXAMPLES_DIR / relative).read_text(encoding="utf-8"))
    if not isinstance(record, dict):
        raise ValueError("request JSON must contain one object")
    return record


def _request_prompt_paths(relative: str) -> dict:
    """Resolve prompt_audio / prompt_text / generation from a packaged request.

    `prompt_audio` in the packaged JSON is relative to the *request file's*
    directory (mirroring sembridge.requests.load_request), i.e. "../audio/x.wav"
    means examples/audio/x.wav.
    """
    record = _read_json(relative)
    prompt_audio = record.get("prompt_audio")
    if prompt_audio is None:
        raise ValueError(f"{relative} has no prompt_audio")
    request_dir = (EXAMPLES_DIR / relative).parent
    prompt_path = Path(prompt_audio)
    if not prompt_path.is_absolute():
        prompt_path = request_dir / prompt_path
    if not prompt_path.is_file():
        raise FileNotFoundError(prompt_path)
    out = {
        "prompt_audio": str(prompt_path),
        "prompt_text": record.get("prompt_text", ""),
        "generation": record.get("generation") or {},
    }
    return out


VOICE_CHOICES = [  # (label, request path) for the reference-voice dropdown
    ("ZH speech 1 — Seed-TTS", "tts/zh_r001.json"),
    ("ZH speech 2 — Seed-TTS", "tts/zh_r002.json"),
    ("EN speech 1 — Common Voice", "tts/en_r001.json"),
    ("EN speech 2 — Common Voice", "tts/en_r002.json"),
    ("ZH singer 1 — GTSinger", "svs/zh_r001.json"),
    ("ZH singer 2 — GTSinger", "svs/zh_r002.json"),
    ("EN singer 1 — GTSinger", "svs/en_r001.json"),
    ("EN singer 2 — Opencpop", "svs/en_r002.json"),
    ("ZH host (spoken ref) — 悟空", "unified/zh_novel_wukong.json"),
    ("ZH ballad host (sung ref) — 彩虹", "unified/zh_robust_caihong.json"),
]
VOICE_BY_LABEL = dict(VOICE_CHOICES)
DEFAULT_TTS_VOICE = VOICE_CHOICES[0][0]
DEFAULT_SVS_VOICE = "ZH singer 1 — GTSinger"


# ------------------------------------------------------------ helpers
def _clean_number(value, fallback):
    """Gradio Number inputs give None for blank; fall back to the model default."""
    if value is None or (isinstance(value, float) and value == 0) or value == "":
        return fallback
    return value


def _generation_kwargs(request_generation: dict, flow_steps, cfg_value, temperature, seed):
    """Request `generation` block first, explicit UI values on top; None = no override."""
    legacy = {
        k: v for k, v in (request_generation or {}).items()
        if k in {"max_new_patches", "min_new_patches", "flow_steps", "cfg_value",
                 "temperature", "seed"} and v is not None
    }
    ui = {
        "flow_steps": flow_steps, "cfg_value": cfg_value,
        "temperature": temperature, "seed": seed,
    }
    legacy.update({k: v for k, v in ui.items() if v is not None})
    return legacy


def _parse_score(score_text: str) -> list[dict]:
    """Parse 'lyric|pitch|seconds' lines into sembridge note dicts."""
    notes = []
    for raw in (score_text or "").splitlines():
        line = raw.strip()
        if not line:
            continue
        parts = [p.strip() for p in line.split("|")]
        if len(parts) != 3:
            raise gr.Error(f"Each note line is 'lyric|pitch|seconds' — got: {line!r}")
        lyric, pitch, duration = parts
        try:
            pitch_i = int(pitch)
            duration_f = float(duration)
        except ValueError:
            raise gr.Error(f"Pitch must be an integer, duration a float — got: {line!r}")
        if not 0 <= pitch_i <= 128:
            raise gr.Error(f"MIDI pitch {pitch_i} is outside [0, 128].")
        if duration_f <= 0:
            raise gr.Error(f"Duration must be positive — got: {line!r}")
        notes.append({"text": lyric, "pitch": pitch_i, "duration": duration_f})
    if not notes:
        raise gr.Error("Enter at least one note: one 'lyric|pitch|seconds' per line.")
    return notes


def _score_to_ui(notes: list[dict]) -> str:
    """Render packaged SVS notes as editable UI score text."""
    return "\n".join(
        f"{n.get('text', '')}|{n.get('pitch', 0)}|{n.get('duration', 0)}" for n in notes
    )


def _estimate_duration(task: str, *args, **kwargs) -> int:
    """Callable @spaces.GPU duration: reserve what this request needs, not a blanket worst case.

    Receives the same inputs as the decorated function (positional layout of
    the handler signature when ZeroGPU passes args positionally, keywords as
    a fallback — both are accepted). Measured on ZeroGPU (RTX 6000 Ada):
    ~6.5 ms per patch×flow-step plus a small cold-start allowance and codec
    decode time.
    """
    text = kwargs.get("text") or kwargs.get("score_text") or (args[0] if args else "") or ""
    flow_steps = kwargs.get("flow_steps")
    if flow_steps is None and len(args) >= 5:
        flow_steps = args[4]
    steps = int(flow_steps) if flow_steps else 10
    steps = max(steps, 2)
    if task == "svs" and isinstance(text, str) and "|" in text:
        # Score text: sum note durations directly.
        total_seconds = 0.0
        for line in text.splitlines():
            if line.strip():
                try:
                    total_seconds += float(line.split("|")[2])
                except (IndexError, ValueError):
                    total_seconds += 0.3
        seconds = total_seconds
    else:
        # Speech: characters ≈ syllables ≈ 155 ms each (CJK 1 char/syllable).
        seconds = (len(text) if isinstance(text, str) else 0) * 0.155
    patches = min(MAX_PATCHES, max(2, int(seconds * 12.5)))
    total = 10.0 + patches * steps * 0.008 + seconds * 0.3
    return min(240, int(total * 1.3))


def _tts_duration(*args, **kwargs) -> int:
    return _estimate_duration("tts", *args, **kwargs)


def _svs_duration(*args, **kwargs) -> int:
    return _estimate_duration("svs", *args, **kwargs)


def _save_and_report(waveform, patches, started, status_box: gr.Textbox | None = None):
    out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    _write_pcm16(Path(out.name), waveform, model.codec.sample_rate)
    seconds = waveform.size(-1) / model.codec.sample_rate
    elapsed = time.perf_counter() - started
    return out.name, (
        f"{seconds:.1f}s audio · {patches} patches · generated in {elapsed:.1f}s"
    )


# ---------------------------------------------------------------- GPU fns
@spaces.GPU(duration=_tts_duration)
def tts_generate(text: str, voice: str, ref_audio, ref_text: str,
                 flow_steps=None, cfg_value=None, temperature=None, seed=None,
                 progress=gr.Progress()):
    """Text-to-speech in the reference voice.

    Args:
        text: Sentence to speak (English or Chinese).
        voice: Packaged reference voice preset; empty string means custom upload.
        ref_audio: Optional uploaded reference recording (must be given with ref_text).
        ref_text: Transcript of the reference recording.
        flow_steps: Flow-matching steps; blank = request/model default (10).
        cfg_value: Classifier-free guidance; blank = 2.0.
        temperature: Sampling temperature; blank = 1.0.
        seed: RNG seed; blank = 42.
    """
    started = time.perf_counter()
    if not (text or "").strip():
        raise gr.Error("Enter some text to speak.")
    request_gen = {}
    if voice and voice in VOICE_BY_LABEL:
        preset = _request_prompt_paths(VOICE_BY_LABEL[voice])
        prompt_audio, prompt_text = preset["prompt_audio"], preset["prompt_text"]
        request_gen = preset["generation"]
    elif ref_audio:
        if not (ref_text or "").strip():
            raise gr.Error("An uploaded voice needs its transcript.")
        prompt_audio, prompt_text = ref_audio, ref_text.strip()
    else:
        raise gr.Error("Pick a packaged voice or upload a reference recording.")
    kwargs = _generation_kwargs(request_gen, flow_steps, cfg_value, temperature, seed)
    waveform, patches = model.generate(
        task="tts", text=text, prompt_audio=prompt_audio, prompt_text=prompt_text,
        prompt_task="speech", **kwargs,
    )
    return _save_and_report(waveform, patches.size(1), started)


@spaces.GPU(duration=_svs_duration)
def svs_generate(score_text: str, voice: str, ref_audio, ref_text: str,
                 flow_steps=None, cfg_value=None, temperature=None, seed=None,
                 progress=gr.Progress()):
    """Note-conditioned singing from a score string.

    Args:
        score_text: One note per line, "lyric|pitch|duration_seconds". "<SP>" = rest.
        voice: Packaged sung reference preset; empty string means custom upload.
        ref_audio: Optional uploaded sung reference recording.
        ref_text: Score-formatted transcript of the sung reference.
        flow_steps: Flow-matching steps; blank = 20 for singing.
        cfg_value: Classifier-free guidance; blank = 2.0.
        temperature: Sampling temperature; blank = 1.0.
        seed: RNG seed; blank = 42.
    """
    started = time.perf_counter()
    notes = _parse_score(score_text)
    request_gen = {}
    if voice and voice in VOICE_BY_LABEL:
        preset = _request_prompt_paths(VOICE_BY_LABEL[voice])
        prompt_audio, prompt_text = preset["prompt_audio"], preset["prompt_text"]
        request_gen = preset["generation"]
    elif ref_audio:
        if not (ref_text or "").strip():
            raise gr.Error("An uploaded sung reference needs its score-formatted transcript.")
        prompt_audio, prompt_text = ref_audio, ref_text.strip()
    else:
        raise gr.Error("Pick a packaged voice or upload a sung reference.")
    kwargs = _generation_kwargs(request_gen, flow_steps, cfg_value, temperature, seed)
    waveform, patches = model.generate(
        task="svs", notes=notes, prompt_audio=prompt_audio, prompt_text=prompt_text,
        prompt_task="singing", **kwargs,
    )
    return _save_and_report(waveform, patches.size(1), started)


def _unified_duration(case_id: str, *args, **kwargs) -> int:
    """Duration for the mixed speech+song cases: longest packaged sequence is
    15 s of audio at flow 20 ≈ 26 s measured; reserve ~1.4× for cold starts."""
    return 45


@spaces.GPU(duration=_unified_duration)
def unified_generate(case_id: str, flow_steps=None, cfg_value=None,
                     temperature=None, seed=None, progress=gr.Progress()):
    """Run one packaged speech→singing→speech sequence end to end.

    Args:
        case_id: Packaged case id: unified_zh_novel_wukong or unified_zh_robust_caihong.
        flow_steps: Optional flow-matching steps override.
        cfg_value: Optional guidance override.
        temperature: Optional temperature override.
        seed: Optional seed override.
    """
    started = time.perf_counter()
    if case_id not in {"unified_zh_novel_wukong", "unified_zh_robust_caihong"}:
        raise gr.Error("Case must be unified_zh_novel_wukong or unified_zh_robust_caihong.")
    record = _read_json(f"unified/{case_id.split('_', 1)[1]}.json")
    preset = _request_prompt_paths(f"unified/{case_id.split('_', 1)[1]}.json")
    kwargs = _generation_kwargs(
        preset["generation"], flow_steps, cfg_value, temperature, seed
    )
    waveform, patches = model.generate(
        task="unified", segments=record["segments"],
        prompt_audio=preset["prompt_audio"], prompt_text=preset["prompt_text"],
        prompt_task=record.get("prompt_task", "speech"), **kwargs,
    )
    return _save_and_report(waveform, patches.size(1), started)


# --------------------------------------------------------------------- UI
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""
CUSTOM = "⬆️ Custom upload"


def _voice_row(request: str) -> str:
    return request


def on_voice_change(voice: str):
    """Preview the packaged recording; prefill its transcript and generation settings."""
    if voice not in VOICE_BY_LABEL:
        return [gr.update(value=None, visible=True, interactive=True),
                gr.update(value=""), None, None, None, None]
    preset = _request_prompt_paths(VOICE_BY_LABEL[voice])
    gen = preset["generation"]
    return [
        gr.update(value=preset["prompt_audio"], visible=True, interactive=False),
        gr.update(value=preset["prompt_text"]),
        gen.get("flow_steps"),
        gen.get("cfg_value"),
        gen.get("temperature"),
        gen.get("seed"),
    ]


def tts_examples() -> list[list]:
    rows = []
    for label, request in VOICE_CHOICES[:4]:
        record = _read_json(request)
        rows.append([record.get("text", ""), label])
    return rows


def svs_examples() -> list[list]:
    rows = []
    for label, request in VOICE_CHOICES[4:8]:
        record = _read_json(request)
        rows.append([_score_to_ui(record.get("notes", [])), label])
    return rows


with gr.Blocks(title="SemBridge — Speech & Singing") as demo:
    gr.HTML(
        "<div style='text-align:center'>"
        "<h1>🎤 SemBridge</h1>"
        "<p style='margin:4px 0'><b>Semantic Token Anchoring for Continuous-Latent "
        "Autoregressive Speech Generation</b> — speech, singing, one voice.</p>"
        "<p style='margin:4px 0'>Model: <a href='https://huggingface.co/ASLP-lab/SemBridge'>"
        "ASLP-lab/SemBridge</a> · <a href='https://github.com/ASLP-lab/SemBridge'>Code</a> · "
        "<a href='https://arxiv.org/abs/2608.07462'>Paper</a> · "
        "<a href='https://tiamojames.github.io/SemBridge_Demo/'>Demo page</a></p>"
        "</div>"
    )
    gr.Markdown(
        "🔊 The model generates from a **reference voice** — choose one of the ten "
        "packaged speakers/singers (exact references from the authors' demo, press ▶ on "
        "the preview to listen) or upload ~4–10 s of clean audio with its transcript."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Tabs():
            with gr.Tab("🗣️ Text to Speech"):
                tts_text = gr.Textbox(
                    label="Text to speak (English or Chinese)",
                    placeholder="Reading between the lines requires understanding.",
                    lines=2, autofocus=True,
                )
                tts_run = gr.Button("Generate speech", variant="primary")
                tts_out = gr.Audio(label="Generated speech", type="filepath")
                tts_status = gr.Textbox(label="Info", show_label=False, interactive=False)
                with gr.Accordion("Reference voice (required)", open=True):
                    tts_voice = gr.Dropdown(list(VOICE_BY_LABEL) + [CUSTOM],
                                            value=DEFAULT_TTS_VOICE, label="Voice")
                    tts_ref_audio = gr.Audio(
                        label="Reference recording (preview; upload when custom)",
                        type="filepath", sources=["upload"],
                    )
                    tts_ref_text = gr.Textbox(label="Reference transcript", lines=2,
                                              placeholder="Upload + transcript when set to Custom upload")
                with gr.Accordion("Generation settings", open=False):
                    with gr.Row():
                        tts_flow = gr.Number(label="Flow steps (default 10)", precision=0)
                        tts_cfg = gr.Number(label="CFG (default 2.0)")
                        tts_temp = gr.Number(label="Temperature (default 1.0)")
                        tts_seed = gr.Number(label="Seed (default 42)", precision=0)
                gr.Examples(
                    examples=tts_examples(),
                    inputs=[tts_text, tts_voice],
                    outputs=[tts_out, tts_status],
                    fn=tts_generate, cache_examples=True, cache_mode="lazy",
                    label="Official demo sentences (Seed-TTS references)",
                )

            with gr.Tab("🎵 Singing (SVS)"):
                svs_score = gr.Textbox(
                    label="Score — one note per line: lyric | MIDI pitch | duration (s). "
                          "'<SP>' = rest.",
                    placeholder="like|58|0.30\nthis|61|0.36\n<SP>|50|0.25",
                    lines=8, autofocus=True,
                )
                svs_run = gr.Button("Generate singing", variant="primary")
                svs_out = gr.Audio(label="Generated singing", type="filepath")
                svs_status = gr.Textbox(label="Info", show_label=False, interactive=False)
                with gr.Accordion("Sung reference voice (required)", open=True):
                    svs_voice = gr.Dropdown(list(VOICE_BY_LABEL) + [CUSTOM],
                                            value=DEFAULT_SVS_VOICE, label="Voice")
                    svs_ref_audio = gr.Audio(
                        label="Sung reference (preview; upload when custom)",
                        type="filepath", sources=["upload"],
                    )
                    svs_ref_text = gr.Textbox(label="Score-formatted reference transcript",
                                              lines=4)
                with gr.Accordion("Generation settings", open=False):
                    with gr.Row():
                        svs_flow = gr.Number(label="Flow steps (default 20)", precision=0)
                        svs_cfg = gr.Number(label="CFG (default 2.0)")
                        svs_temp = gr.Number(label="Temperature (default 1.0)")
                        svs_seed = gr.Number(label="Seed (default 42)", precision=0)
                gr.Examples(
                    examples=svs_examples(),
                    inputs=[svs_score, svs_voice],
                    outputs=[svs_out, svs_status],
                    fn=svs_generate, cache_examples=True, cache_mode="lazy",
                    label="Official demo scores (GTSinger / Opencpop)",
                )

            with gr.Tab("🎭 Mixed speech + song"):
                gr.Markdown(
                    "The unified model speaks, sings, then speaks again — in one "
                    "sequence, one voice. These are the two official packaged cases."
                )
                case_id = gr.Radio(
                    ["unified_zh_novel_wukong", "unified_zh_robust_caihong"],
                    value="unified_zh_novel_wukong",
                    label="Case (speech → singing → speech)",
                )
                uni_run = gr.Button("Generate mixed sequence", variant="primary")
                uni_out = gr.Audio(label="Generated mixed audio", type="filepath")
                uni_status = gr.Textbox(label="Info", show_label=False, interactive=False)
                with gr.Accordion("Generation overrides (optional)", open=False):
                    with gr.Row():
                        uni_flow = gr.Number(label="Flow steps", precision=0)
                        uni_cfg = gr.Number(label="CFG")
                        uni_temp = gr.Number(label="Temperature")
                        uni_seed = gr.Number(label="Seed", precision=0)

        gr.Markdown(
            "📝 Output is mono 44.1 kHz WAV, capped at 500 patches (~40 s). Singing "
            "durations are quantized to 25 ms frames; the score format mirrors the "
            "authors' JSON (lyric / MIDI pitch / seconds). Packaged references are "
            "Seed-TTS / GTSinger / Opencpop samples from the official demo and keep "
            "their original terms."
        )

    tts_voice.change(on_voice_change, tts_voice,
                     [tts_ref_audio, tts_ref_text, tts_flow, tts_cfg,
                      tts_temp, tts_seed])
    svs_voice.change(on_voice_change, svs_voice,
                     [svs_ref_audio, svs_ref_text, svs_flow, svs_cfg,
                      svs_temp, svs_seed])
    tts_run.click(tts_generate,
                  inputs=[tts_text, tts_voice, tts_ref_audio, tts_ref_text,
                          tts_flow, tts_cfg, tts_temp, tts_seed],
                  outputs=[tts_out, tts_status], api_name="tts")
    svs_run.click(svs_generate,
                  inputs=[svs_score, svs_voice, svs_ref_audio, svs_ref_text,
                          svs_flow, svs_cfg, svs_temp, svs_seed],
                  outputs=[svs_out, svs_status], api_name="svs")
    uni_run.click(unified_generate,
                  inputs=[case_id, uni_flow, uni_cfg, uni_temp, uni_seed],
                  outputs=[uni_out, uni_status], api_name="unified")

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)