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

The registry drives the engine1 → engine2 path: demucs (htdemucs, what PLN ran by
hand), Roformer (BS-/Mel-Band via `audio-separator`), and the **merge** backend —
PLN's "merge-of-stems multimodel": SOTA Roformer vocal isolation cascaded into
demucs on the instrumental for drums/bass/other, best-of-both in the 4-stem contract.

Backends shell out to their OWN venv binary, so this module runs cleanly under
system python3. A backend declares what it `produces`; composite backends override
`run_composite` to orchestrate several sub-backends.
"""
from __future__ import annotations

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

from .model import STEMS, Catch, Stem

# demucs / audio-separator print a tqdm-ish bar to stderr; grab the integer percent.
_PCT = re.compile(r"(\d+)%\|")


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

    name = "base"
    default_model = ""
    models: list[str] = []          # selectable model variants (for the GUI/CLI)
    produces: tuple = STEMS         # the stem names this backend emits
    composite = False               # True ⇒ orchestrates sub-backends via run_composite

    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 files."""
        raise NotImplementedError

    def normalize(self, src_stems: Path, stems_dir: Path, catch_dir: Path) -> list[Stem]:
        """Move the backend's raw outputs → stems_dir/{name}.wav. Default = by name."""
        found: list[Stem] = []
        for name in self.produces:
            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))))
        return found

    def run_composite(self, catch, workspace, source, raw_out, stems_dir,
                      progress, log, **opts) -> list[Stem]:
        raise NotImplementedError


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

    name = "demucs"
    default_model = "htdemucs"
    models = ["htdemucs", "htdemucs_ft", "hdemucs_mmi", "htdemucs_6s"]

    @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: BS-/Mel-Band Roformer via `audio-separator` (py3.12 venv).

    A 2-STEM vocal/instrumental SOTA (spike #8: fits 6 GB, ~0.55× realtime). Not a
    4-stem replacement for demucs — see MultiModelBackend for the 4-stem merge.
    """

    name = "roformer"
    default_model = "model_bs_roformer_ep_317_sdr_12.9755.ckpt"
    models = ["model_bs_roformer_ep_317_sdr_12.9755.ckpt",
              "melband_roformer_inst_v2.ckpt",
              "model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt"]
    produces = ("vocals", "instrumental")

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

    def available(self) -> tuple[bool, str]:
        if Path(self.bin).exists():
            return True, self.bin
        return False, (f"audio-separator not found at {self.bin} — needs the py3.12 "
                       "venv (py3.14 breaks beartype). See README 'engine2'.")

    def build_cmd(self, source: Path, outdir: Path, **opts) -> list[str]:
        model = opts.get("model") or self.default_model
        outdir.mkdir(parents=True, exist_ok=True)
        return [self.bin, str(source), "--model_filename", model,
                "--output_dir", str(outdir), "--output_format", "WAV",
                "--log_level", "WARNING"]

    def stem_dir(self, source: Path, outdir: Path, model: str) -> Path:
        return Path(outdir)   # audio-separator writes straight into --output_dir

    def normalize(self, src_stems: Path, stems_dir: Path, catch_dir: Path) -> list[Stem]:
        # audio-separator names files "<src>_(Vocals)_<model…>.wav" — glob by tag.
        found: list[Stem] = []
        for tag, name in (("Vocals", "vocals"), ("Instrumental", "instrumental")):
            hits = glob.glob(str(src_stems / f"*({tag})*.wav"))
            if hits:
                dest = stems_dir / f"{name}.wav"
                shutil.move(hits[0], str(dest))
                found.append(Stem(name=name, path=str(dest.relative_to(catch_dir))))
        return found


class MultiModelBackend(Backend):
    """engine2 flagship — PLN's "merge of stems": Roformer vocals + demucs rest.

    Cascade: Roformer isolates vocals + instrumental (SOTA vocal SDR), then demucs
    runs on the INSTRUMENTAL (vocals already gone → cleaner) for drums/bass/other.
    Merge = Roformer vocals ⊕ demucs drums/bass/other in the standard 4-stem contract.
    """

    name = "merge"
    default_model = "roformer-vocals+demucs-rest"
    composite = True

    def available(self) -> tuple[bool, str]:
        for sub in ("roformer", "demucs"):
            ok, hint = REGISTRY[sub].available()
            if not ok:
                return False, f"merge needs {sub}: {hint}"
        return True, "roformer (vocals) + demucs (drums/bass/other)"

    def run_composite(self, catch, workspace, source, raw_out, stems_dir,
                      progress, log, **opts) -> list[Stem]:
        rof, dem = REGISTRY["roformer"], REGISTRY["demucs"]

        # 1. Roformer → vocals + instrumental
        log("⚙  merge 1/2: roformer (vocals)")
        rof_out = raw_out / "roformer"
        _run_cmd(rof.build_cmd(source, rof_out, model=rof.default_model), progress, log)
        rof_stems = rof.normalize(rof.stem_dir(source, rof_out, rof.default_model),
                                  stems_dir, catch.dir(workspace))
        inst = next((stems_dir / "instrumental.wav" for s in rof_stems
                     if s.name == "instrumental"), None)
        vocals = [s for s in rof_stems if s.name == "vocals"]
        if not inst or not inst.exists():
            raise RuntimeError("merge: roformer produced no instrumental to cascade")

        # 2. demucs on the instrumental → drums/bass/other (discard its vocals)
        log("⚙  merge 2/2: demucs on the instrumental (drums/bass/other)")
        dmodel = opts.get("model") or "htdemucs"
        if dmodel not in dem.models:
            dmodel = "htdemucs"
        dem_out = raw_out / "demucs"
        _run_cmd(dem.build_cmd(inst, dem_out, model=dmodel,
                               shifts=opts.get("shifts", 2)), progress, log)
        dsrc = dem.stem_dir(inst, dem_out, dmodel)
        found: list[Stem] = list(vocals)
        for name in ("drums", "bass", "other"):
            wav = dsrc / 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(workspace)))))
        inst.unlink(missing_ok=True)   # the instrumental was a means, not a stem
        return found


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


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 _run_cmd(cmd: list[str], progress, log) -> None:
    """Run a separation subprocess, forwarding best-effort tqdm percent to progress."""
    log(f"  $ {' '.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 and progress:
            pct = int(m.group(1))
            if pct != last:
                progress(pct)
                last = pct
    if proc.wait() != 0:
        raise RuntimeError(f"{cmd[0]} exited {proc.returncode}")


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"
    stems_dir = catch_dir / "stems"
    stems_dir.mkdir(exist_ok=True)

    if be.composite:
        log(f"⚙  {backend}: {model}")
        found = be.run_composite(catch, workspace, source, raw_out, stems_dir,
                                 progress, log, model=model, **opts)
    else:
        cmd = be.build_cmd(source, raw_out, model=model, **opts)
        log(f"⚙  {backend}/{model}: {' '.join(cmd)}")
        _run_cmd(cmd, progress, log)
        found = be.normalize(be.stem_dir(source, raw_out, model), stems_dir, catch_dir)

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