#!/usr/bin/env python3
# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Fast checks of the wheel, CPU only (CUDA_VISIBLE_DEVICES is set empty for every step); about a minute.

  python3 tools/wheel_check.py [--wheel WHEEL] [--data-from DIR] [--serial] [--keep-venv]

(a) This tree, in the project venv (.venv) and the canonical environment, as `make` runs the tests: the existing tests
    that cover what the wheel bundles. OpenBLAS: the row-parallel VBx products against numpy's own (shapes up to
    5,000 rows; the 20,774-row ones take 16 s and are left to `make test-fast`) and the engine's bits with every BLAS
    thread variable at 64. libsndfile: decoding against soundfile over every format, 2,750 damaged files, and the
    release's load_audio. The assets and everything after the networks: 40 cached records' RTTMs against the
    canonical oracle, through diarize_record and through the engine.
(b) The wheel in a fresh venv: Python 3.12, torch 2.11 (CPU wheels), numpy 2.x and the wheel; none of
    env/canonical.sh's variables and no ANUVANI_DIAR_* variable set, so the package alone finds its libraries.
    torch and numpy are imported first (the Diarizer loads torch before the engine), then anuvani_diar. The private
    OpenBLAS must run 1 thread with core Zen, from the package, without starting a thread; two golden FLACs must
    decode to the reference environment's samples (sha256 of the float32 data, computed there with soundfile); the
    1 h, 3 h and 6 h files' raw-audio oracle records must give RTTMs whose sha256 equal assets/long_golden.json.

The data (cached records, oracle records, test audio) are read from this tree or, with --data-from, from another
checkout of this project; in a git worktree, which has none of the untracked data, the main checkout is used by
default. Part (a) then runs in a private mount namespace (unshare) where the other checkout's folders are bound
read-only over this tree's, so the tests run unchanged and nothing is written to either tree.

(a) and (b) run at the same time unless --serial. The result is printed and written next to the wheel
(<wheel>.check.json). Exit status 0 only if every check passed.
"""
import argparse
import hashlib
import json
import os
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

PROJ = Path(__file__).resolve().parents[1]
# a Python of the release's environment (soundfile 0.14.0): ANUVANI_DIAR_REFERENCE_PYTHON, else this one
REFERENCE_PYTHON = Path(os.environ.get("ANUVANI_DIAR_REFERENCE_PYTHON") or sys.executable)
TORCH = "torch==2.11.*"
TORCH_INDEX = "https://download.pytorch.org/whl/cpu"
NUMPY = "numpy>=2,<3"
LONG = ("long_1h", "long_3h", "long_6h")
FLACS = ("bengali_009", "telugu_020")
TESTS = [
    "tests/test_audio_decode.py::test_synthetic_matrix_is_broad",
    "tests/test_audio_decode.py::test_sf_read_matches_soundfile",
    "tests/test_audio_decode.py::test_load_audio_matches_release",
    "tests/test_audio_decode.py::test_damaged_files_decode_as_soundfile",
    "tests/test_audio_decode.py::test_damaged_files_load_as_the_release",
    "tests/test_audio_decode.py::test_memory_error_text_matches_numpy",
    "tests/test_npblas_par.py",
    "tests/test_engine.py::test_engine_bits_do_not_depend_on_openblas_env",
    "tests/test_engine.py::test_engine_records_sample",
    "tests/test_engine.py::test_a_low_sample_rate_cannot_resample_past_the_limit",
    "tests/test_identity_cached.py::test_identity_cached_sample",
]
TESTS_K = "not 20774"

# run inside the fresh venv with `python -I`: argv[1] is a JSON spec
VENV_CHECK = r'''
import hashlib, json, os, sys, sysconfig, time
spec = json.loads(sys.argv[1])
out = dict(problems=[])
def bad(msg):
    out["problems"].append(msg)
def tasks():
    return len(os.listdir("/proc/self/task"))
set_vars = sorted(v for v in spec["must_be_unset"] if v in os.environ)
if set_vars:
    bad(f"variables set: {set_vars}")
import numpy as np
import torch
out["python"] = sys.version.split()[0]
out["numpy"] = np.__version__
out["torch"] = torch.__version__
before = tasks()
t = time.perf_counter()
import anuvani_diar
from anuvani_diar import _core
blas = _core.openblas_info()
out["openblas"] = dict(blas, load_s=round(time.perf_counter() - t, 3), threads_started=tasks() - before)
site = sysconfig.get_paths()["platlib"]
out["package_dir"] = os.path.dirname(anuvani_diar.__file__)
if not out["package_dir"].startswith(site):
    bad(f"anuvani_diar imported from {out['package_dir']}, not the venv")
if not blas["path"].startswith(out["package_dir"]):
    bad(f"OpenBLAS loaded from {blas['path']}")
if blas["num_threads"] != 1 or blas["corename"] != "Zen" or out["openblas"]["threads_started"] != 0:
    bad(f"OpenBLAS: {blas['num_threads']} threads, core {blas['corename']}, {out['openblas']['threads_started']} threads started")
snd = _core.sndfile_info()
out["sndfile"] = snd
if not (snd["checked"] and snd["version"] == "libsndfile-1.2.2" and snd["path"].startswith(out["package_dir"])):
    bad(f"libsndfile: {snd}")
out["decoded"] = {}
for sid, f in spec["flacs"].items():
    data, sr = _core.sf_read_f32(f["path"])
    if data is None:
        bad(f"{sid}: {sr}")
        continue
    got = dict(sha256=hashlib.sha256(np.ascontiguousarray(data).tobytes()).hexdigest(), samples=int(data.shape[0]), samplerate=int(sr))
    out["decoded"][sid] = dict(got, identical=got == f["reference"])
    if got != f["reference"]:
        bad(f"{sid}: decoded {got}, reference {f['reference']}")
out["long"] = {}
for sid, r in spec["records"].items():
    d = dict(np.load(r["npz"]))
    t = time.perf_counter()
    text = _core.diarize_record(sid, np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"],
                                list(map(float, d["chunk_sw"])), list(map(float, d["frames_sw"])), float(d["duration_s"]),
                                spec["clustering"], spec["decoder"])
    got = hashlib.sha256(text.encode()).hexdigest()
    out["long"][sid] = dict(chunks=int(d["P"].shape[0]), wall_s=round(time.perf_counter() - t, 2), rttm_sha256=got,
                            identical=got == r["rttm"])
    if got != r["rttm"]:
        bad(f"{sid}: RTTM sha256 {got}, golden {r['rttm']}")
after = _core.openblas_info()
if after["num_threads"] != 1:
    bad(f"OpenBLAS runs {after['num_threads']} threads after the records")
print(json.dumps(out))
'''

REFERENCE_DECODE = r'''
import hashlib, json, sys
import numpy as np, soundfile as sf
out = {}
for sid, path in json.loads(sys.argv[1]).items():
    data, sr = sf.read(path, dtype="float32", always_2d=False)
    out[sid] = dict(sha256=hashlib.sha256(np.ascontiguousarray(data).tobytes()).hexdigest(), samples=int(data.shape[0]),
                    samplerate=int(sr))
print(json.dumps(dict(out, soundfile=sf.__version__, libsndfile=sf.__libsndfile_version__)))
'''


def git(*args):
    return subprocess.run(["git", "-C", str(PROJ), *args], capture_output=True, text=True, check=True).stdout.strip()


def main_checkout():
    """This project in the main checkout when this tree is a git worktree, else None."""
    try:
        common = Path(git("rev-parse", "--path-format=absolute", "--git-common-dir"))
        top = Path(git("rev-parse", "--show-toplevel"))
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None
    if common.parent == top:
        return None
    return common.parent / PROJ.relative_to(top)


def canonical_vars():
    return re.findall(r"^export (\w+)=", (PROJ / "env/canonical.sh").read_text(), re.M)


def run_tests(data, log):
    """Part (a); returns (ok, summary line, seconds)."""
    venv = PROJ / ".venv"
    cmd = (f"cd {shlex.quote(str(PROJ))} && source env/canonical.sh && unset CONDA_PREFIX CONDA_DEFAULT_ENV && "
           f"export VIRTUAL_ENV={shlex.quote(str(venv))} PATH={shlex.quote(str(venv / 'bin'))}:$PATH && "
           f"python -s -m pytest -q {' '.join(map(shlex.quote, TESTS))} -k {shlex.quote(TESTS_K)} -p no:cacheprovider")
    if data != PROJ:
        binds = [(data.parents[1] / "diar", PROJ.parents[1] / "diar"),
                 (data.parent / "anuvani_diar_v1", PROJ.parent / "anuvani_diar_v1"),
                 (data / "artifacts", PROJ / "artifacts")]
        mounts = " && ".join(f"mount --bind {shlex.quote(str(s))} {shlex.quote(str(d))} && "
                             f"mount -o remount,bind,ro {shlex.quote(str(d))}" for s, d in binds)
        argv = ["unshare", "--user", "--map-root-user", "--mount", "bash", "-c", f"set -e; {mounts}; {cmd}"]
    else:
        argv = ["bash", "-c", cmd]
    t = time.perf_counter()
    with open(log, "w") as f:
        p = subprocess.Popen(argv, stdout=f, stderr=subprocess.STDOUT)
    return p, t


def finish_tests(p, t, log):
    p.wait()
    secs = time.perf_counter() - t
    text = Path(log).read_text()
    lines = [ln for ln in text.splitlines() if re.search(r"\d+ (passed|failed|error)", ln)]
    summary = lines[-1].strip("= ") if lines else text[-2000:]
    return dict(ok=p.returncode == 0 and "passed" in summary and "failed" not in summary, summary=summary,
                seconds=round(secs, 1), log=str(log))


def check_wheel(wheel, data, tmp, keep):
    """Part (b)."""
    t0 = time.perf_counter()
    res = dict(problems=[])
    # the decoder and the clustering setting (the networks are not run): $ANUVANI_DIAR_BUNDLE, the project's
    # artifacts/bundle, or the release's bundle (the same decoder file)
    bundle = Path(os.environ.get("ANUVANI_DIAR_BUNDLE") or next(
        b for b in (data / "artifacts/bundle", data.parent / "anuvani_diar_v1/bundle") if (b / "MANIFEST.json").is_file()))
    res["bundle"] = str(bundle)
    model = json.loads((bundle / "MANIFEST.json").read_text())["modes"]["indic"]
    golden_long = json.loads((PROJ / "assets/long_golden.json").read_text())["files"]
    inputs = {i["sid"]: i for i in json.loads((data / "artifacts/golden12_inputs.json").read_text())}
    flacs = {}
    for sid in FLACS:
        path = Path(inputs[sid]["path"])
        if hashlib.sha256(path.read_bytes()).hexdigest() != inputs[sid]["sha256"]:
            res["problems"].append(f"{path}: the file is not the golden input")
        flacs[sid] = str(path)
    ref = subprocess.run([str(REFERENCE_PYTHON), "-I", "-c", REFERENCE_DECODE, json.dumps(flacs)], capture_output=True,
                         text=True, check=True, env=dict(os.environ, OPENBLAS_NUM_THREADS="1"))
    reference = json.loads(ref.stdout)
    res["reference_decoder"] = dict(soundfile=reference.pop("soundfile"), libsndfile=reference.pop("libsndfile"),
                                    python=str(REFERENCE_PYTHON))
    venv = tmp / "venv"
    t = time.perf_counter()
    subprocess.run(["uv", "venv", "-q", "--python", "3.12", str(venv)], check=True)
    py = str(venv / "bin" / "python")
    subprocess.run(["uv", "pip", "install", "-q", "--python", py, "--index-url", TORCH_INDEX, TORCH], check=True)
    subprocess.run(["uv", "pip", "install", "-q", "--python", py, NUMPY], check=True)
    subprocess.run(["uv", "pip", "install", "-q", "--python", py, str(wheel)], check=True)
    res["venv_s"] = round(time.perf_counter() - t, 1)
    spec = dict(
        must_be_unset=sorted(set(canonical_vars()) | {"ANUVANI_DIAR_OPENBLAS_DIR", "ANUVANI_DIAR_SNDFILE",
                                                       "ANUVANI_DIAR_ASSETS", "ANUVANI_DIAR_SNDFILE_UNCHECKED"}),
        flacs={sid: dict(path=flacs[sid], reference=reference[sid]) for sid in FLACS},
        records={sid: dict(npz=str(data / f"artifacts/oracle_audio/long/records/{sid}.npz"), rttm=golden_long[sid]["rttm"])
                 for sid in LONG},
        clustering=model["clustering"], decoder=str(bundle / model["decoder"]))
    env = dict(PATH="/usr/bin:/bin", HOME=os.environ.get("HOME", "/tmp"), CUDA_VISIBLE_DEVICES="", LANG="C.UTF-8")
    t = time.perf_counter()
    r = subprocess.run([py, "-I", "-c", VENV_CHECK, json.dumps(spec)], capture_output=True, text=True, env=env, cwd=tmp)
    res["run_s"] = round(time.perf_counter() - t, 1)
    if r.returncode != 0:
        res["problems"].append(f"the check in the venv failed: {r.stderr[-3000:]}")
    else:
        out = json.loads(r.stdout.strip().splitlines()[-1])
        res["problems"] += out.pop("problems")
        res.update(out)
        res["stderr"] = r.stderr[-2000:]
    if keep:
        res["venv"] = str(venv)
    res["seconds"] = round(time.perf_counter() - t0, 1)
    res["ok"] = not res["problems"]
    return res


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--wheel", type=Path, help="default: the newest artifacts/wheels/*.whl")
    ap.add_argument("--data-from", type=Path, help="another checkout of this project (anuvani_diar_rs) holding the data")
    ap.add_argument("--serial", action="store_true", help="run (a), then (b)")
    ap.add_argument("--keep-venv", action="store_true")
    a = ap.parse_args(argv)
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    wheel = a.wheel or max((PROJ / "artifacts/wheels").glob("*.whl"), key=lambda p: p.stat().st_mtime)
    data = (a.data_from or (None if (PROJ / "artifacts/dev_data").is_dir() else main_checkout())
            or PROJ).resolve()
    print(f"wheel {wheel.name}\ndata  {data}" + ("" if data == PROJ else " (bound read-only for the tests)"), flush=True)
    t0 = time.perf_counter()
    tmp = Path(tempfile.mkdtemp(prefix="anuvani_wheel_check_"))
    log = tmp / "tests.log"
    try:
        (PROJ / "artifacts").mkdir(exist_ok=True)
        p, t = run_tests(data, log)
        if a.serial:
            tests = finish_tests(p, t, log)
            venv = check_wheel(wheel, data, tmp, a.keep_venv)
        else:
            venv = check_wheel(wheel, data, tmp, a.keep_venv)
            tests = finish_tests(p, t, log)
        if not tests["ok"]:
            print(Path(log).read_text()[-4000:])
    finally:
        if not a.keep_venv:
            shutil.rmtree(tmp / "venv", ignore_errors=True)
    total = round(time.perf_counter() - t0, 1)
    result = dict(wheel=wheel.name, wheel_sha256=hashlib.sha256(wheel.read_bytes()).hexdigest(), data=str(data),
                  commit=git("rev-parse", "HEAD"), host=os.uname().nodename, cpus=os.cpu_count(),
                  mode="serial" if a.serial else "parallel", seconds=total, tests=tests, wheel_in_fresh_venv=venv,
                  ok=tests["ok"] and venv["ok"])
    (wheel.parent / f"{wheel.name}.check.json").write_text(json.dumps(result, indent=1) + "\n")
    print(f"(a) tests in .venv: {'PASS' if tests['ok'] else 'FAIL'}  {tests['summary']}  ({tests['seconds']} s)")
    print(f"(b) wheel in a fresh venv: {'PASS' if venv['ok'] else 'FAIL'}  (venv {venv.get('venv_s')} s, checks "
          f"{venv.get('run_s')} s; Python {venv.get('python')}, torch {venv.get('torch')}, numpy {venv.get('numpy')})")
    if "openblas" in venv:
        b = venv["openblas"]
        print(f"    OpenBLAS {b['num_threads']} thread, core {b['corename']}, forced {b['coretype_forced']}, "
              f"{b['threads_started']} threads started, from {b['path']}")
    for sid, d in venv.get("decoded", {}).items():
        print(f"    {sid}: {d['samples']} samples at {d['samplerate']} Hz, sha256 {d['sha256'][:16]}..., "
              f"{'identical to' if d['identical'] else 'DIFFERENT from'} the reference environment's")
    for sid, d in venv.get("long", {}).items():
        print(f"    {sid}: RTTM sha256 {d['rttm_sha256'][:16]}... {'identical' if d['identical'] else 'DIFFERENT'} "
              f"({d['chunks']} chunks, {d['wall_s']} s)")
    for msg in venv["problems"]:
        print(f"    PROBLEM: {msg}")
    print(f"{'PASS' if result['ok'] else 'FAIL'} in {total} s; record: {wheel.parent / (wheel.name + '.check.json')}")
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())