#!/usr/bin/env python3
"""Check converted NVFP4 files against a local ComfyUI source checkout."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from safetensors import safe_open


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("comfyui", type=Path, help="Path to a ComfyUI source checkout")
    parser.add_argument(
        "files",
        nargs="*",
        type=Path,
        help="Converted *_comfy.safetensors files. Defaults to all *_comfy.safetensors in cwd.",
    )
    return parser.parse_args()


def read_subset(path: Path, limit_layers: int = 3) -> tuple[dict[str, torch.Tensor], dict[str, str]]:
    state_dict: dict[str, torch.Tensor] = {}
    with safe_open(path, framework="pt", device="cpu") as sf:
        metadata = sf.metadata() or {}
        layers = sorted(json.loads(metadata["_quantization_metadata"])["layers"])[:limit_layers]
        for layer in layers:
            for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale", "comfy_quant"):
                key = f"{layer}.{suffix}"
                state_dict[key] = sf.get_tensor(key)
        return state_dict, metadata


def main() -> None:
    args = parse_args()
    sys.path.insert(0, str(args.comfyui.resolve()))

    import comfy.quant_ops as quant_ops
    import comfy.utils as comfy_utils

    nvfp4 = quant_ops.QUANT_ALGOS.get("nvfp4")
    assert nvfp4 is not None, "ComfyUI source has no nvfp4 quant algo"
    assert nvfp4["storage_t"] == torch.uint8, nvfp4
    assert {"weight_scale", "weight_scale_2", "input_scale"} <= nvfp4["parameters"], nvfp4
    assert nvfp4["comfy_tensor_layout"] == "TensorCoreNVFP4Layout", nvfp4

    files = args.files or sorted(Path(".").glob("*_comfy.safetensors"))
    for path in files:
        state_dict, metadata = read_subset(path)
        converted, _ = comfy_utils.convert_old_quants(dict(state_dict), metadata=metadata)
        comfy_quant_keys = sorted(k for k in converted if k.endswith(".comfy_quant"))
        if not comfy_quant_keys:
            raise AssertionError(f"{path.name}: ComfyUI did not expose comfy_quant keys")
        for key in comfy_quant_keys:
            marker = json.loads(bytes(converted[key].tolist()).decode("utf-8"))
            if marker.get("format") != "nvfp4":
                raise AssertionError(f"{path.name}: {key} marker is {marker!r}")
        print(f"OK {path.name}: ComfyUI source recognizes nvfp4 metadata and comfy_quant keys")


if __name__ == "__main__":
    main()