"""Safetensors metadata inspection for local model files."""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
from pathlib import Path

from safetensors import safe_open

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


@dataclass(frozen=True)
class TensorSummary:
    """Small metadata summary for a safetensors checkpoint."""

    path: Path
    source_prefix: str | None
    tensor_count: int
    dtype_counts: dict[str, int]
    top_level_counts: dict[str, int]
    architecture_guess: str


def _topLevel(key: str) -> str:
    return key.split(".", 1)[0]


def guessArchitecture(keys: list[str]) -> str:
    """Guess a checkpoint role from key patterns."""

    key_set = set(keys)
    joined = "\n".join(keys[:5000])
    if any(key.startswith("decoder.") for key in key_set):
        return "vae"
    if "model.embed_tokens.weight" in key_set:
        return "qwen_text_encoder"
    if "embed_tokens.weight" in key_set:
        return "qwen_text_encoder"
    if (
        "img_in.weight" in key_set
        or "x_embedder.proj.weight" in key_set
        or "x_embedder.weight" in key_set
    ):
        return "z_image_diffusion"
    if "diffusion_model.img_in.weight" in key_set:
        return "z_image_diffusion"
    if any(key.startswith("model.diffusion_model.") for key in key_set):
        return "z_image_diffusion"
    if "double_blocks.0.img_attn.qkv.weight" in joined:
        return "z_image_diffusion"
    return "unknown"


def inspectSafetensors(path: Path) -> TensorSummary:
    """Read safetensors metadata without loading tensor data into memory."""

    return inspectModelSource(ModelSource(path=path, role="unknown"))


def inspectModelSource(source: ModelSource) -> TensorSummary:
    """Read metadata for one resolved model source."""

    dtype_counts: Counter[str] = Counter()
    top_counts: Counter[str] = Counter()
    keys: list[str] = []

    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
                summary_key = key.removeprefix(source.checkpoint_prefix)
            else:
                summary_key = key
            keys.append(summary_key)
            metadata = tensors.get_slice(key)
            dtype_counts[metadata.get_dtype()] += 1
            top_counts[_topLevel(summary_key)] += 1

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

    return TensorSummary(
        path=source.path,
        source_prefix=source.checkpoint_prefix,
        tensor_count=len(keys),
        dtype_counts=dict(sorted(dtype_counts.items())),
        top_level_counts=dict(top_counts.most_common(12)),
        architecture_guess=guessArchitecture(keys),
    )


def formatSummary(name: str, summary: TensorSummary) -> str:
    """Format a checkpoint summary for terminal output."""

    dtype_text = ", ".join(
        f"{dtype}: {count}" for dtype, count in summary.dtype_counts.items()
    )
    top_text = ", ".join(
        f"{prefix}: {count}"
        for prefix, count in summary.top_level_counts.items()
    )
    return "\n".join(
        [line for line in [
            f"{name}: {summary.path}",
            (
                f"  source prefix: {summary.source_prefix}"
                if summary.source_prefix is not None
                else None
            ),
            f"  tensors: {summary.tensor_count}",
            f"  dtypes: {dtype_text}",
            f"  top-level keys: {top_text}",
            f"  architecture guess: {summary.architecture_guess}",
        ] if line is not None]
    )