"""separate — split a source into stems via a pluggable backend.

The whole point of the registry is the engine1 → engine2 path: today demucs
(htdemucs, the model PLN already runs by hand), tomorrow a Roformer SOTA model
(BS-/Mel-Band Roformer via `audio-separator`) behind the *same* interface. The
CLI/GUI ask for "a backend"; they never hardcode demucs.

Backends shell out to their OWN venv binary, so this module — and everything
that imports it — runs cleanly under system python3.
"""
from __future__ import annotations

import os
import re
import shutil
import subprocess
from pathlib import Path
from typing import Callable, Optional

from .model import STEMS, Catch, Stem

# demucs prints a tqdm bar to stderr; grab the integer percent best-effort.
_PCT = re.compile(r"(\d+)%\|")


class Backend:
    """A separation backend. Subclasses fill in the three hooks."""

    name = "base"
    default_model = ""

    def available(self) -> tuple[bool, str]:
        """(is_runnable, human hint). Override."""
        return False, "not implemented"

    def build_cmd(self, source: Path, outdir: Path, **opts) -> list[str]:
        raise NotImplementedError

    def stem_dir(self, source: Path, outdir: Path, model: str) -> Path:
        """Where this backend writes its {stem}.wav files."""
        raise NotImplementedError


class DemucsBackend(Backend):
    """engine1: Meta's Demucs (htdemucs by default). Fits our 6 GB RTX 2060."""

    name = "demucs"
    default_model = "htdemucs"

    @property
    def bin(self) -> str:
        return os.environ.get(
            "FOUNDRY_DEMUCS",
            str(Path.home() / ".virtualenvs" / "demucs" / "bin" / "demucs"),
        )

    def available(self) -> tuple[bool, str]:
        if Path(self.bin).exists() or shutil.which(self.bin):
            return True, self.bin
        return False, f"demucs binary not found at {self.bin} (set $FOUNDRY_DEMUCS)"

    def build_cmd(self, source: Path, outdir: Path, **opts) -> list[str]:
        model = opts.get("model") or self.default_model
        shifts = int(opts.get("shifts", 4))
        cmd = [self.bin, "-n", model, "--shifts", str(shifts), "-o", str(outdir)]
        if opts.get("jobs"):
            cmd += ["-j", str(int(opts["jobs"]))]
        if opts.get("device"):
            cmd += ["-d", str(opts["device"])]
        cmd.append(str(source))
        return cmd

    def stem_dir(self, source: Path, outdir: Path, model: str) -> Path:
        # demucs writes outdir/<model>/<source-stem-name>/{stem}.wav
        return Path(outdir) / model / source.stem


class RoformerBackend(Backend):
    """engine2 (stub): BS-/Mel-Band Roformer via `audio-separator`.

    Wiring deferred — see tools/foundry/README.md "engine2". Registered now so
    the registry, CLI `--backend`, and GUI selector already know it exists.
    """

    name = "roformer"
    default_model = "model_bs_roformer_ep_317_sdr_12.9755"

    def available(self) -> tuple[bool, str]:
        return False, (
            "engine2 not wired yet — needs an `audio-separator` venv + a VRAM "
            "spike on the 6 GB card. See README 'engine2'."
        )


REGISTRY: dict[str, Backend] = {b.name: b for b in (DemucsBackend(), RoformerBackend())}


def get_backend(name: str) -> Backend:
    if name not in REGISTRY:
        raise KeyError(f"unknown backend {name!r}; have {sorted(REGISTRY)}")
    return REGISTRY[name]


def separate(
    catch: Catch,
    workspace: Path,
    *,
    backend: str = "demucs",
    progress: Optional[Callable[[int], None]] = None,
    log: Callable[[str], None] = print,
    **opts,
) -> Catch:
    """Separate a (downloaded) Catch in place; updates and re-saves it."""
    be = get_backend(backend)
    ok, hint = be.available()
    if not ok:
        raise RuntimeError(f"backend {backend!r} unavailable: {hint}")

    catch_dir = catch.dir(workspace)
    source = catch_dir / catch.source_path
    if not source.exists():
        raise FileNotFoundError(source)

    model = opts.pop("model", None) or be.default_model
    raw_out = catch_dir / "_separated"
    cmd = be.build_cmd(source, raw_out, model=model, **opts)
    log(f"⚙  {backend}/{model}: {' '.join(cmd)}")

    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True, bufsize=1)
    last = -1
    for line in proc.stderr or []:
        m = _PCT.search(line)
        if m:
            pct = int(m.group(1))
            if pct != last and progress:
                progress(pct)
                last = pct
    if proc.wait() != 0:
        raise RuntimeError(f"{backend} exited {proc.returncode}")

    # Normalize backend output → catch_dir/stems/{name}.wav
    src_stems = be.stem_dir(source, raw_out, model)
    stems_dir = catch_dir / "stems"
    stems_dir.mkdir(exist_ok=True)
    found: list[Stem] = []
    for name in STEMS:
        wav = src_stems / f"{name}.wav"
        if wav.exists():
            dest = stems_dir / f"{name}.wav"
            shutil.move(str(wav), str(dest))
            found.append(Stem(name=name, path=str(dest.relative_to(catch_dir))))
    shutil.rmtree(raw_out, ignore_errors=True)

    if not found:
        raise RuntimeError(f"no stems produced under {src_stems}")
    catch.backend, catch.model, catch.stems = backend, model, found
    catch.save(workspace)
    log(f"✓ {len(found)} stems → {stems_dir}")
    return catch
