"""abv1 SA3 engine — shared Stable Audio 3 backend for the music-AI demo apps.

Two surfaces on one Space:

* **UI** — Text→audio and Audio→audio tabs for humans (visitors burn their own
  ZeroGPU quota, so these are never token-gated).
* **Headless API** — ``/generate`` and ``/generate_a2a``, registered off hidden
  components so gradio_client / plain HTTP can drive batch jobs (the trend-radio
  cron) and the demo pages' server side.

Model wiring follows the official ``stabilityai/stable-audio-3`` Space exactly:
``stable_audio_tools.get_pretrained_model`` → ``.to("cuda").to(float16)`` at
module level, ``generate_diffusion_cond_inpaint`` under a bare ZeroGPU shim.

Set ``SA3_STUB=1`` to run the whole app (UI + both endpoints) with no torch and
no weights — gens return a synthesized placeholder WAV. That is the only way to
exercise this file on a laptop.
"""

from __future__ import annotations

import os
import random
import tempfile
import time
from dataclasses import dataclass
from typing import Optional

STUB = os.environ.get("SA3_STUB", "").strip() in ("1", "true", "yes")
API_TOKEN = os.environ.get("SA3_API_TOKEN", "").strip()

# ZeroGPU shim. Absent locally, so fall back to a decorator that ignores its
# kwargs and returns the function untouched.
try:
    import spaces  # type: ignore
except Exception:  # pragma: no cover — local / stub runs
    class _NoSpaces:
        @staticmethod
        def GPU(*d_args, **d_kwargs):
            if d_args and callable(d_args[0]):
                return d_args[0]
            def _wrap(fn):
                return fn
            return _wrap
    spaces = _NoSpaces()  # type: ignore

import gradio as gr
import numpy as np
import soundfile as sf


# ---------------------------------------------------------------------------
# Variants
# ---------------------------------------------------------------------------

@dataclass
class Variant:
    key: str
    repo: str
    label: str
    # Hard cap we expose in the UI/API. The real ceiling is the model's
    # sample_size // sample_rate; we take the min of the two at load time.
    max_seconds: int
    default_seconds: int
    placeholder: str


VARIANTS: list[Variant] = [
    Variant(
        key="small-music",
        repo="stabilityai/stable-audio-3-small-music",
        label="Small Music — 0.6B, fast (seconds per gen)",
        max_seconds=120,
        default_seconds=30,
        placeholder="Cinematic neo-soul groove with electric piano, brushed drums, "
                    "walking upright bass, smoky vibe 92 BPM",
    ),
    Variant(
        key="medium",
        repo="stabilityai/stable-audio-3-medium",
        label="Medium — general audio, higher quality",
        max_seconds=180,
        default_seconds=45,
        placeholder="A dream-like synthpop instrumental for a surrealist dream "
                    "sequence, warm analog pads, 120 BPM",
    ),
]

BY_KEY = {v.key: v for v in VARIANTS}
VARIANT_CHOICES = [(v.label, v.key) for v in VARIANTS]
DEFAULT_VARIANT = "small-music"

STUB_SAMPLE_RATE = 44100


# ---------------------------------------------------------------------------
# Model preload (skipped entirely in stub mode)
# ---------------------------------------------------------------------------

@dataclass
class Loaded:
    variant: Variant
    model: object
    sample_rate: int
    sample_size: int
    max_seconds: int


LOADED: dict[str, Loaded] = {}


def _ensure_stable_audio_tools() -> None:
    import subprocess
    import sys
    try:
        import stable_audio_tools  # noqa: F401
        return
    except ImportError:
        pass
    # stable-audio-tools strict-pins torch==2.7.1, which lacks sm_120 kernels.
    # Install --no-deps; its transitive deps are in requirements.txt.
    print("[startup] installing stable-audio-tools (--no-deps) …", flush=True)
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "--quiet", "--no-deps",
         "stable-audio-tools"],
    )


if not STUB:
    _ensure_stable_audio_tools()

    import torch
    import torchaudio
    from einops import rearrange
    from stable_audio_tools import get_pretrained_model
    from stable_audio_tools.inference.generation import generate_diffusion_cond_inpaint

    for _v in VARIANTS:
        print(f"[startup] loading {_v.repo} …", flush=True)
        _t0 = time.time()
        _model, _config = get_pretrained_model(_v.repo)
        _sr = int(_config["sample_rate"])
        _ss = int(_config["sample_size"])
        _model = _model.to("cuda").to(torch.float16)
        LOADED[_v.key] = Loaded(
            variant=_v,
            model=_model,
            sample_rate=_sr,
            sample_size=_ss,
            max_seconds=min(_v.max_seconds, _ss // _sr),
        )
        print(f"[startup] {_v.key} ready in {time.time() - _t0:.1f}s · "
              f"sr={_sr} · sample_size={_ss} (~{_ss // _sr}s model max)", flush=True)
else:
    print("[startup] SA3_STUB=1 — no torch, no weights, placeholder audio only.",
          flush=True)


def variant_max_seconds(key: str) -> int:
    if key in LOADED:
        return LOADED[key].max_seconds
    return BY_KEY[key].max_seconds if key in BY_KEY else 120


# ---------------------------------------------------------------------------
# Input hygiene
# ---------------------------------------------------------------------------

def _clamp_inputs(variant_key: str, seconds, steps, cfg_scale, seed):
    key = (variant_key or DEFAULT_VARIANT).strip()
    if key not in BY_KEY:
        raise gr.Error(f"Unknown model_variant {variant_key!r}. "
                       f"Use one of: {', '.join(BY_KEY)}")
    seconds = max(1, min(int(float(seconds or 30)), variant_max_seconds(key)))
    steps = max(1, min(int(float(steps or 8)), 100))
    cfg_scale = max(0.0, min(float(cfg_scale if cfg_scale is not None else 1.0), 15.0))
    seed = int(float(seed if seed is not None else -1))
    if seed < 0:
        seed = random.randint(0, 2**31 - 1)
    return key, seconds, steps, cfg_scale, seed


def _check_token(token: Optional[str]) -> None:
    """API-only gate. Unset secret == open (dev). UI paths never call this."""
    if not API_TOKEN:
        return
    if (token or "").strip() != API_TOKEN:
        raise gr.Error("Bad or missing API token.")


def _gpu_seconds(seconds: int) -> int:
    """ZeroGPU budget request. small-music does 120s of audio in <2s of GPU and
    medium a few seconds; 20s base + a quarter of the requested length is ample
    headroom without hogging a slot."""
    return int(min(120, 20 + int(seconds) // 4))


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------

def _stub_wav(seconds: int, seed: int) -> str:
    """Band-limited sweep + noise bed so the shape of the payload (stereo wav,
    right duration, real sample rate) matches a genuine gen."""
    rng = np.random.default_rng(seed)
    sr = STUB_SAMPLE_RATE
    n = int(seconds * sr)
    t = np.arange(n, dtype=np.float32) / sr
    f0, f1 = 110.0, 1760.0
    phase = 2 * np.pi * f0 * seconds / np.log(f1 / f0) * (
        np.power(f1 / f0, t / max(seconds, 1e-6)) - 1.0
    )
    sweep = 0.35 * np.sin(phase).astype(np.float32)
    noise = 0.05 * rng.standard_normal(n).astype(np.float32)
    # cheap one-pole lowpass on the noise so it isn't harsh white
    for _ in range(2):
        noise = np.convolve(noise, np.ones(16, dtype=np.float32) / 16, mode="same")
    env = np.minimum(1.0, np.minimum(t * 4.0, (seconds - t) * 4.0)).astype(np.float32)
    mono = (sweep + noise) * env
    stereo = np.stack([mono, np.roll(mono, 64)], axis=1)
    out_path = os.path.join(tempfile.mkdtemp(), "sa3_stub.wav")
    sf.write(out_path, stereo, sr, subtype="PCM_16")
    return out_path


def _load_init_audio(path: str, target_sr: int, dtype):
    """Filepath → (sample_rate, tensor[C,N]) at the model's rate and dtype.
    Decode via soundfile, not torchaudio.load — torchaudio 2.x delegates
    load() to torchcodec, which isn't in the ZeroGPU image (same trap as the
    audiofix Space). torchaudio.functional.resample is pure torch, so it's
    still fine to use.
    Pre-resampling in fp32 keeps prepare_audio's fp32 resample kernel a no-op,
    which otherwise trips on the fp16 model."""
    data, sr = sf.read(path, dtype="float32", always_2d=True)  # [N, C]
    wav = torch.from_numpy(data.T.copy())                      # [C, N]
    if sr != target_sr:
        wav = torchaudio.functional.resample(wav, sr, target_sr)
    return target_sr, wav.to(dtype)


def _generate(variant_key: str,
              prompt: str,
              negative_prompt: str,
              seconds: int,
              steps: int,
              cfg_scale: float,
              seed: int,
              init_audio_path: Optional[str] = None,
              init_noise_level: Optional[float] = None) -> tuple[str, dict]:
    """Returns (wav_path, meta). Inputs must already be clamped."""
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    t0 = time.time()
    if STUB:
        out_path = _stub_wav(seconds, seed)
        sample_rate = STUB_SAMPLE_RATE
    else:
        lv = LOADED[variant_key]
        conditioning = [{"prompt": prompt, "seconds_total": int(seconds)}]
        negative_conditioning = None
        neg = (negative_prompt or "").strip()
        if neg:
            negative_conditioning = [{"prompt": neg, "seconds_total": int(seconds)}]

        gen_kwargs: dict = dict(
            steps=steps,
            cfg_scale=cfg_scale,
            conditioning=conditioning,
            negative_conditioning=negative_conditioning,
            sample_size=lv.sample_size,
            sampler_type="pingpong",
            seed=seed,
            device="cuda",
        )
        if init_audio_path:
            model_dtype = next(lv.model.parameters()).dtype
            gen_kwargs["init_audio"] = _load_init_audio(
                init_audio_path, lv.sample_rate, model_dtype)
            gen_kwargs["init_noise_level"] = float(
                init_noise_level if init_noise_level is not None else 0.4)

        output = generate_diffusion_cond_inpaint(lv.model, **gen_kwargs)
        output = rearrange(output, "b d n -> d (b n)")
        output = (output.to(torch.float32)
                  .div(torch.max(torch.abs(output)).clamp(min=1e-9))
                  .clamp(-1, 1).mul(32767).to(torch.int16).cpu())
        output = output[:, : seconds * lv.sample_rate]
        sample_rate = lv.sample_rate
        out_path = os.path.join(tempfile.mkdtemp(), "sa3.wav")
        sf.write(out_path, output.numpy().T, sample_rate, subtype="PCM_16")

    meta = {
        "ok": True,
        "seed": seed,
        "model_variant": variant_key,
        "seconds": seconds,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "sample_rate": sample_rate,
        "gen_wall_s": round(time.time() - t0, 2),
        "stub": STUB,
    }
    if init_audio_path:
        meta["init_noise_level"] = float(
            init_noise_level if init_noise_level is not None else 0.4)
    print(f"[gen] {meta}", flush=True)
    return out_path, meta


# --- UI handlers ------------------------------------------------------------

def _ui_duration_t2a(variant_key, prompt, negative_prompt, seconds, steps, cfg_scale, seed):
    return _gpu_seconds(seconds)


@spaces.GPU(duration=_ui_duration_t2a)
def ui_text_to_audio(variant_key, prompt, negative_prompt, seconds, steps, cfg_scale, seed):
    key, seconds, steps, cfg_scale, seed = _clamp_inputs(
        variant_key, seconds, steps, cfg_scale, seed)
    path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, seed)
    return path, f"seed **{meta['seed']}** · {meta['gen_wall_s']}s · {key}"


def _ui_duration_a2a(variant_key, init_audio, prompt, negative_prompt,
                     init_noise_level, seconds, steps, cfg_scale, seed):
    return _gpu_seconds(seconds)


@spaces.GPU(duration=_ui_duration_a2a)
def ui_audio_to_audio(variant_key, init_audio, prompt, negative_prompt,
                      init_noise_level, seconds, steps, cfg_scale, seed):
    if not init_audio:
        raise gr.Error("Upload an audio file to transform.")
    key, seconds, steps, cfg_scale, seed = _clamp_inputs(
        variant_key, seconds, steps, cfg_scale, seed)
    noise = max(0.0, min(float(init_noise_level if init_noise_level is not None else 0.4), 1.0))
    path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale,
                           seed, init_audio_path=init_audio, init_noise_level=noise)
    return path, f"seed **{meta['seed']}** · {meta['gen_wall_s']}s · {key} · noise {noise}"


# --- API handlers -----------------------------------------------------------

def _api_duration_gen(token, variant_key, prompt, negative_prompt,
                      seconds, steps, cfg_scale, seed):
    return _gpu_seconds(seconds)


@spaces.GPU(duration=_api_duration_gen)
def api_generate(token, variant_key, prompt, negative_prompt,
                 seconds, steps, cfg_scale, seed):
    _check_token(token)
    key, seconds, steps, cfg_scale, seed = _clamp_inputs(
        variant_key, seconds, steps, cfg_scale, seed)
    path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, seed)
    return path, meta


def _api_duration_a2a(token, variant_key, init_audio, prompt, negative_prompt,
                      init_noise_level, seconds, steps, cfg_scale, seed):
    return _gpu_seconds(seconds)


@spaces.GPU(duration=_api_duration_a2a)
def api_generate_a2a(token, variant_key, init_audio, prompt, negative_prompt,
                     init_noise_level, seconds, steps, cfg_scale, seed):
    _check_token(token)
    if not init_audio:
        raise gr.Error("init_audio is required for generate_a2a.")
    key, seconds, steps, cfg_scale, seed = _clamp_inputs(
        variant_key, seconds, steps, cfg_scale, seed)
    noise = max(0.0, min(float(init_noise_level if init_noise_level is not None else 0.4), 1.0))
    path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale,
                           seed, init_audio_path=init_audio, init_noise_level=noise)
    return path, meta


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

HEADER = """
# abv1 · SA3 engine

Text→audio and audio→audio with **Stable Audio 3**. Shared backend for the abv1
music-AI demos; also exposes headless `/generate` and `/generate_a2a` endpoints.
"""

FOOTER = """
---
**Powered by Stability AI.** Models are used under the
[Stability AI Community License](https://stability.ai/community-license-agreement);
see `NOTICE.md` in this repo. All audio here is **AI-generated — not real
releases** by any artist.
"""

NOISE_INFO = ("Low (0.1–0.3) = stay close to your upload · "
              "mid (0.4–0.6) = recognisable but reimagined · "
              "high (0.8–1.0) = barely related / go wild")


def _on_variant_change(variant_key):
    mx = variant_max_seconds(variant_key)
    v = BY_KEY.get(variant_key, BY_KEY[DEFAULT_VARIANT])
    return (
        gr.update(maximum=mx, value=min(v.default_seconds, mx),
                  label=f"Seconds · max {mx}"),
        gr.update(placeholder=v.placeholder),
    )


_start = BY_KEY[DEFAULT_VARIANT]
_start_max = variant_max_seconds(DEFAULT_VARIANT)

with gr.Blocks(theme=gr.themes.Citrus(), title="abv1 sa3 engine") as demo:
    gr.Markdown(HEADER)

    with gr.Tabs():
        # ---------------- Text → audio ----------------
        with gr.Tab("Text → audio"):
            t_variant = gr.Radio(VARIANT_CHOICES, value=DEFAULT_VARIANT, label="Model")
            with gr.Row():
                with gr.Column(scale=2):
                    t_prompt = gr.Textbox(label="Prompt", lines=3,
                                          placeholder=_start.placeholder)
                    t_negative = gr.Textbox(label="Negative prompt (optional)", lines=1,
                                            placeholder="lo-fi, distorted, vocals")
                    t_seconds = gr.Slider(1, _start_max, value=_start.default_seconds,
                                          step=1, label=f"Seconds · max {_start_max}")
                    with gr.Accordion("Advanced", open=False):
                        t_steps = gr.Slider(1, 100, value=8, step=1, label="Steps")
                        t_cfg = gr.Slider(0.0, 15.0, value=1.0, step=0.1, label="CFG scale")
                        t_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
                    t_btn = gr.Button("Generate", variant="primary", size="lg")
                with gr.Column(scale=1):
                    t_audio = gr.Audio(label="Output", type="filepath")
                    t_meta = gr.Markdown("")

            t_variant.change(_on_variant_change, [t_variant], [t_seconds, t_prompt])
            t_btn.click(
                ui_text_to_audio,
                [t_variant, t_prompt, t_negative, t_seconds, t_steps, t_cfg, t_seed],
                [t_audio, t_meta],
            )

        # ---------------- Audio → audio ----------------
        with gr.Tab("Audio → audio"):
            a_variant = gr.Radio(VARIANT_CHOICES, value=DEFAULT_VARIANT, label="Model")
            with gr.Row():
                with gr.Column(scale=2):
                    a_init = gr.Audio(label="Your audio", type="filepath")
                    a_noise = gr.Slider(0.0, 1.0, value=0.4, step=0.05,
                                        label="How far to travel from your audio",
                                        info=NOISE_INFO)
                    a_prompt = gr.Textbox(label="Prompt", lines=3,
                                          placeholder=_start.placeholder)
                    a_negative = gr.Textbox(label="Negative prompt (optional)", lines=1)
                    a_seconds = gr.Slider(1, _start_max, value=_start.default_seconds,
                                          step=1, label=f"Seconds · max {_start_max}")
                    with gr.Accordion("Advanced", open=False):
                        a_steps = gr.Slider(1, 100, value=8, step=1, label="Steps")
                        a_cfg = gr.Slider(0.0, 15.0, value=1.0, step=0.1, label="CFG scale")
                        a_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
                    a_btn = gr.Button("Transform", variant="primary", size="lg")
                with gr.Column(scale=1):
                    a_audio = gr.Audio(label="Output", type="filepath")
                    a_meta = gr.Markdown("")

            a_variant.change(_on_variant_change, [a_variant], [a_seconds, a_prompt])
            a_btn.click(
                ui_audio_to_audio,
                [a_variant, a_init, a_prompt, a_negative, a_noise,
                 a_seconds, a_steps, a_cfg, a_seed],
                [a_audio, a_meta],
            )

    gr.Markdown(FOOTER)

    # ── Headless API surface ────────────────────────────────────────────
    # Hidden components + hidden buttons carry the api_name routes. Visibility
    # is UI-only; the endpoints are registered server-side regardless, and this
    # plays nicely with gradio_client.handle_file for the a2a upload.
    with gr.Group(visible=False):
        _g_token = gr.Textbox(value="", label="token")
        _g_variant = gr.Textbox(value=DEFAULT_VARIANT, label="model_variant")
        _g_prompt = gr.Textbox(value="", label="prompt")
        _g_negative = gr.Textbox(value="", label="negative_prompt")
        _g_seconds = gr.Number(value=30, label="seconds")
        _g_steps = gr.Number(value=8, label="steps")
        _g_cfg = gr.Number(value=1.0, label="cfg_scale")
        _g_seed = gr.Number(value=-1, label="seed")
        _g_file = gr.File(label="audio out")
        _g_json = gr.JSON(label="meta")
        _g_btn = gr.Button("generate")
        _g_btn.click(
            fn=api_generate,
            inputs=[_g_token, _g_variant, _g_prompt, _g_negative,
                    _g_seconds, _g_steps, _g_cfg, _g_seed],
            outputs=[_g_file, _g_json],
            api_name="generate",
            show_progress="hidden",
        )

        _a_token = gr.Textbox(value="", label="token")
        _a_variant = gr.Textbox(value=DEFAULT_VARIANT, label="model_variant")
        _a_init = gr.File(type="filepath", label="init_audio",
                          file_types=["audio", ".wav", ".mp3", ".flac", ".ogg",
                                      ".m4a", ".aiff", ".aif", ".opus"])
        _a_prompt = gr.Textbox(value="", label="prompt")
        _a_negative = gr.Textbox(value="", label="negative_prompt")
        _a_noise = gr.Number(value=0.4, label="init_noise_level")
        _a_seconds = gr.Number(value=30, label="seconds")
        _a_steps = gr.Number(value=8, label="steps")
        _a_cfg = gr.Number(value=1.0, label="cfg_scale")
        _a_seed = gr.Number(value=-1, label="seed")
        _a_file = gr.File(label="audio out")
        _a_json = gr.JSON(label="meta")
        _a_btn = gr.Button("generate_a2a")
        _a_btn.click(
            fn=api_generate_a2a,
            inputs=[_a_token, _a_variant, _a_init, _a_prompt, _a_negative,
                    _a_noise, _a_seconds, _a_steps, _a_cfg, _a_seed],
            outputs=[_a_file, _a_json],
            api_name="generate_a2a",
            show_progress="hidden",
        )


if __name__ == "__main__":
    demo.queue(max_size=32).launch(
        server_name=os.environ.get("SA3_HOST", "0.0.0.0"),
        server_port=int(os.environ.get("SA3_PORT", "7860")),
    )