"""单一六轨模型：原生 Demucs 优先，PyTorch Demucs 兜底。"""
import os
import re
import subprocess
import sys

STEMS_6S = ("vocals", "drums", "bass", "other", "guitar", "piano")
DEMUCS_6S = "htdemucs_6s"


def separate(audio_path: str, out_dir: str, on_progress=None) -> dict:
    """对 audio_path 跑 Demucs，返回 {stem: wav_path}。

    已有产出时直接复用（重试不重复分离）。
    """
    native = os.environ.get("FENXUAN_DEMUCS_BIN", "")
    track = os.path.splitext(os.path.basename(audio_path))[0]
    stem_root = out_dir if native else os.path.join(out_dir, DEMUCS_6S, track)
    stem_paths = {s: os.path.join(stem_root, f"{s}.wav") for s in STEMS_6S}
    if all(os.path.exists(p) for p in stem_paths.values()):
        if on_progress:
            on_progress(0.99, "AI 分离（复用已有结果）")
        return stem_paths

    os.makedirs(out_dir, exist_ok=True)
    if native:
        cmd = [native, "--model", DEMUCS_6S, "--output", out_dir, audio_path]
    else:
        cmd = [sys.executable, "-m", "demucs", "-n", DEMUCS_6S,
               "--out", out_dir, "--filename", "{track}/{stem}.wav"]
        device = os.environ.get("FENXUAN_DEMUCS_DEVICE")
        if device is None:
            import torch
            if torch.cuda.is_available():
                device = "cuda"
            else:
                device = "mps" if torch.backends.mps.is_available() else ""
        if device:
            cmd.extend(["--device", device])
        cmd.append(audio_path)
    env = dict(os.environ)
    env["PYTHONUNBUFFERED"] = "1"
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, env=env,
    )
    buf = ""
    tail = []
    assert proc.stdout is not None
    for ch in iter(lambda: proc.stdout.read(1), ""):
        buf += ch
        if ch not in "\r\n":
            continue
        line = buf.strip()
        buf = ""
        if not line:
            continue
        tail.append(line)
        tail = tail[-12:]
        if native and on_progress:
            if "Loading model" in line:
                on_progress(0.12, "载入分轨模型")
            elif "Separating" in line:
                on_progress(0.2, "正在一次拆出人声和全部乐器")
        # PyTorch Demucs 的 tqdm 输出形如 "45%|...|"
        for m in re.finditer(r"(\d+(?:\.\d+)?)%", line):
            frac = min(float(m.group(1)) / 100.0, 1.0)
            if on_progress:
                on_progress(min(frac, 0.99), f"AI 分离中 {frac * 100:.0f}%")
    proc.wait()
    if proc.returncode != 0:
        detail = tail[-1] if tail else f"exit {proc.returncode}"
        raise RuntimeError(f"demucs 分离失败：{detail}")
    for stem, path in stem_paths.items():
        if not os.path.exists(path):
            raise RuntimeError(f"demucs 未产出 {stem}: {path}")
    return stem_paths
