# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Raw-audio oracle against the cache oracle and the research caches, and its committed manifest.

For every replayed file:
- the front-end record from raw audio vs the research cache entry the release was frozen from (all six arrays:
  dtype, shape, bytes);
- the RTTM vs the cache oracle's (oracle/cache_canonical.json and its artifacts);
- errors, GPU, rank and processing order.
Writes oracle/audio_canonical.json (sha256 of every RTTM and record array, errors, placement) and prints a summary.

  source env/canonical.sh && python -s tools/compare_oracles.py
"""
import hashlib
import json
import sys
from collections import Counter
from pathlib import Path

import numpy as np

PROJ = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJ / "tests"))
KEYS = ("P", "masks", "emb", "chunk_sw", "frames_sw", "duration_s")


def sha(b):
    return hashlib.sha256(b).hexdigest()


def main():
    import records
    out = PROJ / "artifacts/oracle_audio"
    log = {}
    for f in sorted(out.glob("log_rank*.jsonl")):
        for line in f.read_text().splitlines():
            r = json.loads(line)
            log[r["key"]] = r
    inputs = json.loads((PROJ / "artifacts/oracle_audio_full_inputs.json").read_text())
    cache_or = json.loads((PROJ / "oracle/cache_canonical.json").read_text())["files"]
    cached = {f"{s}/{p.stem}": p for s, p in records.all_records()}
    manifest, stats = {}, Counter()
    rec_diff, rttm_diff = [], []
    for it in inputs:
        key = f"{it['set']}/{it['sid']}"
        r = log.get(key)
        if r is None:
            stats["missing"] += 1
            continue
        m = dict(gpu=r["gpu"], rank=r["rank"], order=r["order"])
        if not r["ok"]:
            m["error"] = r["error"]
            stats["errors"] += 1
            manifest[key] = m
            continue
        d = np.load(out / it["set"] / "records" / f"{it['sid']}.npz")
        m["record_sha256"] = {k: sha(np.ascontiguousarray(d[k]).tobytes()) for k in KEYS}
        m["chunks"] = int(d["P"].shape[0])
        m["duration_s"] = float(d["duration_s"])
        if "error" in r["indic"]:
            m["indic"] = dict(error=r["indic"]["error"])
            stats["indic_errors"] += 1
        else:
            m["indic"] = dict(sha256=sha((out / it["set"] / "indic" / f"{it['sid']}.rttm").read_bytes()))
        if key in cached:
            stats["cached"] += 1
            c = np.load(cached[key])
            bad = [k for k in KEYS if c[k].dtype != d[k].dtype or c[k].shape != d[k].shape or c[k].tobytes() != d[k].tobytes()]
            if bad:
                rec_diff.append((key, bad))
            else:
                stats["records_identical"] += 1
            if m["indic"].get("sha256") == cache_or[key]["indic"].get("sha256"):
                stats["indic_rttm_identical"] += 1
            else:
                rttm_diff.append(key)
        manifest[key] = m
    res = dict(created="7 October 2026", source="tools/oracle_from_audio.py (release 1.0, canonical environment)",
               inputs="artifacts/oracle_audio_full_inputs.json", files=manifest,
               summary=dict(stats), record_differences=rec_diff, rttm_differences=rttm_diff,
               gpus=dict(Counter(m["gpu"] for m in manifest.values())))
    (PROJ / "oracle/audio_canonical.json").write_text(json.dumps(res, indent=1, sort_keys=True) + "\n")
    print(json.dumps(dict(summary=res["summary"], record_differences=rec_diff[:10], rttm_differences=rttm_diff[:10],
                          gpus=res["gpus"]), indent=1))
    errs = Counter(m["error"].split(":")[0] for m in manifest.values() if "error" in m)
    print("errors by type:", dict(errs))
    for k, m in manifest.items():
        if "error" in m:
            print(" ", k, "->", m["error"][:150])


if __name__ == "__main__":
    main()