#!/usr/bin/env python
# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Install check: diarise the 12 IndicDiarBench check recordings and compare with the published output fingerprints.

  python tools/verify_golden.py [--audio-dir AUDIO_DIR] [--device cuda:0]

AUDIO_DIR defaults to the repository's golden/audio (one level above this program's folder).

PASS on all 12 outputs means your installation reproduces our outputs byte for byte (and therefore our scores).
A FAIL with a matching audio fingerprint usually means a different library version, GPU type or non-default
--batch-size / --threads; report the run.json written next to the outputs.
"""
import argparse
import hashlib
import json
import subprocess
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parents[1]


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--audio-dir", type=Path, default=HERE.parent / "golden" / "audio")
    p.add_argument("--device", default="cuda:0")
    a = p.parse_args()
    exp = json.loads((HERE / "golden/expected.json").read_text())["files"]
    found = {f.stem: f for f in a.audio_dir.rglob("*.flac") if f.stem in exp}
    missing = sorted(set(exp) - set(found))
    if missing:
        raise SystemExit(f"audio missing for {missing}")
    bad_audio = [s for s, f in found.items() if hashlib.sha256(f.read_bytes()).hexdigest() != exp[s]["flac_sha256"]]
    if bad_audio:
        print(f"WARNING: audio differs from the scored audio for {bad_audio}")
    ok = 0
    with tempfile.TemporaryDirectory() as t:
        inp = Path(t) / "in"; inp.mkdir()
        for s, f in found.items():
            (inp / f"{s}.flac").symlink_to(f.resolve())
        out = Path(t) / "out"
        subprocess.run([sys.executable, str(HERE / "diarize.py"), "--device", a.device, "--in-dir", str(inp),
                        "--out", str(out)], check=True, stdout=subprocess.DEVNULL)
        for s in sorted(exp):
            got = hashlib.sha256((out / f"{s}.rttm").read_bytes()).hexdigest()
            good = got == exp[s]["indic_rttm_sha256"]
            ok += good
            print(f"{s:14s} {'PASS' if good else 'FAIL'}")
    print(f"{ok}/{len(exp)} outputs identical")
    sys.exit(0 if ok == len(exp) else 1)


if __name__ == "__main__":
    main()