# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""The engine (crate diar-engine) against the release, byte for byte, without a GPU:

1. cached front-end records through the engine's post-processing pool vs the canonical cache oracle;
2. audio -> RTTM with the stand-in networks of fake_nn.py in both the engine and the release's own `analyse` and
   `indic_mode` (run in a subprocess): records, the RTTM and every error must be identical, with and without the
   repacked voiceprint trunk, and every network call must be one of the release's batch shapes.
"""
import json
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest

import fake_nn
import records
from conftest import PROJ, REL

MODEL = json.loads((REL / "bundle/MANIFEST.json").read_text())["modes"]["indic"]
DECODER = str(REL / "bundle" / MODEL["decoder"])
ORACLE = json.loads((PROJ / "oracle/cache_canonical.json").read_text())["files"]
CACHE_RTTM = PROJ / "artifacts/oracle_cache/canonical"
WIN, STEP = 256000, 25600


def make_engine(core, nn, **kw):
    return core.Engine(nn, clustering=MODEL["clustering"], decoder_path=DECODER, **kw)


def record_jobs(items):
    jobs = []
    for set_name, path in items:
        d = records.load(path)
        jobs.append((f"{set_name}/{path.stem}", np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"],
                     list(map(float, d["frames_sw"])), float(d["duration_s"])))
    return jobs


def check_records_against_cache(core, items, post_threads):
    eng = make_engine(core, fake_nn.FakeNN(), post_threads=post_threads)
    got = {r["uri"]: r for r in eng.run_records(record_jobs(items))}
    assert len(got) == len(items)
    bad = []
    for set_name, path in items:
        key = f"{set_name}/{path.stem}"
        r = got[key]
        # the engine writes the uri it was given; the release writes the file id
        want = (CACHE_RTTM / set_name / "indic" / f"{path.stem}.rttm").read_bytes()
        if r["error"] is not None or r["rttm"].replace(f"SPEAKER {key} ", f"SPEAKER {path.stem} ").encode() != want:
            bad.append(key)
    return bad


def test_engine_records_sample(core):
    assert check_records_against_cache(core, records.sample(40), post_threads=8) == []


@pytest.mark.slow
def test_engine_records_all(core):
    assert check_records_against_cache(core, records.all_records(), post_threads=48) == []


# ---------------------------------------------------------------- audio -> RTTM with stand-in networks

@pytest.mark.parametrize("repack", [True, False])
def test_engine_audio_with_stand_in_networks(core, fake_reference, repack):
    items, ref = fake_reference
    log = json.loads((ref / "log.json").read_text())
    nn = fake_nn.FakeNN(devices=(0, 0, 1, 1))
    eng = make_engine(core, nn, repack=repack, keep_records=True, decode_threads=3, post_threads=6, max_ready=3)
    jobs = [(it["sid"], str(Path(it["path"])), it["path"]) for it in items]
    got = {r["uri"]: r for r in eng.run(jobs)}
    assert len(got) == len(items)
    bad = []
    for it in items:
        sid, r, want = it["sid"], got[it["sid"]], log[it["sid"]]
        if "error" in want or r["error"]:
            if r["error"] != want.get("error"):
                bad.append((sid, "error", want.get("error"), r["error"]))
            continue
        d = np.load(ref / f"{sid}.npz")
        rec = r["record"]
        same = (rec["P_bits"].tobytes() == np.ascontiguousarray(d["P"]).view(np.uint16).tobytes()
                and rec["P_bits"].shape == d["P"].shape
                and rec["masks"].tobytes() == d["masks"].tobytes() and rec["masks"].shape == d["masks"].shape
                and rec["emb"].tobytes() == d["emb"].tobytes() and rec["emb"].shape == d["emb"].shape
                and rec["chunk_sw"].tobytes() == d["chunk_sw"].tobytes()
                and rec["frames_sw"].tobytes() == d["frames_sw"].tobytes()
                and np.float64(rec["duration_s"]).tobytes() == d["duration_s"].tobytes())
        if not same:
            bad.append((sid, "record"))
        if r["rttm"] is None or r["rttm"].encode() != (ref / f"{sid}.rttm").read_bytes():
            bad.append((sid, "rttm"))
    assert not bad, bad
    # every network call is one of the release's shapes
    for call in nn.calls:
        if call[0] == "segment":
            assert call[1][1:] == (1, WIN) and 1 <= call[1][0] <= 8
        elif call[0] == "embed":
            assert call[1][1:] == (1, WIN) and call[1][0] in ((4,) if repack else (4, 8)) and call[2] == (call[1][0], 799)
        else:
            assert repack and call[1] == (8, 1, WIN) and call[2][1] == 8 and call[3][1:] == (8, 799)


# ---------------------------------------------------------------- robustness of the engine itself

class SlowFake(fake_nn.FakeNN):
    def segment(self, slot, x):
        import time
        time.sleep(0.05)
        return super().segment(slot, x)


class FailingFake(fake_nn.FakeNN):
    def segment(self, slot, x):
        # fails on any non-silent window
        if x[0, 0, :1000].sum() > 0:
            raise RuntimeError("CUDA out of memory (simulated)")
        return super().segment(slot, x)


def test_engine_stops_when_the_iterator_is_dropped(core, fake_reference):
    import time
    items, _ = fake_reference
    eng = make_engine(core, SlowFake(devices=(0,)), decode_threads=2, post_threads=2)
    it = eng.run([(i["sid"], i["path"], i["path"]) for i in items] * 3)
    first = next(it)
    t = time.perf_counter()
    # Drop: cancel and join
    del it
    assert time.perf_counter() - t < 30 and first["uri"]


def test_engine_reports_backend_errors_per_file(core, fake_reference):
    items, ref = fake_reference
    log = json.loads((ref / "log.json").read_text())
    eng = make_engine(core, FailingFake(), decode_threads=2, post_threads=2)
    got = list(eng.run([(i["sid"], i["path"], i["path"]) for i in items]))
    assert len(got) == len(items)
    for r in got:
        want = log[r["uri"]].get("error")
        assert r["error"] == want or r["error"] == "RuntimeError: CUDA out of memory (simulated)", (r["uri"], r["error"])


def test_engine_bits_do_not_depend_on_openblas_env(tmp_path):
    """The engine pins its private OpenBLAS to one thread itself: run with OPENBLAS_NUM_THREADS=64 in a fresh process."""
    import os
    import subprocess
    code = f"""
import json, sys, numpy as np
sys.path.insert(0, {str(PROJ / 'tests')!r})
import records, fake_nn
from anuvani_diar import _core
m = json.loads(open({str(REL / 'bundle/MANIFEST.json')!r}).read())["modes"]["indic"]
eng = _core.Engine(fake_nn.FakeNN(), clustering=m["clustering"], decoder_path={DECODER!r}, post_threads=8)
jobs = []
for s, p in records.sample(40):
    d = records.load(p)
    jobs.append((s + "/" + p.stem, np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"],
                 list(map(float, d["frames_sw"])), float(d["duration_s"])))
bad = 0
for r in eng.run_records(jobs):
    s, sid = r["uri"].split("/", 1)
    want = open({str(CACHE_RTTM)!r} + f"/{{s}}/indic/{{sid}}.rttm", "rb").read()
    bad += r["rttm"].replace("SPEAKER " + r["uri"] + " ", "SPEAKER " + sid + " ").encode() != want
print("BAD", bad)
"""
    env = dict(os.environ, OPENBLAS_NUM_THREADS="64", OMP_NUM_THREADS="64", MKL_NUM_THREADS="64")
    out = subprocess.run([sys.executable, "-s", "-c", code], env=env, capture_output=True, text=True, check=True)
    assert "BAD 0" in out.stdout, out.stdout + out.stderr


def test_a_low_sample_rate_cannot_resample_past_the_limit(core, tmp_path):
    """The size limit applies to the audio after resampling too, computed from the header before anything is
    allocated: 1,000 samples declared at 1 Hz are 4 KB decoded but 64 MB at 16 kHz."""
    import soundfile as sf
    low, ok = tmp_path / "one_hertz.wav", tmp_path / "eight_khz.wav"
    sf.write(low, np.full(1000, 0.1, np.float32), 1, subtype="FLOAT")
    sf.write(ok, np.full(16000, 0.1, np.float32), 8000, subtype="FLOAT")
    nn = fake_nn.FakeNN()
    eng = make_engine(core, nn, max_audio_bytes=1 << 20, decode_threads=1, post_threads=2)
    got = {r["uri"]: r for r in eng.run([("low", str(low), str(low)), ("ok", str(ok), str(ok))])}
    assert got["low"]["error"] == ("RefusedError: audio of 64000000 bytes (resampled to 16 kHz) is larger than this "
                                   "engine's limit of 1048576 bytes")
    assert got["ok"]["error"] is None


def test_a_rate_that_does_not_reduce_cannot_build_a_huge_resampling_filter(core, tmp_path):
    """The resampler's filter grows with the declared rate, not the file's length: 100 samples declared at 16,001 Hz
    would build a 2 GB table. Refused from the header before anything is resampled (docs/SECURITY.md); 44.1 kHz runs."""
    import soundfile as sf
    odd, ok = tmp_path / "odd.wav", tmp_path / "cd.wav"
    sf.write(odd, np.full(100, 0.1, np.float32), 16001, subtype="FLOAT")
    sf.write(ok, np.full(44100, 0.1, np.float32), 44100, subtype="FLOAT")

    class Counting(fake_nn.FakeNN):
        rates = []

        def resample(self, audio, sr):
            self.rates.append(sr)
            return super().resample(audio, sr)
    nn = Counting()
    eng = make_engine(core, nn, decode_threads=1, post_threads=2)
    got = {r["uri"]: r for r in eng.run([("odd", str(odd), str(odd)), ("ok", str(ok), str(ok))])}
    assert got["odd"]["error"] == ("RefusedError: resampling 16001 Hz audio to 16 kHz needs a filter of 2049920000 bytes, "
                                   "larger than this engine's limit of 67108864 bytes")
    assert got["ok"]["error"] is None and nn.rates == [44100]