#!/usr/bin/env python3
"""Validate a packaged Mage-VL FP8 or NVFP4 Hugging Face repository."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

from safetensors import safe_open


ROOT = Path(__file__).resolve().parent


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


def fail(message: str) -> None:
    raise RuntimeError(message)


def main() -> None:
    config = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    quantization = config.get("mage_vl_quantization")
    if not isinstance(quantization, dict):
        fail("config.json has no mage_vl_quantization dictionary")
    format_name = quantization.get("format")
    if format_name not in {"scaled_fp8_w8a8", "native_nvfp4_w4a4"}:
        fail(f"unsupported quantization format: {format_name}")
    if format_name == "native_nvfp4_w4a4":
        if quantization.get("runtime_version") != 2:
            fail("NVFP4 release must declare runtime_version 2")
        if quantization.get("default_profile") != "hybrid_fast":
            fail("NVFP4 default_profile must be hybrid_fast")
        profiles = quantization.get("runtime_profiles")
        if not isinstance(profiles, dict):
            fail("NVFP4 release has no runtime_profiles dictionary")
        if set(profiles) != {"hybrid_fast", "pure_w4a4_v2"}:
            fail(f"unexpected NVFP4 runtime profiles: {sorted(profiles)}")
        hybrid = profiles["hybrid_fast"]
        pure = profiles["pure_w4a4_v2"]
        if hybrid.get("smallm_backend") != "w4a16_gemv":
            fail("hybrid_fast must select w4a16_gemv")
        if hybrid.get("shared_gate_up_activation") or hybrid.get(
            "shared_qkv_activation"
        ):
            fail("hybrid_fast must not enable shared W4A4 activations")
        if pure.get("smallm_backend") != "off":
            fail("pure_w4a4_v2 must disable the W4A16 small-M backend")
        if not pure.get("shared_gate_up_activation") or not pure.get(
            "shared_qkv_activation"
        ):
            fail("pure_w4a4_v2 must enable both shared W4A4 activations")

    index = json.loads(
        (ROOT / "model.safetensors.index.json").read_text(encoding="utf-8")
    )
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict) or not weight_map:
        fail("checkpoint index has no weight map")
    shard_names = sorted(set(weight_map.values()))
    discovered: dict[str, str] = {}
    logical_size = 0
    for shard_name in shard_names:
        shard_path = ROOT / shard_name
        if not shard_path.is_file():
            fail(f"missing checkpoint shard: {shard_name}")
        with safe_open(shard_path, framework="pt", device="cpu") as handle:
            for key in handle.keys():
                if key in discovered:
                    fail(f"duplicate tensor across shards: {key}")
                discovered[key] = shard_name
                tensor_slice = handle.get_slice(key)
                shape = tensor_slice.get_shape()
                dtype = tensor_slice.get_dtype()
                element_sizes = {
                    "BF16": 2,
                    "F32": 4,
                    "F8_E4M3": 1,
                    "U8": 1,
                }
                if dtype not in element_sizes:
                    fail(f"unsupported validation dtype {dtype} for {key}")
                elements = 1
                for dimension in shape:
                    elements *= dimension
                logical_size += elements * element_sizes[dtype]
    if discovered != weight_map:
        missing = sorted(set(weight_map) - set(discovered))
        extra = sorted(set(discovered) - set(weight_map))
        fail(f"checkpoint index mismatch; missing={missing[:3]} extra={extra[:3]}")
    if logical_size != int(index["metadata"]["total_size"]):
        fail(
            f"logical checkpoint size mismatch: {logical_size} != "
            f"{index['metadata']['total_size']}"
        )

    qdata = [key for key in weight_map if key.endswith(".qdata")]
    scales = [key for key in weight_map if key.endswith(".weight_scale")]
    block_scales = [
        key for key in weight_map if key.endswith(".weight_block_scale")
    ]
    if len(qdata) != 252 or len(scales) != 252:
        fail(
            f"expected 252 qdata/scales, got {len(qdata)}/{len(scales)}"
        )
    if format_name == "native_nvfp4_w4a4" and len(block_scales) != 252:
        fail(f"expected 252 NVFP4 block scales, got {len(block_scales)}")
    if format_name == "scaled_fp8_w8a8" and block_scales:
        fail("FP8 checkpoint unexpectedly contains NVFP4 block scales")

    manifest = json.loads((ROOT / "MANIFEST.json").read_text(encoding="utf-8"))
    for record in manifest.get("files", []):
        path = ROOT / record["path"]
        if not path.is_file():
            fail(f"manifest file is missing: {record['path']}")
        if path.stat().st_size != record["size"]:
            fail(f"manifest size mismatch: {record['path']}")
        if sha256(path) != record["sha256"]:
            fail(f"manifest hash mismatch: {record['path']}")

    print(
        json.dumps(
            {
                "status": "pass",
                "format": format_name,
                "checkpoint_tensors": len(weight_map),
                "checkpoint_shards": len(shard_names),
                "checkpoint_logical_bytes": logical_size,
                "quantized_language_projections": len(qdata),
                "manifest_files": len(manifest.get("files", [])),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()