"""Byte-compare isolated SM89 extensions on adversarial short NHD shapes."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
from pathlib import Path

import torch


CASES = (
    (1, 1),
    (31, 63),
    (32, 64),
    (33, 65),
    (63, 127),
    (64, 128),
    (65, 129),
    (127, 63),
    (128, 64),
    (129, 65),
    (191, 191),
    (192, 192),
    (193, 193),
)


def load_extension(path: Path):
    spec = importlib.util.spec_from_file_location("sageattention._qattn_sm89", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load extension from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def digest(value: torch.Tensor) -> str:
    data = value.detach().contiguous().view(torch.uint8).cpu().numpy()
    return hashlib.sha256(memoryview(data)).hexdigest()


def launch(extension, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale):
    output = torch.empty(q_int8.shape, dtype=torch.bfloat16, device=q_int8.device)
    extension.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
        q_int8,
        k_int8,
        v_fp8,
        output,
        q_scale,
        k_scale,
        v_scale,
        0,
        0,
        2,
        128**-0.5,
        0,
    )
    return output


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--baseline", type=Path, required=True)
    parser.add_argument("--candidate", type=Path, required=True)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--heads", type=int, default=2)
    parser.add_argument("--seed", type=int, default=73021)
    parser.add_argument("--quick", action="store_true", help="Run only one one-tile and one multi-tile case.")
    args = parser.parse_args()

    import sageattention.core as sage_core

    baseline = load_extension(args.baseline)
    candidate = load_extension(args.candidate)
    cases = ((1, 1), (129, 193)) if args.quick else CASES
    results = []

    with torch.inference_mode():
        for index, (q_len, kv_len) in enumerate(cases):
            generator = torch.Generator(device="cuda").manual_seed(args.seed + index)
            q = torch.randn(
                (1, q_len, args.heads, 128),
                dtype=torch.bfloat16,
                device="cuda",
                generator=generator,
            )
            k = torch.randn(
                (1, kv_len, args.heads, 128),
                dtype=torch.bfloat16,
                device="cuda",
                generator=generator,
            )
            v = torch.randn(
                (1, kv_len, args.heads, 128),
                dtype=torch.bfloat16,
                device="cuda",
                generator=generator,
            )
            k_mean = k.mean(dim=1, keepdim=True)
            q_int8, q_scale, k_int8, k_scale = sage_core.per_warp_int8_cuda(
                q,
                k,
                k_mean,
                BLKQ=128,
                WARPQ=32,
                BLKK=64,
                tensor_layout="NHD",
            )
            v_fp8, v_scale, _ = sage_core.per_channel_fp8(
                v,
                tensor_layout="NHD",
                scale_max=2.25,
                smooth_v=False,
            )
            expected = launch(
                baseline, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale,
            )
            actual = launch(
                candidate, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale,
            )
            torch.cuda.synchronize()
            delta = actual.float() - expected.float()
            result = {
                "q_len": q_len,
                "kv_len": kv_len,
                "equal": torch.equal(actual, expected),
                "different_elements": int(torch.count_nonzero(actual != expected).item()),
                "max_abs": delta.abs().max().item(),
                "baseline_sha256": digest(expected),
                "candidate_sha256": digest(actual),
            }
            results.append(result)
            print(json.dumps(result), flush=True)
            if not result["equal"]:
                raise RuntimeError(f"short-shape parity failed: {result}")

    report = {
        "device": torch.cuda.get_device_name(),
        "baseline": str(args.baseline),
        "candidate": str(args.candidate),
        "heads": args.heads,
        "seed": args.seed,
        "cases": results,
    }
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()