"""Lazy loading for the current Comfy-format H3 safetensors checkpoint."""

from pathlib import Path

import torch

from .nvfp4 import Nvfp4Linear, load_nvfp4_linear


class H3Checkpoint:
    """Load individual tensors/modules without materializing the whole checkpoint."""

    def __init__(self, path: str | Path, device: str | torch.device = "cuda"):
        self.path = str(path)
        self.device = str(device)

    def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
        from safetensors import safe_open

        with safe_open(self.path, framework="pt", device=self.device) as checkpoint:
            value = checkpoint.get_tensor(name)
        return value.to(dtype=dtype) if dtype is not None else value

    def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
        names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias", "pre_quant_scale")
        tensors = {}
        from safetensors import safe_open

        with safe_open(self.path, framework="pt", device=self.device) as checkpoint:
            available = set(checkpoint.keys())
            for suffix in names:
                name = f"{prefix}.{suffix}"
                if name in available:
                    tensors[name] = checkpoint.get_tensor(name)
        return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)