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

import os
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)
        self._no_mmap_tensors: dict[str, torch.Tensor] | None = None

    def _disable_mmap(self) -> bool:
        return os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}

    def _all_tensors_no_mmap(self) -> dict[str, torch.Tensor]:
        if self._no_mmap_tensors is None:
            from safetensors.torch import load

            with open(self.path, "rb") as file:
                tensors = load(file.read())
            self._no_mmap_tensors = {name: value.to(self.device) for name, value in tensors.items()}
        return self._no_mmap_tensors

    def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
        if self._disable_mmap():
            value = self._all_tensors_no_mmap()[name]
            return value.to(dtype=dtype) if dtype is not None else value

        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 = {}
        if self._disable_mmap():
            available_tensors = self._all_tensors_no_mmap()
            for suffix in names:
                name = f"{prefix}.{suffix}"
                if name in available_tensors:
                    tensors[name] = available_tensors[name]
            return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)

        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)