# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Identity check of every audio input: the engine with the real networks on the local GPUs (anuvani_diar.Diarizer)
against the raw-audio oracle (the release's own code on the same audio): every front-end record and the RTTM
byte-identical, and every failure the same failure.

  make test-gpu          # golden 12 + 200-file stratified sample + long files + robustness set
  make test-gpu-all      # every replay input (1,973 files)
"""
import json
from pathlib import Path

import numpy as np
import pytest

import records
from conftest import PROJ

ORACLE = PROJ / "artifacts/oracle_audio"
INPUTS = {f"{i['set']}/{i['sid']}": i for i in json.loads((PROJ / "artifacts/oracle_audio_full_inputs.json").read_text())}
RECORD_KEYS = ("P", "masks", "emb", "chunk_sw", "frames_sw", "duration_s")


def oracle_log():
    log = {}
    for f in sorted(ORACLE.glob("log_rank*.jsonl")):
        for line in f.read_text().splitlines():
            r = json.loads(line)
            log[r["key"]] = r
    return log


def identity_sample():
    keys = [f"bench/{i['sid']}" for i in json.loads((PROJ / "artifacts/golden12_inputs.json").read_text())]
    keys += [f"{s}/{p.stem}" for s, p in records.sample(200)]
    keys += [k for k, i in INPUTS.items() if i["set"] in ("long", "robust")]
    return list(dict.fromkeys(keys))


def compare(keys, diarizer):
    log = oracle_log()
    missing = [k for k in keys if k not in log]
    assert not missing, f"oracle replay incomplete: {missing[:5]}"
    jobs = [INPUTS[k] for k in keys]
    by_uri = {}
    eng_jobs = [(i["sid"], str(Path(i["path"])), i["path"]) for i in jobs]
    for r in diarizer.engine.run(eng_jobs):
        by_uri[(jobs[r["index"]]["set"], r["uri"])] = r
    bad = []
    for k in keys:
        s, sid = k.split("/", 1)
        want, got = log[k], by_uri[(s, sid)]
        if not want["ok"]:
            if got["error"] != want["error"]:
                bad.append((k, "error", want["error"], got["error"]))
            continue
        # the release's own post-processing raised
        want_err = want["indic"].get("error")
        if got["error"] != want_err:
            bad.append((k, "error", want_err, got["error"]))
            continue
        d = np.load(ORACLE / s / "records" / f"{sid}.npz")
        rec = got["record"]
        mine = dict(P=rec["P_bits"].view(np.float16), masks=rec["masks"], emb=rec["emb"], chunk_sw=rec["chunk_sw"],
                    frames_sw=rec["frames_sw"], duration_s=np.array(rec["duration_s"]))
        for key in RECORD_KEYS:
            a, b = mine[key], d[key]
            if a.dtype != b.dtype or a.shape != b.shape or a.tobytes() != b.tobytes():
                bad.append((k, key))
        if want_err is None and (got["rttm"] is None or got["rttm"].encode() != (ORACLE / s / "indic" / f"{sid}.rttm").read_bytes()):
            bad.append((k, "rttm"))
    return bad


@pytest.fixture(scope="module")
def diarizer():
    from anuvani_diar import Diarizer
    return Diarizer(keep_records=True, strict=True)


@pytest.mark.gpu
def test_identity_audio_sample(diarizer):
    keys = identity_sample()
    assert len(keys) >= 230
    assert compare(keys, diarizer) == []


@pytest.mark.gpu
@pytest.mark.slow
def test_identity_audio_all(diarizer):
    assert compare(list(INPUTS), diarizer) == []