"""Selective safetensors loading for component and checkpoint sources."""

from __future__ import annotations

import torch
from safetensors import safe_open
from safetensors.torch import load_file

from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError

SAFETENSOR_DTYPE_MAP = {
    "BF16": torch.bfloat16,
    "F16": torch.float16,
    "F32": torch.float32,
}


def loadStateDict(source: ModelSource) -> dict[str, torch.Tensor]:
    """Load tensors for one model component from a safetensors source."""

    if source.checkpoint_prefix is None:
        return load_file(source.path, device="cpu")

    state_dict: dict[str, torch.Tensor] = {}
    with safe_open(source.path, framework="pt", device="cpu") as tensors:
        for key in tensors.keys():
            if key.startswith(source.checkpoint_prefix):
                stripped_key = key.removeprefix(source.checkpoint_prefix)
                state_dict[stripped_key] = tensors.get_tensor(key)

    if not state_dict:
        raise DiffusionCliError(
            f"Checkpoint does not contain {source.role}: {source.path}"
        )
    return state_dict


def inspectSourceTorchDtype(source: ModelSource) -> torch.dtype | None:
    """Return a single floating torch dtype from source metadata if clear."""

    dtypes: set[str] = set()
    matched = False
    with safe_open(source.path, framework="pt", device="cpu") as tensors:
        for key in tensors.keys():
            if source.checkpoint_prefix is not None:
                if not key.startswith(source.checkpoint_prefix):
                    continue
            matched = True
            dtypes.add(tensors.get_slice(key).get_dtype())

    if source.checkpoint_prefix is not None and not matched:
        raise DiffusionCliError(
            f"Checkpoint does not contain {source.role}: {source.path}"
        )

    mapped = {SAFETENSOR_DTYPE_MAP[dtype] for dtype in dtypes
              if dtype in SAFETENSOR_DTYPE_MAP}
    if len(mapped) == 1 and len(mapped) == len(dtypes):
        return next(iter(mapped))
    return None