"""Run a pipeline step in its virtual environment (stdlib only, any Python 3.11+).

usage: python run.py <step> [args]      python run.py all      python run.py check
Two environments, because qwen-tts pins transformers 4.x and the QC models need 5.x:
  .venv-tts : extract, segments, render, refbench, mod
  .venv-qc  : refs, qc, pick, assemble
"""
import os
import pathlib
import subprocess
import sys

ROOT = pathlib.Path(__file__).resolve().parent
ENV = {"extract": "tts", "segments": "tts", "render": "tts", "refbench": "tts", "mod": "tts",
       "refs": "qc", "qc": "qc", "pick": "qc", "assemble": "qc"}
ALL = [["extract"], ["refs"], ["segments"], ["render"], ["qc"], ["pick"],
       ["render", "--weak"], ["qc"], ["pick"], ["assemble"], ["mod"]]


def py(env):
    d = ROOT / f".venv-{env}"
    exe = d / ("Scripts/python.exe" if os.name == "nt" else "bin/python")
    if not exe.exists():
        raise SystemExit(f"{d} is missing: see README, 'Установка'")
    return str(exe)


def run(step, args):
    if step == "refbench" and "--report" in args:
        env = "qc"
    else:
        env = ENV[step]
    e = dict(os.environ, PYTHONUTF8="1", PYTHONIOENCODING="utf-8")
    print(f"== {step} {' '.join(args)} [{env}]", flush=True)
    r = subprocess.run([py(env), "-m", f"ensvoice.{step}", *args], cwd=ROOT, env=e)
    if r.returncode:
        raise SystemExit(f"step {step} failed ({r.returncode})")


def check():
    code = ("import torch,sys;print(sys.version.split()[0],'torch',torch.__version__,'cuda',torch.cuda.is_available(),"
            "torch.cuda.get_device_name(0) if torch.cuda.is_available() else '-')")
    for env, mods in (("tts", "qwen_tts,ruaccent,soundfile"), ("qc", "gigaam,utmosv2,transformers,librosa,soxr,jiwer")):
        subprocess.run([py(env), "-c", f"import {mods};" + code], check=True)
    from ensvoice.config import load  # noqa: E402  (stdlib + tomllib only)
    cfg = load()
    for f in ("enshrouded.kfc",):
        assert (cfg.game / f).exists(), f"game not found: {cfg.game / f}"
    assert (cfg.tools / "kfc-parser.exe").exists(), f"missing {cfg.tools / 'kfc-parser.exe'}: python install.py downloads it"
    print("CHECK_OK")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    step, args = sys.argv[1], sys.argv[2:]
    if step == "check":
        check()
    elif step == "all":
        for s in ALL:
            run(s[0], s[1:])
    else:
        run(step, args)