# SPDX-License-Identifier: Apache-2.0
"""Validate the local HF package and caller-owned deployment contract on CUDA."""
from __future__ import annotations

import base64
import hashlib
import json
from pathlib import Path

import torch
from kernels import get_local_kernel

import benchmark as bench


def main() -> None:
    here = Path(__file__).resolve().parent
    runtime = here / "build/torch-cuda"
    metadata = json.loads((runtime / "metadata.json").read_text())
    for filename, expected in metadata["digest"]["files"].items():
        actual = base64.b64encode(hashlib.sha256((runtime / filename).read_bytes()).digest()).decode()
        assert actual == expected, f"digest mismatch: {filename}"
    op = get_local_kernel(here, trust_remote_code=True)
    assert all(hasattr(op, name) for name in ("Pool", "flush_inplace", "run"))
    namespace = {}
    exec(json.loads((here / "definition.json").read_text())["reference"], namespace)
    reference = namespace["run"]
    device = torch.device("cuda")
    ns = 8
    pool = op.Pool(ns, device=device)
    is_kda = hasattr(pool, "gbuf")
    idx = torch.tensor([0, 7, 2, 5, 1, 3], dtype=torch.int32, device=device)
    counters = torch.where(torch.arange(ns * 32, device=device).reshape(ns, 32) % 3 == 0, 16, 7).int()
    selected = torch.zeros(ns, dtype=torch.bool, device=device)
    selected[idx.long()] = True
    selected[0] = False
    due = selected[:, None] & (counters == 16)
    fields = ["codes", "s_k", "s_v", "uq", "kbuf", "ubuf", "w", "hcnt"] + (["gbuf", "pcum"] if is_kda else ["p"])
    q0 = torch.empty(128, 4, device=device)

    def restore(inputs):
        pool.codes.copy_(inputs[0]); pool.s_k.copy_(inputs[1]); pool.s_v.copy_(inputs[2])
        pool.uq[..., :128].copy_(inputs[3]); pool.uq[..., 128:].copy_(inputs[4])
        pool.kbuf.copy_(inputs[5]); pool.ubuf.copy_(inputs[6])
        if is_kda:
            pool.gbuf.copy_(inputs[7]); pool.w.copy_(inputs[8]); pool.pcum.copy_(inputs[9])
        else:
            pool.w.copy_(inputs[7]); pool.p[..., 0].copy_(inputs[8])
        pool.hcnt.copy_(counters)
        q0.copy_(inputs[-1])

    def snapshot():
        return {name: getattr(pool, name).clone() for name in fields}

    def verify(inputs, before):
        ref = reference(*[tensor.clone() for tensor in inputs])
        actual = (pool.codes, pool.s_k, pool.s_v, pool.uq[..., :128], pool.uq[..., 128:])
        # Each selected (slot, head) is a separate checked head, never averaged across slots.
        out = tuple(tensor[due].unsqueeze(0) for tensor in actual)
        ref = tuple(tensor[due].unsqueeze(0) for tensor in ref)
        exact = bench.exact_state(inputs)[due].unsqueeze(0)
        passed, ratio = bench.judge(out, ref, exact)
        assert passed, f"deployment accuracy failed: {ratio}"
        for name in fields:
            assert torch.equal(getattr(pool, name)[~due], before[name][~due]), f"non-due/slot-0 write: {name}"
        for name in ("kbuf", "ubuf"):
            assert torch.equal(getattr(pool, name), before[name]), f"update buffer changed: {name}"
        # The serving contract leaves hcnt at L ("flushed, empty"); the decode
        # step consumes that sentinel. Clearing it in flush would race scanners.
        assert torch.equal(pool.hcnt, before["hcnt"]), "flush changed fill counters"
        assert bool((pool.w[due] == 0).all())
        if is_kda:
            assert bool((pool.gbuf[due] == 1).all())
            assert bool((pool.pcum[due] == 0).all())
        else:
            assert bool((pool.p[due] == 1).all())
        return float(ratio)

    stream = torch.cuda.Stream()
    inputs = bench.structured_inputs(ns, device, seed=91)
    restore(inputs)
    torch.cuda.synchronize()
    before = snapshot()
    torch.cuda.synchronize()
    with torch.cuda.stream(stream):
        op.flush_inplace(pool, idx, q0)
    stream.synchronize()
    ratios = [verify(inputs, before)]
    before = snapshot()
    op.flush_inplace(pool, torch.empty(0, dtype=torch.int32, device=device), q0)
    torch.cuda.synchronize()
    assert all(torch.equal(getattr(pool, name), before[name]) for name in fields), "empty idx is not a no-op"

    restore(inputs)
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream):
        op.flush_inplace(pool, idx, q0)
    for seed in (92, 93):
        inputs = bench.structured_inputs(ns, device, seed=seed)
        # structured_inputs always uses the same q0; vary it explicitly here.
        generator = torch.Generator(device=device).manual_seed(seed)
        inputs = (*inputs[:-1], torch.randn(128, 4, device=device, generator=generator))
        restore(inputs)
        before = snapshot()
        torch.cuda.synchronize()
        graph.replay()
        torch.cuda.synchronize()
        ratios.append(verify(inputs, before))
        # Simulate the caller advancing the counters so no head is due.
        pool.hcnt[due] = 0
        after = snapshot()
        torch.cuda.synchronize()
        graph.replay()
        torch.cuda.synchronize()
        assert all(torch.equal(getattr(pool, name), after[name]) for name in fields), "replay with no due heads is not a no-op"
    print(json.dumps({
        "status": "pass", "gpu": torch.cuda.get_device_name(),
        "compute_capability": list(torch.cuda.get_device_capability()),
        "checks": ["runtime digests", "HF get_local_kernel", "mixed due heads", "unselected slots", "reserved slot 0", "window resets", "fill counters unchanged", "empty idx", "non-default stream", "graph capture", "graph replay with changed inputs and q0", "no-op replay with no due heads"],
        "deployment_worst_head_error_ratios": ratios,
    }, indent=2))


if __name__ == "__main__":
    main()