# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Shared test fixtures. Every comparison is bitwise: arrays are compared through their raw bytes, so -0.0 vs 0.0
and NaN payloads count as differences."""
import os
import sys
from pathlib import Path

import numpy as np
import pytest
from hypothesis import settings

PROJ = Path(__file__).resolve().parents[1]
# untracked reference data (the release's weights, the benchmark manifest, artifacts/) lives in the main checkout; a
# linked git worktree reads it there with ANUVANI_DIAR_MAIN_TREE=<the main checkout's SPEAKER-ID> (read only)
_MAIN = os.environ.get("ANUVANI_DIAR_MAIN_TREE")
SID = Path(_MAIN) if _MAIN else PROJ.parents[1]
# frozen reference package (read-only)
REL = SID / "release/anuvani_diar_v1"
ART = SID / "release/anuvani_diar_rs/artifacts" if _MAIN else PROJ / "artifacts"
# development data (benchmark manifests, cached front-end records, the research pipeline's outputs): not part of the
# repository; ANUVANI_DIAR_DEV_DATA names it, by default artifacts/dev_data (docs/TESTING.md)
DEV = Path(os.environ.get("ANUVANI_DIAR_DEV_DATA") or ART / "dev_data")

# deterministic property tests: fixed derandomised seed, no example database
settings.register_profile("det", derandomize=True, database=None, max_examples=2000, deadline=None, print_blob=True)
settings.load_profile("det")

# the differential tests compute their references with numpy in this process, whose BLAS bits depend on its thread
# count; a test that only compares engine output with stored oracles may run with other thread settings on purpose
# (it sets ANUVANI_DIAR_ANY_BLAS_THREADS=1 for its own subprocess)
if ((os.environ.get("OPENBLAS_NUM_THREADS") != "1" and not os.environ.get("ANUVANI_DIAR_ANY_BLAS_THREADS"))
        or not sys.flags.no_user_site):
    raise RuntimeError("run tests through `make` (canonical environment: env/canonical.sh, python -s)")


def bits_equal(a, b):
    """Same dtype, shape and bits for every non-NaN value (so -0.0 != 0.0), and NaN in exactly the same places.

    NaN payload and sign bits are not compared: the compiled numpy loops may pick a different NaN payload when two
    NaNs meet, and no NaN payload can reach an RTTM (every downstream use only tests NaN-ness)."""
    a, b = np.asarray(a), np.asarray(b)
    if a.dtype != b.dtype or a.shape != b.shape:
        return False
    if a.dtype.kind != "f":
        return a.tobytes() == b.tobytes()
    na, nb = np.isnan(a), np.isnan(b)
    if not np.array_equal(na, nb):
        return False
    return a[~na].tobytes() == b[~nb].tobytes()


@pytest.fixture(scope="session")
def core():
    from anuvani_diar import _core
    return _core


# ---------------------------------------------------------------- the release with stand-in networks (CPU)

def fake_inputs(tmp):
    """The robustness files, window-count edges cut from one benchmark file, and five whole benchmark files."""
    import json
    import soundfile as sf
    win, step = 256000, 25600
    items = [dict(sid=i["sid"], path=i["path"]) for i in json.loads((ART / "test_audio/MANIFEST.json").read_text())
             if i["set"] == "robust"]
    rows = json.loads((DEV / "manifests/bench_manifest.json").read_text())["rows"]
    src = sorted(rows, key=lambda r: r["sample_id"])
    a, _ = sf.read(src[0]["flac_path"], dtype="float32")
    for n in (win - 1, win, win + 1, win + 7 * step, win + 8 * step - 1, win + 8 * step, win + 8 * step + 1,
              win + 15 * step + 5, win + 16 * step, 30 * 16000 + 7):
        p = tmp / f"len_{n}.flac"
        sf.write(p, a[:n], 16000, subtype="PCM_16")
        items.append(dict(sid=p.stem, path=str(p)))
    # five whole benchmark files
    for r in src[1:6]:
        items.append(dict(sid=r["sample_id"], path=r["flac_path"]))
    return items


@pytest.fixture(scope="session")
def fake_reference(tmp_path_factory):
    """(inputs, folder): the release's own `analyse` and `indic_mode` on `fake_inputs` with the stand-in networks of
    fake_nn.py; the folder holds log.json (errors), <sid>.npz (records) and <sid>.rttm."""
    import json
    import subprocess
    tmp = tmp_path_factory.mktemp("fake")
    items = fake_inputs(tmp)
    (tmp / "inputs.json").write_text(json.dumps(items))
    subprocess.run([sys.executable, "-s", "-W", "ignore", str(PROJ / "tests/ref_fake_run.py"), str(tmp / "inputs.json"),
                    str(tmp / "ref")], check=True, capture_output=True)
    return items, tmp / "ref"