# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Identity check through the serving stack (Triton, the Rust gateway, the Python client) against the raw-audio oracle,
byte for byte: every RTTM, every error, and every front-end record the gateway dumps. Jobs arrive concurrently in a
shuffled order; most name a file on the server, the rest are uploaded through the gateway's spool.

Run through triton/identity_check.sh (it starts Triton and the gateway and sets GATEWAY_URL; its variants set
IDENTITY_SEED for another arrival order and RESTART_AT for a Triton restart in the middle of the run).
"""
import json
import os
import random
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import pytest

from conftest import PROJ
from test_identity_audio import INPUTS, identity_sample, oracle_log

ORACLE = PROJ / "artifacts/oracle_audio"
RECORDS = PROJ / "artifacts/identity_serving/records"
pytestmark = pytest.mark.skipif("GATEWAY_URL" not in os.environ, reason="needs triton/identity_check.sh")


def run_through_gateway(keys, uploads=20, seed=7):
    from anuvani_diar import Client
    c = Client(os.environ["GATEWAY_URL"], timeout=3600)
    order = list(keys)
    random.Random(seed).shuffle(order)
    # uploads only for files that succeed: a failing upload's message names the spool file, not the client's path
    log = oracle_log()
    upload = set([k for k in order if log[k]["ok"]][:uploads])

    def one(k):
        it = INPUTS[k]
        job = c.submit(it["path"], uri=it["sid"], server_path=k not in upload)
        return k, c.wait(job, poll=1.0)
    with ThreadPoolExecutor(64) as ex:
        return dict(ex.map(one, order))


def compare(keys, got):
    log = oracle_log()
    bad = []
    for k in keys:
        s, sid = k.split("/", 1)
        want, r = log[k], got[k]
        if not want["ok"]:
            if r["error"] != want["error"]:
                bad.append((k, "error", want["error"], r["error"]))
            continue
        # the release's own post-processing raised
        want_err = want["indic"].get("error")
        if r["error"] != want_err:
            bad.append((k, "error", want_err, r["error"]))
            continue
        if want_err is None and (r["status"] != "done" or (r["rttm"] or "").encode() != (ORACLE / s / "indic" / f"{sid}.rttm").read_bytes()):
            bad.append((k, r["status"], "rttm"))
        d = np.load(ORACLE / s / "records" / f"{sid}.npz")
        P = np.fromfile(RECORDS / f"{sid}.P", dtype="<u2")
        masks = np.fromfile(RECORDS / f"{sid}.masks", dtype=np.uint8)
        emb = np.fromfile(RECORDS / f"{sid}.emb", dtype="<f4")
        if (P.tobytes() != np.ascontiguousarray(d["P"]).view(np.uint16).tobytes() or masks.tobytes() != d["masks"].tobytes()
                or emb.tobytes() != d["emb"].tobytes()):
            bad.append((k, "record"))
    return bad


def gateway_metric(name):
    with urllib.request.urlopen(os.environ["GATEWAY_URL"] + "/metrics") as r:
        for line in r.read().decode().splitlines():
            if line.startswith(name + " "):
                return float(line.split()[1])


SEED = int(os.environ.get("IDENTITY_SEED", "7"))


def quick_sample():
    """check-exact's set (tools/check_exact.py): the golden 12, every robustness input, the 1 h file, 20 sampled files."""
    import records
    keys = [f"bench/{i['sid']}" for i in json.loads((PROJ / "artifacts/golden12_inputs.json").read_text())]
    keys += [k for k, i in INPUTS.items() if i["set"] == "robust"] + ["long/long_1h"]
    keys += [f"{s}/{p.stem}" for s, p in records.sample(20)]
    return [k for k in dict.fromkeys(keys) if k in INPUTS]


@pytest.mark.gpu
def test_identity_serving_sample():
    keys = quick_sample() if os.environ.get("IDENTITY_QUICK") else identity_sample()
    got = run_through_gateway(keys, seed=SEED)
    assert len(got) == len(keys)
    assert compare(keys, got) == []
    # the restart hit work in flight, which was sent again
    if os.environ.get("RESTART_AT"):
        assert gateway_metric("anuvani_diar_triton_retries_total") > 0


@pytest.mark.gpu
@pytest.mark.slow
def test_identity_serving_all():
    keys = list(INPUTS)
    got = run_through_gateway(keys, uploads=100, seed=SEED)
    assert compare(keys, got) == []
    if os.environ.get("RESTART_AT"):
        assert gateway_metric("anuvani_diar_triton_retries_total") > 0