"""ZeroGPU strip-pass benchmark — throwaway test app.

Standalone Gradio app, NOT part of the production site. Runs the exact same
vocal-isolation model the live Voice Remover's Step 1 uses
(mel_band_roformer_vocals_becruily) and times it on this Space's ZeroGPU
hardware, so we can decide with real numbers whether ZeroGPU is viable as a
backend before touching anything production-facing.

Deploy: create/point a ZeroGPU-hardware Space at this folder (sdk: gradio,
see README.md). Upload a song, click "Run strip pass".
"""
import os
import shutil
import subprocess
import tempfile
import time
import traceback
from pathlib import Path

import spaces
import gradio as gr
import torch
import torchaudio
from audio_separator.separator import Separator

MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/tmp/models"))
MODELS_DIR.mkdir(parents=True, exist_ok=True)

# Same model the live site's Step 1 ("Isolate vocals — becruily") uses.
VOCAL_MODEL = "mel_band_roformer_vocals_becruily.ckpt"

# Fetch the checkpoint once at Space startup — CPU-only, no GPU needed — so the
# timed run below measures disk->GPU load, not an internet download. Wrapped in
# try/except so a transient network hiccup at boot can't crash the whole Space;
# worst case the first real run below pays the download cost instead.
try:
    print(f"[startup] fetching {VOCAL_MODEL} ...", flush=True)
    _warm = Separator(output_dir="/tmp", model_file_dir=str(MODELS_DIR))
    _warm.load_model(model_filename=VOCAL_MODEL)
    del _warm
    print("[startup] model cached on disk.", flush=True)
except Exception:
    print("[startup] warm-up download failed, will retry on first run:", flush=True)
    print(traceback.format_exc(), flush=True)


def _normalize(src: str) -> str:
    """Approximates production's _normalize_for_sep (stereo / 32-bit float /
    min 30s) so the timed input shape matches what RoFormer sees live.

    Goes through ffmpeg first — same as production's load_audio — because
    libsndfile (torchaudio/soundfile's backend) doesn't reliably decode mp3/m4a
    uploads, only wav/flac. Falls back to torchaudio directly if ffmpeg isn't
    on this Space's image (covers wav/flac uploads either way)."""
    out = tempfile.mktemp(suffix=".wav")
    ff = shutil.which("ffmpeg")
    if ff:
        # -c:a pcm_f32le (not just -sample_fmt flt): ffmpeg's default codec for a
        # .wav target is pcm_s16le, which only accepts s16 and hard-fails if you
        # ask it for float samples. Naming the float encoder directly sidesteps
        # that mismatch instead of relying on ffmpeg to pick a compatible one.
        try:
            subprocess.run(
                [ff, "-y", "-i", src,
                 "-ac", "2", "-c:a", "pcm_f32le",
                 "-af", "apad=whole_dur=30",
                 out],
                check=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            err = e.stderr.decode("utf-8", "replace")[-1000:] if e.stderr else "(no stderr captured)"
            raise RuntimeError(f"ffmpeg failed on the uploaded file:\n{err}") from e
        return out

    wf, sr = torchaudio.load(src)
    if wf.shape[0] == 1:
        wf = wf.repeat(2, 1)
    elif wf.shape[0] > 2:
        wf = wf[:2]
    min_samples = sr * 30
    if wf.shape[1] < min_samples:
        pad = torch.zeros(wf.shape[0], min_samples - wf.shape[1], dtype=wf.dtype)
        wf = torch.cat([wf, pad], dim=1)
    torchaudio.save(out, wf, sr, encoding="PCM_F", bits_per_sample=32)
    return out


# Quota is billed on actual seconds used, not this ceiling — but ZeroGPU ranks
# shorter declared durations higher in its queue, so keep this close to real
# observed usage rather than padded to the 120s max. 60s covers up to roughly a
# 6-7 min song at the ~0.13s-inference-per-audio-second rate we measured on a
# ~3 min test song; bump it back up if you test something much longer.
@spaces.GPU(duration=60)
def run_benchmark(song_path):
    if song_path is None:
        return "Upload a song first."

    work_dir = tempfile.mkdtemp(prefix="bench_")
    try:
        gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "no GPU visible"

        t0 = time.perf_counter()
        sep = Separator(output_dir=work_dir, output_format="WAV", model_file_dir=str(MODELS_DIR))
        sep.load_model(model_filename=VOCAL_MODEL)
        t1 = time.perf_counter()

        sep_input = _normalize(song_path)
        out_names = sep.separate(sep_input)
        t2 = time.perf_counter()

        load_s, sep_s, total_s = t1 - t0, t2 - t1, t2 - t0
        runs_per_day = max(1, int(120 // total_s)) if total_s > 0 else "?"

        return (
            f"GPU: {gpu_name}\n\n"
            f"Model load + move-to-device: {load_s:.1f}s\n"
            f"Separation (inference):      {sep_s:.1f}s\n"
            f"Total:                       {total_s:.1f}s\n\n"
            f"Output stems: {out_names}\n\n"
            f"2 min/day free quota ÷ this total ≈ {runs_per_day} run(s)/day per visitor "
            f"(before counting any region splits on top of this strip pass)."
        )
    except Exception:
        return "Error:\n" + traceback.format_exc()
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)


with gr.Blocks(title="ZeroGPU strip-pass benchmark") as demo:
    gr.Markdown(
        "### ZeroGPU strip-pass benchmark\n"
        "Throwaway test app — not the production site. Upload a song and run "
        "the same vocal-isolation model the live site's Step 1 uses, timed on "
        "this Space's GPU, to get real numbers on whether ZeroGPU is viable."
    )
    audio_in = gr.Audio(label="Song", type="filepath")
    run_btn = gr.Button("Run strip pass", variant="primary")
    result = gr.Textbox(label="Timing", lines=9)
    run_btn.click(fn=run_benchmark, inputs=audio_in, outputs=result)

demo.queue().launch()