#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import html
import importlib.metadata
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import sysconfig
import traceback
import zipfile
from collections import defaultdict, deque
from pathlib import Path
from urllib.parse import quote


APP_DIR = Path(__file__).resolve().parent
OUTPUT_ROOT = APP_DIR / "separated"
WORK_ROOT = APP_DIR / ".stem_work"

MODEL_BY_MODE = {
    "2": "roformer-model-bs-roformer-leap-vocals-by-pcunwa",
    "4": "roformer-model-bs-roformer-musdb18hq-by-zfturbo",
    "6": "roformer-model-bs-roformer-sw-by-jarredou",
    "53": "roformer-model-bs-roformer-mvsep-mega-53-stems",
}

EXPECTED_STEMS = {
    "2": ["vocals", "background"],
    "4": ["vocals", "drums", "bass", "other"],
    "6": ["vocals", "drums", "bass", "guitar", "piano", "other"],
}

MODE_CHOICES = [
    "Voice + music — 2 tracks",
    "Core instruments — 4 tracks",
    "Full band — 6 tracks",
    "Every detail — 53 tracks",
]

MODE_BY_CHOICE = dict(zip(MODE_CHOICES, ["2", "4", "6", "53"]))


def _run(cmd: list[str], cwd: Path | None = None) -> None:
    print("+", " ".join(map(str, cmd)), flush=True)
    subprocess.run([str(item) for item in cmd], cwd=cwd, check=True)


def validate_runtime() -> tuple[str, str]:
    try:
        version = importlib.metadata.version("bs-roformer-infer")
    except importlib.metadata.PackageNotFoundError as exc:
        raise RuntimeError(
            "Dependencies are not installed. Run: python -m pip install -r requirements.txt"
        ) from exc
    backend = "mlx" if sys.platform == "darwin" and platform.machine() == "arm64" else "torch"
    return version, backend


RUNTIME_COMMIT, COMPUTE_BACKEND = validate_runtime()

import gradio as gr
import imageio_ffmpeg
import numpy as np
import soundfile as sf
if COMPUTE_BACKEND == "mlx":
    import mlx.core as mx
else:
    import torch
if os.environ.get("SPACE_ID"):
    import spaces
else:
    class _LocalSpaces:
        @staticmethod
        def GPU(*_args, **_kwargs):
            return lambda fn: fn

    spaces = _LocalSpaces()
from bs_roformer import MODEL_REGISTRY, ensure_model_assets


def sha256_file(path: Path, block_size: int = 8 * 1024 * 1024) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while chunk := f.read(block_size):
            h.update(chunk)
    return h.hexdigest()


def safe_stem_name(name: str) -> str:
    cleaned = "".join(c if c.isalnum() or c in "._-" else "_" for c in name)
    cleaned = cleaned.strip("._")
    return cleaned or "audio"


def normalized_token(path: Path) -> str:
    return re.sub(r"[^a-z0-9]+", "_", path.stem.lower()).strip("_")


def output_stem_name(path: Path) -> str:
    """Derive the model's stem label from its rendered WAV filename."""
    return re.sub(r"^input[_ -]", "", path.stem.lower())


def best_match(stem_name: str, paths: list[Path]) -> Path | None:
    stem = stem_name.lower()
    matches = []
    for path in paths:
        token = normalized_token(path)
        parts = token.split("_")
        if stem in parts:
            rank = 0
        elif token.endswith(stem) or token.startswith(stem):
            rank = 1
        elif stem in token:
            rank = 2
        else:
            continue
        matches.append((rank, len(path.name), path))

    if not matches:
        return None

    matches.sort(key=lambda x: (x[0], x[1]))
    return matches[0][2]


def copy_atomic(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_suffix(dst.suffix + ".tmp")
    tmp.unlink(missing_ok=True)
    shutil.copy2(src, tmp)
    tmp.replace(dst)


def stage_audio(input_audio: Path, input_hash: str) -> Path:
    stage_dir = WORK_ROOT / "inputs" / input_hash[:16]
    stage_dir.mkdir(parents=True, exist_ok=True)
    stage_wav = stage_dir / "input.wav"
    stage_meta = stage_dir / "input.json"

    valid = False
    if stage_wav.exists() and stage_meta.exists():
        try:
            meta = json.loads(stage_meta.read_text())
            valid = meta.get("sha256") == input_hash
        except Exception:
            pass

    if valid:
        return stage_wav

    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
    tmp_wav = stage_dir / "input.tmp.wav"
    tmp_wav.unlink(missing_ok=True)

    cmd = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "warning",
        "-y",
        "-i",
        str(input_audio),
        "-vn",
        "-ac",
        "2",
        "-ar",
        "44100",
        "-c:a",
        "pcm_f32le",
        str(tmp_wav),
    ]
    _run(cmd)
    tmp_wav.replace(stage_wav)

    stage_meta.write_text(json.dumps({
        "source_name": input_audio.name,
        "sha256": input_hash,
        "sample_rate": 44100,
        "channels": 2,
    }, indent=2))
    return stage_wav


def cli_path() -> Path:
    found = shutil.which("bs-roformer-infer")
    if found:
        return Path(found)

    candidate = Path(sysconfig.get_path("scripts")) / "bs-roformer-infer"
    if candidate.exists():
        return candidate

    raise RuntimeError("Could not locate the bs-roformer-infer console script.")


def run_inference(stage_wav: Path, raw_dir: Path, model: str) -> None:
    infer_input_dir = WORK_ROOT / "jobs" / stage_wav.parent.name / model / "input"
    infer_input_dir.mkdir(parents=True, exist_ok=True)
    staged_link = infer_input_dir / "input.wav"

    if staged_link.exists() or staged_link.is_symlink():
        staged_link.unlink()
    try:
        staged_link.symlink_to(stage_wav.resolve())
    except OSError:
        shutil.copy2(stage_wav, staged_link)

    if raw_dir.exists():
        shutil.rmtree(raw_dir)
    raw_dir.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(cli_path()),
        "--input_folder",
        str(infer_input_dir),
        "--store_dir",
        str(raw_dir),
        "--backend",
        COMPUTE_BACKEND,
    ]
    if COMPUTE_BACKEND == "mlx":
        checkpoint_path, config_path = mlx_compatible_assets(model)
        cmd.extend([
            "--model_path",
            str(checkpoint_path),
            "--config_path",
            str(config_path),
        ])
    else:
        cmd.extend(["--model", model])

    print("+", " ".join(cmd), flush=True)
    tail: deque[str] = deque(maxlen=120)

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    assert process.stdout is not None
    for line in process.stdout:
        line = line.rstrip()
        if line:
            print(line, flush=True)
            tail.append(line)

    return_code = process.wait()
    if return_code != 0:
        raise RuntimeError(
            f"bs-roformer-infer exited with code {return_code}.\n\n"
            + "\n".join(tail)
        )


def ensure_healthy_model_assets(model: str) -> None:
    checkpoint, _ = ensure_model_assets(model)
    if zipfile.is_zipfile(checkpoint):
        return

    print(f"Cached checkpoint is corrupt; downloading a fresh copy: {checkpoint}")
    checkpoint.unlink(missing_ok=True)
    ensure_model_assets(model)


def mlx_compatible_assets(model: str) -> tuple[Path, Path]:
    """Return model assets, patching MLX-incompatible chunk sizes when needed."""
    checkpoint, config_path = map(Path, ensure_model_assets(model))
    config_text = config_path.read_text()
    chunk_matches = list(re.finditer(
        r"(?m)^([ \t]*chunk_size:[ \t]*)(\d+)([^\n]*)$",
        config_text,
    ))
    hop_match = re.search(
        r"(?m)^[ \t]*stft_hop_length:[ \t]*(\d+)",
        config_text,
    )
    if not chunk_matches or hop_match is None:
        raise RuntimeError("The detailed separation settings are incomplete.")

    # Published configs keep the chunked-inference setting in the audio section.
    # Use the first entry instead of the model's unrelated training metadata.
    chunk_match = chunk_matches[0]
    chunk_size = int(chunk_match.group(2))
    hop_length = int(hop_match.group(1))
    if chunk_size % hop_length == 0:
        return checkpoint, config_path

    aligned_chunk_size = chunk_size - (chunk_size % hop_length)
    patched_text = (
        config_text[:chunk_match.start(2)]
        + str(aligned_chunk_size)
        + config_text[chunk_match.end(2):]
    )
    patched_dir = WORK_ROOT / "configs"
    patched_dir.mkdir(parents=True, exist_ok=True)
    patched_path = patched_dir / f"{safe_stem_name(model)}.mlx.yaml"
    if not patched_path.exists() or patched_path.read_text() != patched_text:
        patched_path.write_text(patched_text)
    return checkpoint, patched_path


def normalize_outputs(
    mode: str,
    stage_wav: Path,
    raw_dir: Path,
    job_dir: Path,
) -> tuple[list[Path], dict[str, Path]]:
    raw_wavs = sorted(raw_dir.rglob("*.wav"))
    if not raw_wavs:
        raise RuntimeError(f"No WAV outputs found in {raw_dir}")

    previews: dict[str, Path] = {}

    if mode == "2":
        final_dir = job_dir / "stems"
        final_dir.mkdir(parents=True, exist_ok=True)

        vocals_src = best_match("vocals", raw_wavs)
        if vocals_src is None:
            raise RuntimeError(
                "Could not identify the vocals output.\n"
                + "\n".join(str(p) for p in raw_wavs)
            )

        vocals_dst = final_dir / "01_vocals.wav"
        background_dst = final_dir / "02_background.wav"
        copy_atomic(vocals_src, vocals_dst)

        instrumental_src = (
            best_match("instrumental", raw_wavs)
            or best_match("background", raw_wavs)
        )

        if instrumental_src is not None:
            copy_atomic(instrumental_src, background_dst)
        else:
            mix, sr_mix = sf.read(stage_wav, dtype="float32", always_2d=True)
            vocals, sr_vocals = sf.read(vocals_src, dtype="float32", always_2d=True)
            if sr_mix != sr_vocals:
                raise RuntimeError(
                    f"Sample-rate mismatch: mixture={sr_mix}, vocals={sr_vocals}"
                )
            n = min(len(mix), len(vocals))
            background = mix[:n] - vocals[:n]
            sf.write(background_dst, background, sr_mix, subtype="FLOAT")

        files = [vocals_dst, background_dst]
        previews = {
            "vocals": vocals_dst,
            "background": background_dst,
        }
        return files, previews

    if mode in {"4", "6"}:
        final_dir = job_dir / "stems"
        final_dir.mkdir(parents=True, exist_ok=True)

        files = []
        for index, stem in enumerate(EXPECTED_STEMS[mode], 1):
            src = best_match(stem, raw_wavs)
            if src is None:
                raise RuntimeError(
                    f"Could not identify stem {stem!r}.\n"
                    + "\n".join(str(p) for p in raw_wavs)
                )
            dst = final_dir / f"{index:02d}_{stem}.wav"
            copy_atomic(src, dst)
            files.append(dst)
            previews[stem] = dst

        return files, previews

    # The detailed model emits a distinct WAV for every one of its 53 stems.
    # Keep every source in the mixer instead of showing only a small preview.
    for src in raw_wavs:
        stem = output_stem_name(src)
        if stem in previews:
            stem = src.stem.lower()
        previews[stem] = src

    return raw_wavs, previews


def make_zip(files: list[Path], job_dir: Path) -> Path:
    zip_path = job_dir / "stems.zip"
    tmp_zip = job_dir / "stems.zip.tmp"
    tmp_zip.unlink(missing_ok=True)

    with zipfile.ZipFile(tmp_zip, "w", compression=zipfile.ZIP_STORED) as zf:
        seen: set[str] = set()
        for path in files:
            arcname = path.name
            if arcname in seen:
                arcname = f"{path.parent.name}_{path.name}"
            seen.add(arcname)
            zf.write(path, arcname=arcname)

    tmp_zip.replace(zip_path)
    return zip_path


STEM_LABELS = {
    "vocals": "Vocals",
    "background": "Background",
    "drums": "Drums",
    "bass": "Bass",
    "guitar": "Guitar",
    "piano": "Piano",
    "other": "Other",
    "synth": "Synth",
}

GROUP_LABELS = {
    "vocals": "Vocals",
    "rhythm": "Rhythm",
    "guitars_keys": "Guitars & keys",
    "orchestral_winds": "Orchestral & winds",
    "other": "Other instruments",
}

STEM_GROUPS = {
    "vocals": {"vocal", "vocals", "lead-vocal", "back-vocal"},
    "rhythm": {"drums", "kick", "snare", "toms", "hh", "percussion", "congas", "tambourine", "timpani", "triangle", "bass", "double-bass"},
    "guitars_keys": {"acoustic-guitar", "electric-guitar", "guitar", "banjo", "dobro", "mandolin", "ukulele", "sitar", "harp", "piano", "digital-piano", "keys", "organ", "harpsichord", "accordion", "synth", "marimba", "glockenspiel", "bells"},
    "orchestral_winds": {"bowed_strings", "strings", "cello", "viola", "violin", "brass", "trumpet", "trombone", "tuba", "french-horn", "flute", "clarinet", "oboe", "saxophone", "bassoon", "harmonica", "wind", "wind-chimes", "woodwind"},
}

# These broad predictions overlap the more specific instrument predictions in
# the same group. They are a convenient mix layer, not an extra source to add.
AGGREGATE_STEMS = {
    "vocal", "vocals",
    "drums", "percussion",
    "guitar", "keys",
    "bowed_strings", "strings", "brass", "wind", "woodwind",
}

ICON_PATHS = {
    "sparkles": '<path d="m12 3-1.9 4.8L5 10l5.1 2.2L12 17l1.9-4.8L19 10l-5.1-2.2L12 3Z"/><path d="M5 3v4M3 5h4M19 17v4m-2-2h4"/>',
    "play": '<path d="m8 5 11 7-11 7Z"/>',
    "pause": '<path d="M9 5v14M15 5v14"/>',
    "download": '<path d="M12 3v12m0 0 4-4m-4 4-4-4M5 19h14"/>',
    "speaker": '<path d="M11 5 6 9H3v6h3l5 4Z"/><path d="M15.5 8.5a5 5 0 0 1 0 7M18 6a8.5 8.5 0 0 1 0 12"/>',
    "speaker-muted": '<path d="M11 5 6 9H3v6h3l5 4Z"/><path d="m16 9 5 5m0-5-5 5"/>',
    "vocals": '<rect x="9" y="2" width="6" height="12" rx="3"/><path d="M5 10a7 7 0 0 0 14 0M12 17v4M8 21h8"/>',
    "background": '<path d="m12 2 9 5-9 5-9-5 9-5Z"/><path d="m3 12 9 5 9-5M3 17l9 5 9-5"/>',
    "drums": '<ellipse cx="12" cy="6" rx="8" ry="3"/><path d="M4 6v11c0 1.7 3.6 3 8 3s8-1.3 8-3V6M8 9v10M16 9v10"/>',
    "bass": '<path d="M3 12h3l2-6 4 12 3-9 2 6h4"/>',
    "guitar": '<path d="m14 5 5-3 3 3-3 5M13 6l5 5M9 10c-3-1-7 2-7 6 0 3 3 6 6 6 4 0 7-4 6-7l4-4-5-5-4 4Z"/>',
    "piano": '<rect x="3" y="5" width="18" height="14" rx="2"/><path d="M7 5v9M11 5v9M15 5v9M19 5v9M5 14h14"/>',
    "other": '<circle cx="8" cy="8" r="4"/><rect x="12" y="12" width="8" height="8" rx="2"/><path d="m5 20 4-7 4 7Z"/>',
    "synth": '<path d="M4 5h16M4 12h16M4 19h16"/><circle cx="9" cy="5" r="2"/><circle cx="15" cy="12" r="2"/><circle cx="7" cy="19" r="2"/>',
}


def icon_svg(name: str, class_name: str = "ui-icon") -> str:
    return (
        f'<svg class="{class_name}" viewBox="0 0 24 24" fill="none" '
        f'stroke="currentColor" stroke-width="1.7" stroke-linecap="round" '
        f'stroke-linejoin="round" aria-hidden="true">{ICON_PATHS[name]}</svg>'
    )


def stem_group(stem: str) -> str:
    for group, stems in STEM_GROUPS.items():
        if stem in stems:
            return group
    return "other"


def stem_label(stem: str) -> str:
    return STEM_LABELS.get(stem, stem.replace("-", " ").replace("_", " ").title())


def mixer_html(previews: dict[str, Path], mode: str) -> str:
    """Build native audio controls; Gradio's waveform player has no controllable audio src."""
    detailed_mode = mode == "53"
    flat_mode = mode == "6"
    grouped: dict[str, list[tuple[str, Path]]] = defaultdict(list)
    for stem, path in previews.items():
        grouped[stem_group(stem)].append((stem, path))

    if not grouped:
        return '<p class="empty-mixer">No playable stems were produced.</p>'

    controls = []
    tracks = []
    for group in GROUP_LABELS:
        group_stems = grouped.get(group)
        if not group_stems:
            continue
        group_label = GROUP_LABELS[group]
        escaped_group_label = html.escape(group_label)
        group_controls = []
        group_tracks = []
        for stem, path in sorted(group_stems, key=lambda item: stem_label(item[0])):
            label = stem_label(stem)
            url = "/gradio_api/file=" + quote(str(path.resolve()), safe="/")
            escaped_label = html.escape(label)
            escaped_filename = html.escape(path.name, quote=True)
            icon_name = stem if stem in ICON_PATHS else "other"
            aggregate = detailed_mode and stem in AGGREGATE_STEMS
            state = "false" if aggregate else "true"
            action = "Unmute" if aggregate else "Mute"
            button_icon = "speaker-muted" if aggregate else "speaker"
            button_text = "Muted" if aggregate else "On"
            tier = "aggregate" if aggregate else "individual"
            group_controls.append(
                f'<div class="mixer-stem" data-tier="{tier}"><span class="stem-name">{icon_svg(icon_name)}{escaped_label}</span>'
                f'<button class="mixer-toggle" data-stem="{stem}" data-label="{escaped_label}" data-tier="{tier}" '
                f'aria-pressed="{state}" aria-label="{action} {escaped_label}" title="{action} {escaped_label}">'
                f'{icon_svg(button_icon)}<span>{button_text}</span></button></div>'
            )
            group_tracks.append(
                f'<article class="stem-track" data-stem="{stem}" data-group="{group}" data-tier="{tier}" '
                f'data-initial-muted="{str(aggregate).lower()}" data-src="{html.escape(url, quote=True)}">'
                f'<div class="stem-track-header"><span class="stem-track-label">{icon_svg(icon_name)}{escaped_label}</span>'
                f'<a class="stem-download" href="{html.escape(url, quote=True)}" '
                f'download="{escaped_filename}">{icon_svg("download")}Save</a></div>'
                f'<div class="track-timeline" aria-hidden="true"><span class="track-playhead"></span></div>'
                f'</article>'
            )
        group_header = (
            f'<section class="mixer-group"><div class="mixer-group-heading"><h4>{escaped_group_label}</h4>'
            f'<button class="mixer-group-toggle" data-group="{group}" data-label="{escaped_group_label}" '
            f'data-detailed="{str(detailed_mode).lower()}" '
            f'aria-pressed="true" aria-label="Mute all {escaped_group_label}" title="Mute all {escaped_group_label}">'
            f'{icon_svg("speaker")}<span>Mute all</span></button></div>'
        )
        if flat_mode:
            controls.extend(group_controls)
            tracks.extend(group_tracks)
        elif detailed_mode:
            controls.append(
                group_header + f'<details><summary>{len(group_stems)} tracks</summary>{"".join(group_controls)}</details></section>'
            )
            tracks.append(
                f'<details class="stem-track-group"{" open" if group == "vocals" else ""}>'
                f'<summary><span>{escaped_group_label}</span><span>{len(group_stems)} tracks</span></summary>'
                f'<div class="stem-track-grid">{"".join(group_tracks)}</div></details>'
            )
        else:
            controls.append(group_header + "".join(group_controls) + "</section>")
            tracks.append(
                f'<section class="stem-track-group is-expanded"><div class="stem-track-group-title">'
                f'<span>{escaped_group_label}</span><span>{len(group_stems)} tracks</span></div>'
                f'<div class="stem-track-grid">{"".join(group_tracks)}</div></section>'
            )

    track_markup = (
        f'<div class="stem-track-grid full-grid">{"".join(tracks)}</div>'
        if flat_mode
        else "".join(tracks)
    )

    return """
    <section id="sync-mixer" aria-label="Synchronized stem mixer">
      <aside id="mixer-sidebar">
        <h3>Mix</h3>
        <div class="transport-row">
          <button class="mixer-transport" data-action="play">%s Play</button>
          <button class="mixer-transport" data-action="pause">%s Pause</button>
        </div>
        %s
      </aside>
      <div id="stem-players">
        <section class="shared-timeline" aria-label="Shared playback timeline">
          <div class="timeline-labels"><span id="timeline-current">0:00</span><span id="timeline-total">0:00</span></div>
          <input id="timeline-scrubber" type="range" min="0" max="1000" value="0" step="1" aria-label="Seek all stems">
        </section>
        <div id="stem-track-list">%s</div>
      </div>
    </section>
    """ % (icon_svg("play"), icon_svg("pause"), "".join(controls), track_markup)


def mode_from_choice(choice: str) -> str:
    if not choice:
        return "6"
    return MODE_BY_CHOICE.get(choice, "6")


@spaces.GPU(duration=120)
def separate(
    audio_path: str | None,
    mode_choice: str,
    force: bool,
    progress=gr.Progress(),
):
    try:
        if not audio_path:
            raise gr.Error("Upload an audio file first.")

        mode = mode_from_choice(mode_choice)
        if mode not in MODEL_BY_MODE:
            raise gr.Error(f"Unsupported mode: {mode}")

        model = MODEL_BY_MODE[mode]
        try:
            MODEL_REGISTRY.get(model)
        except KeyError as exc:
            raise RuntimeError(
                f"The installed bs-roformer-infer package does not provide {model!r}. "
                "Use this project's environment: "
                "source .venv/bin/activate && python -m pip install -r requirements.txt"
            ) from exc

        input_audio = Path(audio_path)
        if not input_audio.exists():
            raise gr.Error("The uploaded audio file is no longer available.")

        progress(0.05, desc="Reading song")
        input_hash = sha256_file(input_audio)
        clean_name = safe_stem_name(input_audio.stem)

        # The filename is only a display convenience. Look for prior results by
        # content hash first, so a re-upload (or a renamed copy) reuses the same
        # separation instead of needlessly running ML again.
        job_dir = OUTPUT_ROOT / clean_name / f"{mode}-stem" / input_hash[:16]
        job_key = {
            "input_sha256": input_hash,
            "model": model,
            "runtime_commit": RUNTIME_COMMIT,
            "backend": COMPUTE_BACKEND,
        }

        if not force:
            for candidate in OUTPUT_ROOT.glob(f"*/{mode}-stem/{input_hash[:16]}"):
                candidate_sentinel = candidate / ".complete.json"
                candidate_raw = candidate / "raw"
                try:
                    previous = json.loads(candidate_sentinel.read_text())
                    valid = (
                        all(previous.get(k) == v for k, v in job_key.items())
                        and any(candidate_raw.rglob("*.wav"))
                    )
                except Exception:
                    valid = False
                if valid:
                    job_dir = candidate
                    break

        raw_dir = job_dir / "raw"
        sentinel = job_dir / ".complete.json"
        job_dir.mkdir(parents=True, exist_ok=True)

        progress(0.15, desc="Preparing audio")
        stage_wav = stage_audio(input_audio, input_hash)

        cached = False
        if sentinel.exists() and not force:
            try:
                previous = json.loads(sentinel.read_text())
                cached = (
                    all(previous.get(k) == v for k, v in job_key.items())
                    and raw_dir.exists()
                    and any(raw_dir.rglob("*.wav"))
                )
            except Exception:
                cached = False

        if cached:
            progress(0.70, desc="Opening tracks")
        else:
            progress(0.25, desc="Getting ready")
            ensure_healthy_model_assets(model)

            progress(0.40, desc="Separating")
            run_inference(stage_wav, raw_dir, model)

            if not any(raw_dir.rglob("*.wav")):
                raise RuntimeError("Inference completed but produced no WAV files.")

            sentinel.write_text(json.dumps({
                **job_key,
                "source_name": input_audio.name,
            }, indent=2))

        progress(0.82, desc="Finishing tracks")
        files, previews = normalize_outputs(mode, stage_wav, raw_dir, job_dir)

        progress(0.92, desc="Preparing downloads")
        zip_path = make_zip(files, job_dir)

        runtime_test = (
            float(mx.sum(mx.array([1.0, 2.0, 3.0])).item())
            if COMPUTE_BACKEND == "mlx"
            else float(torch.tensor([1.0, 2.0, 3.0]).sum().item())
        )
        if runtime_test != 6.0:
            raise RuntimeError("Audio processing runtime check failed.")

        status = f"### Ready to mix\n{len(files)} tracks available."

        progress(1.0, desc="Done")

        return (
            status,
            [str(path) for path in files],
            str(zip_path),
            mixer_html(previews, mode),
        )

    except gr.Error:
        raise
    except Exception as exc:
        traceback.print_exc()
        raise gr.Error("We couldn't split this song. Please try again.") from exc


runtime_probe = (
    float(mx.sum(mx.array([1.0, 2.0, 3.0])).item())
    if COMPUTE_BACKEND == "mlx"
    else float(torch.tensor([1.0, 2.0, 3.0]).sum().item())
)
if runtime_probe != 6.0:
    raise RuntimeError("Audio processing runtime check failed.")


MIXER_HEAD = """
<script>
(() => {
  const engine = {
    context: null,
    mixer: null,
    buffers: new Map(),
    gains: new Map(),
    sources: new Map(),
    muted: new Set(),
    duration: 0,
    offset: 0,
    startedAt: 0,
    playing: false,
    loading: false,
  };

  const formatTime = (seconds) => {
    if (!Number.isFinite(seconds)) return "0:00";
    const whole = Math.max(0, Math.floor(seconds));
    return `${Math.floor(whole / 60)}:${String(whole % 60).padStart(2, "0")}`;
  };

  const speakerIcon = (muted) => `<svg class="ui-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.7" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true">${muted
    ? '<path d="M11 5 6 9H3v6h3l5 4Z"/><path d="m16 9 5 5m0-5-5 5"/>'
    : '<path d="M11 5 6 9H3v6h3l5 4Z"/><path d="M15.5 8.5a5 5 0 0 1 0 7M18 6a8.5 8.5 0 0 1 0 12"/>'
  }</svg>`;

  const setStemMuted = (stem, muted) => {
    const gain = engine.gains.get(stem);
    if (!gain || !engine.mixer) return;
    if (muted) engine.muted.add(stem); else engine.muted.delete(stem);
    const now = engine.context.currentTime;
    if (typeof gain.gain.cancelAndHoldAtTime === "function") {
      gain.gain.cancelAndHoldAtTime(now);
    } else {
      gain.gain.cancelScheduledValues(now);
    }
    gain.gain.setTargetAtTime(muted ? 0 : 1, now, 0.008);
    const control = [...engine.mixer.querySelectorAll(".mixer-toggle")]
      .find((button) => button.dataset.stem === stem);
    if (!control) return;
    const action = muted ? "Unmute" : "Mute";
    const label = control.dataset.label || stem;
    control.setAttribute("aria-pressed", String(!muted));
    control.setAttribute("aria-label", `${action} ${label}`);
    control.title = `${action} ${label}`;
    control.innerHTML = `${speakerIcon(muted)}<span>${muted ? "Muted" : "On"}</span>`;
  };

  const updateGroupControl = (group) => {
    if (!engine.mixer) return;
    const control = engine.mixer.querySelector(`.mixer-group-toggle[data-group="${group}"]`);
    const stems = [...engine.mixer.querySelectorAll(`.stem-track[data-group="${group}"]`)]
      .map((track) => track.dataset.stem);
    if (!control || !stems.length) return;
    const allMuted = stems.every((stem) => engine.muted.has(stem));
    const label = control.dataset.label || group;
    control.setAttribute("aria-pressed", String(!allMuted));
    control.setAttribute("aria-label", `${allMuted ? "Unmute" : "Mute"} all ${label}`);
    control.title = `${allMuted ? "Unmute" : "Mute"} all ${label}`;
    control.innerHTML = `${speakerIcon(allMuted)}<span>${allMuted ? "Unmute all" : "Mute all"}</span>`;
  };

  const position = () => engine.playing
    ? Math.min(engine.duration, engine.offset + engine.context.currentTime - engine.startedAt)
    : engine.offset;

  const updateTimeline = () => {
    const scrubber = document.querySelector("#timeline-scrubber");
    const current = document.querySelector("#timeline-current");
    const total = document.querySelector("#timeline-total");
    if (!scrubber || !current || !total) return;
    const now = position();
    scrubber.max = String(Math.max(1, Math.round(engine.duration * 1000)));
    scrubber.value = String(Math.round(now * 1000));
    current.textContent = formatTime(now);
    total.textContent = formatTime(engine.duration);
    const progress = engine.duration ? (now / engine.duration) * 100 : 0;
    document.querySelectorAll("#stem-track-list .stem-track").forEach((track) => {
      track.style.setProperty("--track-progress", `${progress}%`);
    });
  };

  const stopSources = () => {
    engine.sources.forEach((source) => {
      try { source.stop(); } catch (_) {}
      source.disconnect();
    });
    engine.sources.clear();
  };

  const resetEngine = () => {
    stopSources();
    engine.mixer = null;
    engine.buffers.clear();
    engine.gains.forEach((gain) => gain.disconnect());
    engine.gains.clear();
    engine.muted.clear();
    engine.duration = 0;
    engine.offset = 0;
    engine.startedAt = 0;
    engine.playing = false;
    engine.loading = false;
  };

  const initialize = async (mixer) => {
    if (engine.loading || engine.mixer === mixer) return;
    resetEngine();
    engine.loading = true;
    engine.mixer = mixer;
    engine.context ||= new AudioContext();
    const tracks = [...mixer.querySelectorAll(".stem-track[data-src]")];
    try {
      await Promise.all(tracks.map(async (track) => {
        const stem = track.dataset.stem;
        const response = await fetch(track.dataset.src);
        if (!response.ok) throw new Error(`Could not load ${stem}`);
        const buffer = await engine.context.decodeAudioData(await response.arrayBuffer());
        if (engine.mixer !== mixer) return;
        const gain = engine.context.createGain();
        gain.connect(engine.context.destination);
        engine.buffers.set(stem, buffer);
        engine.gains.set(stem, gain);
        if (track.dataset.initialMuted === "true") {
          gain.gain.value = 0;
          engine.muted.add(stem);
        }
      }));
      if (engine.mixer !== mixer) return;
      engine.duration = Math.min(...[...engine.buffers.values()].map((buffer) => buffer.duration));
      mixer.dataset.audioReady = "true";
      updateTimeline();
    } catch (error) {
      console.error("Stem mixer initialization failed", error);
      mixer.dataset.audioError = "true";
    } finally {
      engine.loading = false;
    }
  };

  const playFrom = async (offset) => {
    if (!engine.buffers.size) return;
    await engine.context.resume();
    stopSources();
    engine.offset = Math.max(0, Math.min(offset, engine.duration));
    const when = engine.context.currentTime + 0.03;
    engine.startedAt = when;
    engine.playing = true;
    engine.buffers.forEach((buffer, stem) => {
      const source = engine.context.createBufferSource();
      source.buffer = buffer;
      source.connect(engine.gains.get(stem));
      source.start(when, engine.offset);
      engine.sources.set(stem, source);
    });
  };

  const pause = () => {
    if (!engine.playing) return;
    engine.offset = position();
    engine.playing = false;
    stopSources();
    updateTimeline();
  };

  const seek = async (time) => {
    const wasPlaying = engine.playing;
    engine.offset = Math.max(0, Math.min(time, engine.duration));
    if (wasPlaying) await playFrom(engine.offset);
    updateTimeline();
  };

  document.addEventListener("click", async (event) => {
    const transport = event.target.closest(".mixer-transport");
    if (transport) {
      if (transport.dataset.action === "play") {
        await playFrom(position() >= engine.duration ? 0 : position());
      } else {
        pause();
      }
      return;
    }
    if (!engine.mixer) return;
    const groupControl = event.target.closest(".mixer-group-toggle");
    if (groupControl) {
      const group = groupControl.dataset.group;
      const stems = [...engine.mixer.querySelectorAll(`.stem-track[data-group="${group}"]`)]
        .map((track) => track.dataset.stem);
      const shouldMute = stems.some((stem) => !engine.muted.has(stem));
      if (shouldMute) {
        stems.forEach((stem) => setStemMuted(stem, true));
      } else if (groupControl.dataset.detailed === "true") {
        // Returning a detailed group to audible state selects its individual
        // sources, keeping broad aggregate predictions muted.
        engine.mixer.querySelectorAll(`.stem-track[data-group="${group}"]`).forEach((track) => {
          setStemMuted(track.dataset.stem, track.dataset.tier === "aggregate");
        });
      } else {
        stems.forEach((stem) => setStemMuted(stem, false));
      }
      updateGroupControl(group);
      return;
    }
    const control = event.target.closest(".mixer-toggle");
    if (!control) return;
    const stem = control.dataset.stem;
    const muted = !engine.muted.has(stem);
    const track = engine.mixer.querySelector(`.stem-track[data-stem="${stem}"]`);
    const group = track?.dataset.group;
    if (!muted && group && control.dataset.tier) {
      const oppositeTier = control.dataset.tier === "aggregate" ? "individual" : "aggregate";
      engine.mixer.querySelectorAll(`.stem-track[data-group="${group}"][data-tier="${oppositeTier}"]`)
        .forEach((other) => setStemMuted(other.dataset.stem, true));
    }
    setStemMuted(stem, muted);
    if (group) updateGroupControl(group);
  });

  document.addEventListener("input", async (event) => {
    if (event.target.id !== "timeline-scrubber") return;
    await seek(Number(event.target.value) / 1000);
  });

  document.addEventListener("click", (event) => {
    if (!event.target.closest("#split-button")) return;
    document.querySelector("#studio-hero")?.classList.add("is-compact");
  });

  const tick = () => {
    const mixer = document.querySelector("#sync-mixer");
    if (mixer && engine.mixer !== mixer && !engine.loading) initialize(mixer);
    if (!mixer && engine.mixer) resetEngine();
    if (engine.playing) {
      if (position() >= engine.duration) {
        engine.offset = engine.duration;
        engine.playing = false;
        stopSources();
      }
      updateTimeline();
    }
    requestAnimationFrame(tick);
  };
  requestAnimationFrame(tick);
})();
</script>
"""


MIXER_CSS = """
:root { color-scheme: dark; --studio-accent: #a78bfa; --studio-warm: #fb7185; --studio-panel: #15131d; --studio-line: #2c2838; }
body, .gradio-container { background: #0b0a10 !important; color: #f7f4ff !important; }
.gradio-container { max-width: 1180px !important; margin: 0 auto !important; padding: 28px 24px 56px !important; }
footer { display: none !important; }
#studio-hero { position: relative; overflow: hidden; margin: 8px 0 28px; padding: 38px 40px; border: 1px solid #292435; border-radius: 28px; background: radial-gradient(circle at 82% 20%, rgba(167,139,250,.26), transparent 32%), radial-gradient(circle at 10% 100%, rgba(251,113,133,.14), transparent 36%), #111018; box-shadow: 0 30px 80px rgba(0,0,0,.28); transition: margin .25s ease, padding .25s ease, border-radius .25s ease; }
#studio-hero::after { content: ""; position: absolute; width: 220px; height: 220px; right: -70px; bottom: -130px; border: 1px solid rgba(255,255,255,.12); border-radius: 50%; }
.ui-icon { width: 1.05rem; height: 1.05rem; flex: 0 0 auto; }
.hero-kicker { display: flex; align-items: center; gap: 8px; margin-bottom: 12px; color: #c4b5fd; font-size: .72rem; font-weight: 800; letter-spacing: .22em; text-transform: uppercase; }
.hero-kicker .ui-icon { width: 1rem; height: 1rem; }
#studio-hero h1 { max-width: 650px; margin: 0; font-size: clamp(2.35rem, 6vw, 4.8rem); line-height: .98; letter-spacing: -.055em; }
#studio-hero p { max-width: 560px; margin: 18px 0 0; color: #aca6b9; font-size: 1.08rem; line-height: 1.55; }
#studio-hero.is-compact { position: sticky; top: 8px; z-index: 10; margin: 8px 0 16px; padding: 13px 18px; border-radius: 16px; background: rgba(17,16,24,.94); box-shadow: 0 10px 28px rgba(0,0,0,.3); backdrop-filter: blur(14px); }
#studio-hero.is-compact::after, #studio-hero.is-compact h1, #studio-hero.is-compact p { display: none; }
#studio-hero.is-compact .hero-kicker { margin: 0; font-size: .78rem; }
#studio-hero.is-compact .hero-kicker::after { content: "Separation workspace"; margin-left: 8px; padding-left: 10px; border-left: 1px solid #494157; color: #aca6b9; font-weight: 600; letter-spacing: .08em; }
#upload-panel { align-items: stretch; gap: 16px; margin-bottom: 16px; }
#song-card, #controls-card { padding: 18px !important; border: 1px solid var(--studio-line) !important; border-radius: 20px !important; background: var(--studio-panel) !important; box-shadow: 0 14px 34px rgba(0,0,0,.16); }
#song-card > div, #controls-card > div { background: transparent !important; }
#split-button { min-height: 48px; margin-top: 8px; border: 0 !important; border-radius: 14px !important; background: linear-gradient(135deg, #8b5cf6, #ec4899) !important; box-shadow: 0 10px 28px rgba(139,92,246,.28); font-weight: 750; transition: transform .16s ease, box-shadow .16s ease; }
#split-button:hover { transform: translateY(-1px); box-shadow: 0 14px 34px rgba(139,92,246,.38); }
#status-card { min-height: 0; margin: 4px 2px 14px; color: #c9c3d4; }
#status-card h3 { margin-bottom: 2px; color: #f7f4ff; }
.tabs { border: 1px solid var(--studio-line) !important; border-radius: 22px !important; background: #111018 !important; overflow: hidden; }
.tab-nav { padding: 8px 10px 0 !important; border-bottom: 1px solid var(--studio-line) !important; background: #111018 !important; }
.tab-nav button { border-radius: 10px 10px 0 0 !important; }
.tabitem { padding: 18px !important; background: #111018 !important; }
#sync-mixer { display: flex; align-items: flex-start; gap: 16px; }
#mixer-sidebar {
  position: sticky;
  top: 1rem;
  min-width: 220px;
  border: 1px solid var(--studio-line);
  border-radius: 18px;
  padding: 18px;
  background: #181620;
}
#mixer-sidebar h3 { margin: 0 0 12px; font-size: 1.15rem; }
.transport-row { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; margin-bottom: 12px; }
.mixer-stem { display: flex; align-items: center; justify-content: space-between; gap: .75rem; padding: 8px 0; border-bottom: 1px solid rgba(255,255,255,.055); }
.mixer-stem:last-child { border-bottom: 0; }
.mixer-group { margin-top: 14px; }
.mixer-group-heading { display: flex; align-items: center; justify-content: space-between; gap: 8px; padding-bottom: 5px; border-bottom: 1px solid rgba(255,255,255,.1); }
.mixer-group-heading h4 { margin: 0; color: #c9c3d4; font-size: .76rem; letter-spacing: .08em; text-transform: uppercase; }
.mixer-group details summary { padding: 7px 0 1px; color: #817a8d; cursor: pointer; font-size: .78rem; }
.mixer-group details summary:hover { color: #c9c3d4; }
.mixer-transport, .mixer-toggle, .mixer-group-toggle { border: 1px solid #373144; border-radius: 10px; background: #211e2a; color: #f7f4ff; cursor: pointer; }
.mixer-transport { display: inline-flex; align-items: center; justify-content: center; gap: 7px; }
.mixer-transport { min-height: 40px; }
.mixer-toggle { display: inline-flex; align-items: center; justify-content: center; gap: 6px; min-width: 5.4rem; padding: 6px 10px; color: #d8b4fe; }
.mixer-group-toggle { display: inline-flex; align-items: center; justify-content: center; gap: 5px; padding: 5px 7px; color: #d8b4fe; font-size: .72rem; white-space: nowrap; }
.mixer-group-toggle .ui-icon { width: .9rem; height: .9rem; }
.mixer-toggle[aria-pressed="false"] { color: #817a8d; background: #111018; }
.mixer-group-toggle[aria-pressed="false"] { color: #817a8d; background: #111018; }
#stem-players { flex: 1; min-width: 0; }
.shared-timeline { margin-bottom: 14px; padding: 16px 18px; border: 1px solid var(--studio-line); border-radius: 18px; background: #181620; }
.timeline-labels { display: flex; justify-content: space-between; margin-bottom: .45rem; font-variant-numeric: tabular-nums; font-size: .85rem; color: var(--body-text-color-subdued); }
#timeline-scrubber { width: 100%; accent-color: var(--studio-accent); cursor: pointer; }
#stem-track-list { display: grid; gap: 1rem; }
.stem-track-group { border: 1px solid var(--studio-line); border-radius: 16px; background: #15131c; overflow: hidden; }
.stem-track-group > summary { display: flex; align-items: center; justify-content: space-between; padding: 13px 16px; color: #c9c3d4; cursor: pointer; font-size: .8rem; font-weight: 700; letter-spacing: .08em; text-transform: uppercase; }
.stem-track-group > summary span:last-child { color: #817a8d; font-size: .72rem; font-weight: 600; letter-spacing: .04em; }
.stem-track-group[open] > summary { border-bottom: 1px solid var(--studio-line); }
.stem-track-group-title { display: flex; align-items: center; justify-content: space-between; padding: 13px 16px; border-bottom: 1px solid var(--studio-line); color: #c9c3d4; font-size: .8rem; font-weight: 700; letter-spacing: .08em; text-transform: uppercase; }
.stem-track-group-title span:last-child { color: #817a8d; font-size: .72rem; font-weight: 600; letter-spacing: .04em; }
.stem-track-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(170px, 1fr)); gap: 8px; padding: 10px; }
.stem-track { border: 1px solid rgba(255,255,255,.07); border-radius: 11px; padding: 10px 11px; background: #181620; transition: border-color .16s ease, transform .16s ease; }
.stem-track:hover { border-color: #494157; transform: translateY(-1px); }
.stem-track-header { display: flex; align-items: center; justify-content: space-between; gap: .75rem; }
.stem-name, .stem-track-label { display: inline-flex; align-items: center; gap: 9px; }
.stem-name .ui-icon, .stem-track-label .ui-icon { color: #a78bfa; }
.stem-track-label { font-size: .88rem; font-weight: 600; }
.stem-download { display: inline-flex; align-items: center; gap: 6px; color: #c4b5fd !important; font-size: .82rem; text-decoration: none; }
.stem-download .ui-icon { width: .9rem; height: .9rem; }
.stem-download:hover { text-decoration: underline; }
.track-timeline { height: 7px; margin-top: 13px; border-radius: 999px; background: #292532; overflow: hidden; }
.track-playhead { display: block; width: 100%; height: 100%; transform: translateX(calc(-100% + var(--track-progress, 0%))); background: linear-gradient(90deg, var(--studio-accent), var(--studio-warm)); transition: transform .08s linear; }
.empty-mixer { padding: 60px 24px; color: #817a8d; text-align: center; }
@media (max-width: 768px) {
  .gradio-container { padding: 14px 12px 36px !important; }
  #studio-hero { padding: 28px 24px; border-radius: 22px; }
  #sync-mixer { display: block; }
  #mixer-sidebar { position: static; margin-bottom: 1rem; }
  .stem-track-grid { grid-template-columns: 1fr; }
}
"""


APP_THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.violet,
    secondary_hue=gr.themes.colors.pink,
    neutral_hue=gr.themes.colors.slate,
)


with gr.Blocks(title="Stem Studio") as demo:
    gr.HTML(
        """
        <header id="studio-hero">
          <div class="hero-kicker">
            <svg class="ui-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.7" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="m12 3-1.9 4.8L5 10l5.1 2.2L12 17l1.9-4.8L19 10l-5.1-2.2L12 3Z"/><path d="M5 3v4M3 5h4M19 17v4m-2-2h4"/></svg>
            Stem Studio
          </div>
          <h1>Hear every layer.</h1>
          <p>Upload a song, separate it, and shape the mix.</p>
        </header>
        """
    )

    with gr.Row(elem_id="upload-panel"):
        with gr.Column(scale=2, elem_id="song-card"):
            audio_input = gr.Audio(
                label="Choose a song",
                type="filepath",
                sources=["upload"],
            )
        with gr.Column(scale=1, elem_id="controls-card"):
            mode_input = gr.Radio(
                choices=MODE_CHOICES,
                value=MODE_CHOICES[2],
                label="Mix detail",
            )
            run_button = gr.Button("Split song", variant="primary", elem_id="split-button")

    force_input = gr.State(False)

    status_output = gr.Markdown("", elem_id="status-card")

    with gr.Tab("Mixer"):
        mixer_output = gr.HTML(
            "<p class='empty-mixer'>Your tracks will appear here.</p>"
        )

    with gr.Tab("Downloads"):
        files_output = gr.File(
            label="Tracks",
            file_count="multiple",
        )
        zip_output = gr.File(
            label="All tracks (.zip)",
            file_count="single",
        )

    outputs = [
        status_output,
        files_output,
        zip_output,
        mixer_output,
    ]

    audio_input.change(
        fn=lambda: None,
        inputs=[],
        outputs=[],
        js="""
        () => {
            document.querySelectorAll("audio").forEach((player) => {
                player.pause();
                player.currentTime = 0;
            });
        }
        """,
    )

    run_button.click(
        fn=separate,
        inputs=[audio_input, mode_input, force_input],
        outputs=outputs,
        concurrency_limit=1,
    )


if __name__ == "__main__":
    is_hf_space = bool(os.environ.get("SPACE_ID"))
    default_server_name = "0.0.0.0" if is_hf_space else "127.0.0.1"
    server_name = os.environ.get("GRADIO_SERVER_NAME", default_server_name)
    server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
    share = os.environ.get("GRADIO_SHARE", "").lower() in {"1", "true", "yes", "on"}

    print(f"bs-roformer-infer commit: {RUNTIME_COMMIT}")
    print(f"Opening Gradio on http://{server_name}:{server_port}")

    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        server_name=server_name,
        server_port=server_port,
        share=share,
        inbrowser=not is_hf_space,
        allowed_paths=[str(OUTPUT_ROOT.resolve())],
        theme=APP_THEME,
        css=MIXER_CSS,
        head=MIXER_HEAD,
    )