"""Chora separation backend — Hugging Face ZeroGPU (Gradio SDK) entrypoint.

The FastAPI server (main.py) is for local dev; this file wraps the same
separator for HF Spaces, where free GPU time is only available to Gradio
apps via @spaces.GPU. The frontend talks to it with @gradio/client.
"""

import tempfile

import gradio as gr

try:
    import spaces  # available on HF Spaces; degrade gracefully locally

    GPU = spaces.GPU
except ImportError:  # local run: no-op decorator
    def GPU(*_args, **_kwargs):
        def wrap(fn):
            return fn
        return wrap

from separator import separate

STEM_ORDER = ("vocals", "drums", "bass", "other")
MAX_DURATION_S = 8 * 60  # refuse absurdly long uploads


def _gpu_seconds(audio_path: str, *_args, **_kwargs) -> int:
    """Reserve GPU time proportional to song length instead of a flat
    2 minutes — ZeroGPU quota is charged per reservation, so this makes
    each call ~3x cheaper for typical songs."""
    try:
        import soundfile as sf

        info = sf.info(audio_path)
        song_s = info.frames / info.samplerate
    except Exception:
        song_s = 240  # unknown container — assume ~4 min
    return int(min(90, max(25, 10 + song_s * 0.2)))


@GPU(duration=_gpu_seconds)  # H200 slice: a 3-4 min song separates in ~15-30 s
def predict(audio_path: str, progress=gr.Progress()) -> list[str]:
    if audio_path is None:
        raise gr.Error("No audio received")

    # Duration guard. soundfile can't read every container (e.g. m4a) —
    # if it can't, skip the check and let demucs handle decoding.
    try:
        import soundfile as sf

        info = sf.info(audio_path)
        duration = info.frames / info.samplerate
        if duration > MAX_DURATION_S:
            raise gr.Error(f"Track too long ({duration/60:.1f} min). Max is {MAX_DURATION_S//60} min.")
    except gr.Error:
        raise
    except Exception:
        pass

    out_dir = tempfile.mkdtemp(prefix="chora-")
    progress(0.0)

    # Throttle: demucs fires progress hundreds of times; every call crosses
    # the ZeroGPU process bridge, so only forward meaningful steps.
    last = -1.0

    def on_progress(f: float) -> None:
        nonlocal last
        if f - last >= 0.02:
            last = f
            progress(f)

    stems = separate(
        input_path=audio_path,
        output_dir=out_dir,
        device="cuda",
        output_format="mp3",  # ~10x smaller download than wav
        on_progress=on_progress,
    )
    return [str(stems[s]) for s in STEM_ORDER]


demo = gr.Interface(
    fn=predict,
    inputs=gr.Audio(type="filepath", label="song"),
    outputs=[gr.Audio(type="filepath", label=s) for s in STEM_ORDER],
    title="chora backend",
    description="4-stem source separation (Demucs htdemucs). Called by the chora frontend over the REST API.",
    flagging_mode="never",
    api_name="predict",  # frontend calls /gradio_api/call/predict
)

if __name__ == "__main__":
    demo.launch()