"""Capture fixed aligned prequantized Sage2 fixtures for VEA-B Phase 2B."""

from __future__ import annotations

import argparse
import hashlib
import json
import platform
from pathlib import Path

import torch


SHAPE = {"batch": 1, "heads": 1, "q": 128, "kv": 192, "d": 128}
SEED = 73021


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


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def save(output_dir: Path, name: str, value: torch.Tensor) -> dict:
    host = value.detach().contiguous().cpu()
    path = output_dir / f"{name}.pt"
    torch.save({"tensor": host}, path)
    loaded = torch.load(path, map_location="cpu", weights_only=True)["tensor"]
    if not torch.equal(host, loaded):
        raise RuntimeError(f"reload mismatch for {name}")
    return {
        "path": str(path),
        "shape": list(host.shape),
        "dtype": str(host.dtype),
        "tensor_sha256": tensor_sha256(host),
        "file_sha256": file_sha256(path),
        "size_bytes": path.stat().st_size,
        "reload_verified": True,
    }


def public_mainloop(q, k, v, q_scale, k_scale, v_scale):
    import sageattention.core as sage_core

    output = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
    sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
        q, k, v, output, q_scale, k_scale, v_scale,
        0, 0, 2, SHAPE["d"] ** -0.5, 0,
    )
    return output


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--image", required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    generator = torch.Generator(device="cuda").manual_seed(SEED)
    q_bf16 = torch.randn((1, 128, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)
    k_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)
    v_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)

    import sageattention.core as sage_core

    with torch.inference_mode():
        k_mean = k_bf16.mean(dim=1, keepdim=True)
        q_int8, q_scale, k_int8, k_scale = sage_core.per_warp_int8_cuda(
            q_bf16, k_bf16, k_mean, BLKQ=128, WARPQ=32, BLKK=64,
            tensor_layout="NHD",
        )
        v_fp8, v_scale, _ = sage_core.per_channel_fp8(
            v_bf16, tensor_layout="NHD", scale_max=2.25, smooth_v=False,
        )
        public_output = public_mainloop(q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale)
        torch.cuda.synchronize()

    tensors = {
        "q_bf16": q_bf16,
        "k_bf16": k_bf16,
        "v_bf16": v_bf16,
        "k_mean_bf16": k_mean,
        "q_int8": q_int8,
        "q_scale_fp32": q_scale,
        "k_int8": k_int8,
        "k_scale_fp32": k_scale,
        "v_fp8_e4m3": v_fp8,
        "v_scale_fp32": v_scale,
        "public_output_bf16": public_output,
    }
    records = {name: save(args.output_dir, name, value) for name, value in tensors.items()}
    manifest = {
        "schema": "vortex-exact-phase2b-aligned-fixtures",
        "version": 1,
        "status": "captured_and_reload_verified",
        "shape": SHAPE,
        "epochs": 3,
        "seed": SEED,
        "scope": "prepared BF16 retained only for provenance; prototype inputs are prequantized tensors",
        "reference": {
            "sageattention_version": "2.2.0",
            "sageattention_commit": "d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5",
            "mainloop": "qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf",
            "tensor_layout": "NHD",
            "causal": False,
            "qk_quant_granularity": "per_warp",
            "v_scale_max": 2.25,
        },
        "environment": {
            "image": args.image,
            "gpu": torch.cuda.get_device_name(),
            "compute_capability": list(torch.cuda.get_device_capability()),
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
            "python": platform.python_version(),
        },
        "tensors": records,
    }
    manifest_path = args.output_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(manifest, indent=2), flush=True)


if __name__ == "__main__":
    main()