"""Demucs 六轨分离与 BS-Roformer 人声分离。"""
import gc
import os
import re
import subprocess
import sys

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


def _roformer_segment():
    """限制单次送入模型的音频长度，避免 Apple 图形内存被一口气吃满。"""
    try:
        value = int(os.environ.get("FENXUAN_ROFORMER_SEGMENT", "512"))
    except ValueError:
        value = 512
    return max(64, min(512, value))


def _roformer_params():
    segment = _roformer_segment()
    return {"segment_size": segment, "batch_size": 1,
            "overlap": max(1, segment // 128),
            "override_model_segment_size": True}


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

    已有产出时直接复用（重试不重复分离）。
    """
    track = os.path.splitext(os.path.basename(audio_path))[0]
    stem_paths = {s: os.path.join(out_dir, DEMUCS_6S, track, 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)
    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
        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:]
        # 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


def separate_roformer(audio_path: str, out_dir: str, on_progress=None) -> dict:
    """BS-Roformer Viperx-1297 人声/伴奏二分（当前开源人声分离最高档）。

    返回 {'vocals': path, 'instrumental': path}。
    """
    from audio_separator.separator import Separator

    os.makedirs(out_dir, exist_ok=True)
    # 已有产出直接复用
    for f in os.listdir(out_dir):
        if "(Vocals)" in f:
            inst = [g for g in os.listdir(out_dir) if "(Instrumental)" in g]
            if inst:
                if on_progress:
                    on_progress(0.95, "Roformer（复用已有结果）")
                return {"vocals": os.path.join(out_dir, f),
                        "instrumental": os.path.join(out_dir, inst[0])}
    model_dir = os.environ.get(
        "FENXUAN_MODELS",
        os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(
            os.path.dirname(os.path.abspath(__file__))))),
            "runtime", "fenxuan", "sep-models"))
    if on_progress:
        on_progress(0.05, "准备人声模型（首次使用会自动下载）")
    # ponytail: 小段降低峰值内存；步长必须随片段一起缩短，否则中间音频会漏掉。
    sep = Separator(
        log_level=50, model_file_dir=model_dir,
        output_dir=out_dir, output_format="WAV",
        mdxc_params=_roformer_params())
    try:
        sep.load_model("model_bs_roformer_ep_317_sdr_12.9755.ckpt")
        if on_progress:
            on_progress(0.15, "正在逐段分离人声与伴奏，长歌曲需要几分钟")
        outputs = sep.separate(audio_path)
    finally:
        del sep
        gc.collect()
        try:
            import torch
            if torch.backends.mps.is_available():
                torch.mps.empty_cache()
        except (ImportError, RuntimeError):
            pass
    if on_progress:
        on_progress(0.95, "Roformer 分离完成")
    vocals = next((os.path.join(out_dir, f) for f in outputs if "(Vocals)" in f), None)
    inst = next((os.path.join(out_dir, f) for f in outputs if "(Instrumental)" in f), None)
    if not vocals or not inst or not os.path.exists(vocals) or not os.path.exists(inst):
        raise RuntimeError(f"Roformer 未产出人声/伴奏: {outputs}")
    return {"vocals": vocals, "instrumental": inst}
