# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""diar-audio against soundfile 0.14.0 (libsndfile 1.2.2) and the release's `load_audio`: decoded float32 samples,
sample rates, channel means and error messages, compared bit for bit.

The engine's input path is `decode_mono` (Rust), then torchaudio's resample when the rate is not 16 kHz (the same
call the release makes, run by the network backend), then `validate` (Rust). `engine_load_audio` below composes
exactly that, so the end-to-end comparison with the release's `load_audio` covers every branch.
"""
import json
import os
import sys
import unicodedata
from multiprocessing import get_context
from pathlib import Path, PurePosixPath

import numpy as np
import pytest
import soundfile as sf
from hypothesis import given, strategies as st
from hypothesis.extra import numpy as hnp

from conftest import PROJ

TEST_AUDIO = PROJ / "artifacts/test_audio"


@pytest.fixture(scope="session")
def load_audio_ref():
    sys.path.insert(0, str(PROJ / "tests"))
    import refpkg
    refpkg.load()
    import torch
    # what load_bundle sets before any load_audio
    torch.set_num_threads(2)
    from anuvani_diar_ref.frontend import load_audio
    return load_audio


class EngineError(Exception):
    """An error the engine reports as the string "ExceptionName: message"."""


def outcome(fn, *a):
    try:
        return fn(*a), None
    except EngineError as e:
        return None, str(e)
    # the release's failure is part of the contract
    except Exception as e:
        return None, f"{type(e).__name__}: {e}"


def engine_load_audio(core, path):
    """decode_mono (Rust) -> torchaudio resample if needed (the release's own call) -> validate (Rust)."""
    import torch
    import torchaudio.functional as F
    name = str(path)
    audio, sr, err = core.decode_mono(name)
    if audio is None:
        raise EngineError(err)
    if sr != 16000:
        audio = F.resample(torch.from_numpy(audio).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32)
        sr = 16000
        err = core.validate_audio(name, audio)
    if err:
        raise EngineError(err)
    return audio, sr


def same_load(core, load_audio_ref, path):
    ref, ref_err = outcome(load_audio_ref, Path(path))
    got, got_err = outcome(engine_load_audio, core, Path(path))
    if ref_err or got_err:
        return ref_err is not None and got_err == ref_err, (ref_err, got_err)
    (ra, rsr), (ga, gsr) = ref, got
    return (rsr == gsr and ra.dtype == ga.dtype and ra.shape == ga.shape and ra.tobytes() == ga.tobytes()), None


# ---------------------------------------------------------------- synthetic files: every format and subtype

def _signal(rng, n, ch, kind):
    x = (rng.standard_normal((n, ch)) * 0.3).astype(np.float64)
    if kind == "clip":
        # full scale
        x[: n // 3] = np.sign(x[: n // 3])
    if kind == "edge":
        x = np.clip(x, -1, 1)
        if n > 8:
            x[0] = 1.0; x[1] = -1.0; x[2] = 0.0; x[3] = -0.0; x[4] = 1e-40; x[5] = 3.0; x[6] = -3.0
    return np.clip(x, -1, 1) if kind != "edge" else x


# (format, subtype, extension, rates)
SPECS = [
    ("WAV", "PCM_U8", "wav", (16000, 8000)), ("WAV", "PCM_16", "wav", (16000, 44100)), ("WAV", "PCM_24", "wav", (16000, 48000)),
    ("WAV", "PCM_32", "wav", (16000,)), ("WAV", "FLOAT", "wav", (16000, 22050)), ("WAV", "DOUBLE", "wav", (16000,)),
    ("WAV", "ULAW", "wav", (16000, 8000)), ("WAV", "ALAW", "wav", (8000,)), ("WAV", "IMA_ADPCM", "wav", (16000,)),
    ("WAV", "MS_ADPCM", "wav", (16000,)), ("WAV", "GSM610", "wav", (8000,)),
    ("FLAC", "PCM_S8", "flac", (16000,)), ("FLAC", "PCM_16", "flac", (16000, 44100)), ("FLAC", "PCM_24", "flac", (16000, 48000)),
    ("OGG", "VORBIS", "ogg", (16000, 44100)), ("OGG", "OPUS", "opus", (16000, 48000)),
    ("MP3", "MPEG_LAYER_III", "mp3", (16000, 44100)),
    ("AIFF", "PCM_16", "aiff", (16000,)), ("AIFF", "FLOAT", "aiff", (16000,)), ("AU", "PCM_16", "au", (16000,)),
    ("W64", "PCM_16", "w64", (16000,)), ("CAF", "PCM_16", "caf", (16000,)), ("RF64", "FLOAT", "rf64", (16000,)),
]
CHANNELS = (1, 2, 3, 6, 8)
LENGTHS = (1, 7, 4801, 3 * 16000 + 5)


@pytest.fixture(scope="session")
def synthetic(tmp_path_factory):
    d = tmp_path_factory.mktemp("audio")
    rng = np.random.default_rng(20261007)
    files = []
    for fmt, sub, ext, rates in SPECS:
        for ch in CHANNELS:
            if sub == "GSM610" and ch > 1:
                continue
            for n in LENGTHS:
                for sr in rates:
                    kind = ("plain", "clip", "edge")[(len(files)) % 3]
                    x = _signal(rng, n, ch, kind if sub in ("FLOAT", "DOUBLE") else ("clip" if kind == "clip" else "plain"))
                    p = d / f"{fmt}_{sub}_{ch}ch_{n}_{sr}.{ext}"
                    try:
                        sf.write(p, x if ch > 1 else x[:, 0], sr, subtype=sub, format=fmt)
                    except Exception:
                        # a combination libsndfile cannot write
                        continue
                    files.append(p)
    # float files carrying NaN / inf (the release rejects them after the mean)
    for name, val in (("nan", np.nan), ("posinf", np.inf), ("neginf", -np.inf)):
        x = _signal(rng, 32000, 2, "plain"); x[777, 1] = val
        p = d / f"WAV_FLOAT_special_{name}.wav"; sf.write(p, x, 16000, subtype="FLOAT"); files.append(p)
    # >= 8 channels: numpy's 8-accumulator pairwise sum
    for ch in (9, 12, 16, 33):
        x = _signal(rng, 9000, ch, "plain")
        p = d / f"WAV_FLOAT_{ch}ch.wav"; sf.write(p, x, 16000, subtype="FLOAT"); files.append(p)
    # empty files at both kinds of rate
    for sr in (16000, 8000):
        p = d / f"empty_{sr}.wav"; sf.write(p, np.zeros(0, np.float32), sr, subtype="PCM_16"); files.append(p)
    # not audio / unreadable
    (d / "text.wav").write_text("not a wav file\n"); files.append(d / "text.wav")
    trunc = d / "truncated.flac"; good = files[[str(f) for f in files].index(str(d / "FLAC_PCM_16_2ch_48005_16000.flac"))]
    trunc.write_bytes(good.read_bytes()[: len(good.read_bytes()) // 2]); files.append(trunc)
    (d / "zero_bytes.wav").write_bytes(b""); files.append(d / "zero_bytes.wav")
    files.append(d / "does_not_exist.wav")
    # a directory
    files.append(d)
    return files


def test_synthetic_matrix_is_broad(synthetic):
    assert len(synthetic) > 250, len(synthetic)


def test_sf_read_matches_soundfile(core, synthetic):
    core.sndfile_init()
    bad = []
    for p in synthetic:
        ref, ref_err = outcome(sf.read, str(p), -1, 0, None, "float32", False)
        got, got_err = core.sf_read_f32(str(p))
        if ref_err or got is None:
            if ref_err != got_err:
                bad.append((p.name, ref_err, got_err))
            continue
        ra, rsr = ref
        if not (rsr == got_err and ra.dtype == got.dtype and ra.shape == got.shape and ra.tobytes() == got.tobytes()):
            bad.append((p.name, "data", ra.shape, getattr(got, "shape", None)))
    assert not bad, bad[:10]


def test_load_audio_matches_release(core, load_audio_ref, synthetic):
    bad = [(p.name, info) for p in synthetic for ok, info in [same_load(core, load_audio_ref, p)] if not ok]
    assert not bad, bad[:10]


def test_robustness_files_match_release(core, load_audio_ref):
    items = json.loads((TEST_AUDIO / "MANIFEST.json").read_text())
    bad = [(i["sid"], info) for i in items for ok, info in [same_load(core, load_audio_ref, i["path"])] if not ok]
    assert len(items) == 28 and not bad, bad


# ---------------------------------------------------------------- channel mean

special = st.sampled_from([0.0, -0.0, 1.0, -1.0, 1e-45, -1e-45, 1e-38, 3.4e38, -3.4e38, np.inf, -np.inf, np.nan, 0.1, 1 / 3])


@given(hnp.arrays(np.float32, st.tuples(st.integers(0, 40), st.integers(1, 40)),
                  elements=st.one_of(st.floats(-1, 1, width=32), special.map(np.float32))))
def test_channel_mean_matches_numpy(core, x):
    got = core.channel_mean(np.ascontiguousarray(x))
    ref = x.mean(axis=1).astype(np.float32) if x.shape[0] else np.zeros(0, np.float32)
    na, nb = np.isnan(ref), np.isnan(got)
    assert np.array_equal(na, nb) and ref[~na].tobytes() == got[~nb].tobytes()


def test_channel_mean_long_rows(core):
    rng = np.random.default_rng(3)
    for ch in (127, 128, 129, 255, 1000, 9000):
        x = rng.standard_normal((5, ch)).astype(np.float32)
        assert core.channel_mean(x).tobytes() == x.mean(axis=1).astype(np.float32).tobytes(), ch


# ---------------------------------------------------------------- names and messages

def test_repr_every_code_point(core):
    bad = []
    for cp in range(0x110000):
        if 0xD800 <= cp <= 0xDFFF:
            continue
        s = chr(cp)
        if core.py_repr_str(s) != repr(s):
            bad.append(hex(cp))
    assert not bad, (len(bad), bad[:20])


@given(st.text(alphabet=st.sampled_from(list("ab'\"\\/ .\t\n\r\x00\x7f\x80\xa0कि​\U0001F600")), max_size=12))
def test_repr_strings(core, s):
    if "\x00" in s:
        return
    assert core.py_repr_str(s) == repr(s)


@pytest.mark.parametrize("p", ["a//b/./c/", "./a", "", ".", "//a", "///a", "a/../b", "/", "x/.wav", "/x/a.b.wav",
                               "x.", "a/b/", "..", "./..", "a/./././b", "/x/noext", "/x/.hidden.flac", "c d/e f.mp3"])
def test_pathlib(core, p):
    assert core.pathlib_str(p) == str(PurePosixPath(p))
    assert core.pathlib_stem(p) == PurePosixPath(p).stem


# ---------------------------------------------------------------- every replay input (slow)

def _check_one(item):
    from anuvani_diar import _core as core
    import refpkg
    refpkg.load()
    import torch
    torch.set_num_threads(2)
    from anuvani_diar_ref.frontend import load_audio
    ok, info = same_load(core, load_audio, item["path"])
    return item["set"] + "/" + item["sid"], ok, info


@pytest.mark.slow
def test_every_replay_input_decodes_identically():
    items = json.loads((PROJ / "artifacts/oracle_audio_full_inputs.json").read_text())
    with get_context("spawn").Pool(16, initializer=os.nice, initargs=(5,)) as pool:
        res = pool.map(_check_one, items, chunksize=4)
    bad = [r for r in res if not r[1]]
    assert len(res) == 1973 and not bad, bad[:10]


# ---------------------------------------------------------------- damaged files (seeded mutations)

MUTANT_BASES = [("WAV", "PCM_16", "wav", 2, 16000), ("WAV", "FLOAT", "wav", 1, 16000), ("FLAC", "PCM_16", "flac", 1, 16000),
                ("FLAC", "PCM_24", "flac", 2, 44100), ("OGG", "VORBIS", "ogg", 2, 16000), ("MP3", "MPEG_LAYER_III", "mp3", 1, 16000),
                ("AIFF", "PCM_16", "aiff", 1, 8000), ("W64", "PCM_16", "w64", 1, 16000), ("CAF", "PCM_16", "caf", 2, 16000),
                ("AU", "PCM_16", "au", 1, 16000), ("RF64", "FLOAT", "rf64", 1, 16000)]
EDGE_WORDS = (0, 1, 2, 0x7F, 0xFF, 0x7FFF, 0xFFFF, 0x7FFFFFFF, 0x80000000, 0xFFFFFFFF)
PLAIN_RATES = (8000, 16000, 22050, 44100, 48000)


def mutate(data, rng):
    """One seeded damage: bits flipped in the header or anywhere, a cut, a header word set to an edge value,
    bytes inserted or removed, or junk appended."""
    b = bytearray(data)
    kind = int(rng.integers(6))
    if kind == 0:
        for _ in range(int(rng.integers(1, 4))):
            i = int(rng.integers(min(len(b), 64)))
            b[i] ^= 1 << int(rng.integers(8))
    elif kind == 1:
        for _ in range(int(rng.integers(1, 8))):
            i = int(rng.integers(len(b)))
            b[i] ^= 1 << int(rng.integers(8))
    elif kind == 2:
        b = b[: int(rng.integers(len(b)))]
    elif kind == 3:
        i = 4 * int(rng.integers(min(len(b), 128) // 4))
        b[i:i + 4] = int(EDGE_WORDS[int(rng.integers(len(EDGE_WORDS)))]).to_bytes(4, "little" if rng.integers(2) else "big")
    elif kind == 4:
        i = int(rng.integers(len(b)))
        if rng.integers(2):
            b[i:i] = rng.integers(0, 256, int(rng.integers(1, 16)), dtype=np.uint8).tobytes()
        else:
            del b[i:i + int(rng.integers(1, 16))]
    else:
        b += rng.integers(0, 256, int(rng.integers(1, 4096)), dtype=np.uint8).tobytes()
    return bytes(b)


@pytest.fixture(scope="session")
def mutants(tmp_path_factory):
    d = tmp_path_factory.mktemp("mutants")
    rng = np.random.default_rng(20261008)
    files = []
    for fmt, sub, ext, ch, sr in MUTANT_BASES:
        x = _signal(rng, 8000, ch, "plain")
        base = d / f"base_{fmt}_{sub}.{ext}"
        sf.write(base, x if ch > 1 else x[:, 0], sr, subtype=sub, format=fmt)
        data = base.read_bytes()
        for k in range(250):
            p = d / f"m_{fmt}_{sub}_{k:03d}.{ext}"
            p.write_bytes(mutate(data, rng))
            files.append(p)
    return files


def test_damaged_files_decode_as_soundfile(core, mutants):
    """2,750 damaged files: the same samples and rate as soundfile, or the same error text."""
    core.sndfile_init()
    bad, kinds = [], set()
    for p in mutants:
        ref, ref_err = outcome(sf.read, str(p), -1, 0, None, "float32", False)
        got, got_err = core.sf_read_f32(str(p))
        if ref_err or got is None:
            kinds.add((ref_err or "").split(":")[0])
            if ref_err != got_err:
                bad.append((p.name, ref_err, got_err))
            continue
        ra, rsr = ref
        if not (rsr == got_err and ra.dtype == got.dtype and ra.shape == got.shape and ra.tobytes() == got.tobytes()):
            bad.append((p.name, "data", ra.shape, getattr(got, "shape", None)))
    assert not bad, bad[:10]
    # the suite does reach the error paths
    assert "LibsndfileError" in kinds


def test_damaged_files_load_as_the_release(core, load_audio_ref, mutants):
    """The whole input path (decode, channel mean, resample at the usual rates, validation) on the damaged files;
    files whose damage gives an unusual rate are compared at decoding only (above), since resampling an arbitrary
    rate is the release's own torchaudio call either way."""
    bad, n = [], 0
    for p in mutants:
        got, sr = core.sf_read_f32(str(p))
        if got is not None and sr not in PLAIN_RATES:
            continue
        n += 1
        ok, info = same_load(core, load_audio_ref, p)
        if not ok:
            bad.append((p.name, info))
    assert n > 1500 and not bad, (n, bad[:10])


def test_memory_error_text_matches_numpy(core):
    from numpy.core._exceptions import _ArrayMemoryError
    sizes = set()
    for k in range(0, 62):
        for d in (-3, -1, 0, 1, 3):
            sizes.add((1 << k) + d)
    for unit in range(1, 7):
        for v in (999.4, 999.5, 999.6, 1023.4, 1023.5, 1023.6, 1.0, 1.005, 9.995, 99.95):
            sizes.add(int(v * (1 << (10 * unit))))
    rng = np.random.default_rng(5)
    sizes.update(int(x) for x in rng.integers(1, 1 << 62, 3000, dtype=np.int64))
    bad = []
    # numpy refuses larger sizes before allocating
    for nb in sorted(s for s in sizes if 0 < s < 1 << 63):
        for ch in (1, 2):
            frames = nb // (4 * ch)
            if frames <= 0:
                continue
            shape = (frames,) if ch == 1 else (frames, ch)
            want = str(_ArrayMemoryError(shape, np.dtype("float32")))
            got = core.numpy_memory_error(frames, ch)
            if got != want:
                bad.append((frames, ch, want, got))
    assert not bad, bad[:5]