#!/usr/bin/env python3
"""Print a summary of a Hugging Face checkpoint.

Weight bytes are not read. config.json, the tokenizer files, and safetensors
headers are enough for the architecture, quantization, and tensor list.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import struct
import sys
from collections import Counter
from dataclasses import dataclass
from pathlib import Path

# Same round-trip check as nvfp4gguf.bundle.chat_template_has_tools.
_TOOL_RESULT_MARKERS = (
    "tool_response",
    "tool_results",
    "role'] == 'tool'",
    "role'] == \"tool\"",
    "role\"] == 'tool'",
    "role\"] == \"tool\"",
    "message.role == 'tool'",
    "message.role == \"tool\"",
    "ipython",
)

_NESTED_CONFIGS = (
    "text_config",
    "llm_config",
    "language_config",
    "thinker_config",
    "decoder",
)

_INDEX_NAMES = (
    "model.safetensors.index.json",
    "diffusion_pytorch_model.safetensors.index.json",
)

_MAX_HEADER = 64 * 1024 * 1024


@dataclass
class TensorInfo:
    name: str
    dtype: str
    shape: list[int]
    nbytes: int | None
    file: str


@dataclass
class ShardInfo:
    path: Path
    header_size: int
    data_end: int
    metadata: dict


@dataclass
class Checkpoint:
    path: Path
    root: Path
    config: dict
    config_name: str
    generation: dict
    tokenizer: dict
    quant: dict
    hf_quant: dict | None
    recipe: str | None
    base_model: str | None
    tensors: list[TensorInfo]
    shards: list[ShardInfo]
    shard_total: int
    partial: bool
    index_meta: dict
    template: str
    template_where: str
    template_count: int
    nvfp4: bool
    comfy: bool


def format_size(num_bytes: int) -> str:
    size = float(num_bytes)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024 or unit == "TiB":
            return f"{int(size)} B" if unit == "B" else f"{size:.2f} {unit}"
        size /= 1024
    return f"{size:.2f} TiB"


def format_params(count: int) -> str:
    if count >= 1_000_000_000:
        return f"{count} ({count / 1_000_000_000:.2f}B)"
    if count >= 1_000_000:
        return f"{count} ({count / 1_000_000:.2f}M)"
    return str(count)


def format_number(value) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, int):
        return str(value)
    if isinstance(value, float):
        if value.is_integer() and abs(value) < 1e15:
            return str(int(value))
        return f"{value:.6g}"
    return _preview(value)


def _preview(value, limit: int = 160) -> str:
    if isinstance(value, str):
        flat = value.replace("\n", "\\n")
        if len(flat) > limit:
            return f"{flat[:limit]}... ({len(value)} chars)"
        return flat
    if isinstance(value, dict):
        text = json.dumps(value, ensure_ascii=False, default=str, separators=(", ", ": "))
        if len(text) > limit:
            return f"{text[:limit]}... ({len(text)} chars)"
        return text
    if isinstance(value, list):
        text = ", ".join(_preview(item, 40) for item in value[:12])
        if len(value) > 12:
            text += f", ... ({len(value)})"
        return "[" + text + "]"
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, float):
        return format_number(value)
    if value is None:
        return "null"
    return str(value)


def _line(label: str, value) -> None:
    print(f"{label:<16}{value}")


def read_json(path: Path):
    try:
        text = path.read_text(encoding="utf-8")
    except (OSError, UnicodeError) as exc:
        raise ValueError(f"{path}: {exc}") from exc
    try:
        return json.loads(text)
    except json.JSONDecodeError as exc:
        raise ValueError(f"{path}: invalid JSON ({exc.msg})") from exc


def read_header(path: Path) -> tuple[dict, int]:
    with path.open("rb") as handle:
        raw = handle.read(8)
        if len(raw) != 8:
            raise ValueError(f"{path} is not a safetensors file")
        header_size = struct.unpack("<Q", raw)[0]
        if header_size <= 0 or header_size > _MAX_HEADER:
            raise ValueError(f"{path} is not a safetensors file")
        payload = handle.read(header_size)
    if len(payload) != header_size:
        raise ValueError(f"{path} has a truncated safetensors header")
    try:
        header = json.loads(payload)
    except json.JSONDecodeError as exc:
        raise ValueError(f"{path} has an unreadable safetensors header") from exc
    if not isinstance(header, dict):
        raise ValueError(f"{path} has an unreadable safetensors header")
    return header, header_size


def _sections(config: dict) -> list[dict]:
    found = [config]
    for key in _NESTED_CONFIGS:
        nested = config.get(key)
        if isinstance(nested, dict):
            found.append(nested)
    return found


def _present(config: dict, keys: tuple[str, ...]):
    for section in _sections(config):
        for key in keys:
            if key in section and section[key] not in (None, ""):
                return section[key]
    return None


def _positive_int(config: dict, keys: tuple[str, ...], limit: int = 10**12) -> int | None:
    value = _present(config, keys)
    if isinstance(value, bool):
        return None
    if isinstance(value, int) and 0 < value <= limit:
        return value
    if isinstance(value, float) and value.is_integer() and 0 < value <= limit:
        return int(value)
    return None


def _text(value) -> str:
    return value if isinstance(value, str) and value else ""


def chat_template_has_tools(template: str) -> bool:
    if not template:
        return False
    if "tools" not in template and "tool_call" not in template:
        return False
    if "tool_calls" not in template and "assistant_tool_call" not in template:
        return False
    return any(marker in template for marker in _TOOL_RESULT_MARKERS)


def _template_parts(value) -> list[str]:
    if isinstance(value, str):
        return [value] if value.strip() else []
    if not isinstance(value, list):
        return []
    parts = []
    for item in value:
        if isinstance(item, str) and item.strip():
            parts.append(item)
        elif isinstance(item, dict) and isinstance(item.get("template"), str) and item["template"].strip():
            parts.append(item["template"])
    return parts


def load_chat_template(root: Path, tokenizer: dict) -> tuple[str, str, int]:
    parts = _template_parts(tokenizer.get("chat_template")) if tokenizer else []
    if parts:
        count = len(tokenizer.get("chat_template")) if isinstance(tokenizer.get("chat_template"), list) else 1
        return "\n".join(parts), "tokenizer_config.json", count
    path = root / "chat_template.jinja"
    if path.is_file():
        try:
            return path.read_text(encoding="utf-8"), path.name, 1
        except (OSError, UnicodeError) as exc:
            raise ValueError(f"{path}: {exc}") from exc
    return "", "", 0


def _token_text(value) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, dict) and isinstance(value.get("content"), str):
        return value["content"]
    return ""


def token_label(token_id, tokenizer: dict) -> str:
    if isinstance(token_id, list):
        return ", ".join(token_label(item, tokenizer) for item in token_id)
    name = ""
    decoder = tokenizer.get("added_tokens_decoder") if tokenizer else None
    if isinstance(decoder, dict):
        entry = decoder.get(str(token_id))
        if isinstance(entry, dict):
            name = _token_text(entry.get("content"))
    if name:
        return f"{token_id} ({name})"
    return str(token_id)


def _token_id(source: dict, key: str):
    if not source or key not in source:
        return None
    value = source.get(key)
    if isinstance(value, bool) or value is None or value == "":
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, list) and value and all(isinstance(item, int) and not isinstance(item, bool) for item in value):
        return value
    return None


def read_base_model(path: Path) -> str | None:
    if not path.is_file():
        return None
    try:
        text = path.read_text(encoding="utf-8", errors="replace")[:65536]
    except OSError:
        return None
    if not text.startswith("---"):
        return None
    end = text.find("\n---", 3)
    if end < 0:
        return None
    values: list[str] = []
    grabbing = False
    for line in text[3:end].splitlines():
        if not grabbing:
            if line.startswith("base_model:"):
                rest = line.split(":", 1)[1].strip().strip("'\"")
                if rest:
                    values.append(rest)
                else:
                    grabbing = True
            continue
        stripped = line.strip()
        if stripped.startswith("- "):
            values.append(stripped[2:].strip().strip("'\""))
            continue
        break
    cleaned = [value for value in values if value]
    return ", ".join(cleaned) if cleaned else None


def read_recipe_scheme(path: Path) -> str | None:
    if not path.is_file():
        return None
    try:
        text = path.read_text(encoding="utf-8", errors="replace")[:65536]
    except OSError:
        return None
    for line in text.splitlines():
        stripped = line.strip()
        if stripped.startswith("scheme:"):
            value = stripped.split(":", 1)[1].strip().strip("'\"")
            return value or None
    return None


def load_quant(config: dict, root: Path) -> tuple[dict, dict | None]:
    raw = config.get("quantization_config")
    quant = dict(raw) if isinstance(raw, dict) else {}
    path = root / "hf_quant_config.json"
    if not path.is_file():
        return quant, None
    extra = read_json(path)
    if not isinstance(extra, dict):
        raise ValueError(f"{path} is not a JSON object")
    block = extra.get("quantization") if isinstance(extra.get("quantization"), dict) else {}
    producer = extra.get("producer") if isinstance(extra.get("producer"), dict) else {}
    for key, value in block.items():
        quant.setdefault(key, value)
    if "quant_method" not in quant and isinstance(producer.get("name"), str) and producer["name"]:
        quant["quant_method"] = producer["name"].lower()
    return quant, extra


def quant_is_nvfp4(quant: dict) -> bool:
    algo = str(quant.get("quant_algo") or "")
    fmt = str(quant.get("format") or "")
    if algo in {"NVFP4", "W4A16_NVFP4"} or algo.endswith("NVFP4"):
        return True
    if "nvfp4" in fmt.lower():
        return True
    groups = quant.get("config_groups")
    if isinstance(groups, dict):
        for group in groups.values():
            if isinstance(group, dict) and "nvfp4" in str(group.get("format") or "").lower():
                return True
    layers = quant.get("quantized_layers")
    if isinstance(layers, dict):
        for entry in layers.values():
            if isinstance(entry, dict) and str(entry.get("quant_algo") or "").endswith("NVFP4"):
                return True
    return False


def _weight_schemes(quant: dict) -> list[tuple]:
    found = []
    groups = quant.get("config_groups")
    if isinstance(groups, dict):
        for group in groups.values():
            if not isinstance(group, dict):
                continue
            weights = group.get("weights") if isinstance(group.get("weights"), dict) else {}
            acts = group.get("input_activations") if isinstance(group.get("input_activations"), dict) else {}
            found.append((weights.get("num_bits"), weights.get("group_size"), acts.get("num_bits")))
    if found:
        return found
    bits = quant.get("bits", quant.get("num_bits", quant.get("w_bit")))
    group_size = quant.get("group_size", quant.get("groupsize"))
    if bits is not None or group_size is not None:
        return [(bits, group_size, None)]
    return []


def find_index(root: Path) -> Path | None:
    for name in _INDEX_NAMES:
        path = root / name
        if path.is_file():
            return path
    found = sorted(path for path in root.glob("*.safetensors.index.json") if path.is_file())
    return found[0] if found else None


def shards_from_index(index_path: Path, payload: dict) -> list[Path]:
    weight_map = payload.get("weight_map")
    if not isinstance(weight_map, dict) or not weight_map:
        raise ValueError(f"{index_path} has no weight_map")
    root = index_path.parent
    names: list[Path] = []
    seen: set[str] = set()
    for name in weight_map.values():
        if not isinstance(name, str) or not name or name.startswith(("/", "\\")) or ".." in Path(name).parts:
            raise ValueError(f"{index_path} has an invalid weight_map entry")
        if name not in seen:
            seen.add(name)
            names.append(root / name)
    return names


def discover_shards(root: Path) -> list[Path]:
    for name in ("model.safetensors", "diffusion_pytorch_model.safetensors"):
        path = root / name
        if path.is_file():
            return [path]
    return sorted(path for path in root.glob("*.safetensors") if path.is_file())


def parse_tensor(name: str, entry, file_name: str) -> tuple[TensorInfo, int]:
    if not isinstance(entry, dict):
        raise ValueError(f"{file_name}: tensor {name!r} has an invalid header entry")
    dtype = entry.get("dtype")
    if not isinstance(dtype, str) or not dtype:
        dtype = "?"
    shape = entry.get("shape")
    if shape is None:
        dims: list[int] = []
    elif isinstance(shape, list) and all(isinstance(dim, int) and not isinstance(dim, bool) for dim in shape):
        dims = [int(dim) for dim in shape]
    else:
        raise ValueError(f"{file_name}: tensor {name!r} has an invalid shape")
    offsets = entry.get("data_offsets")
    nbytes = None
    end = 0
    if isinstance(offsets, list) and len(offsets) == 2 and all(isinstance(item, int) and not isinstance(item, bool) for item in offsets):
        start, end = int(offsets[0]), int(offsets[1])
        if start < 0 or end < start:
            raise ValueError(f"{file_name}: tensor {name!r} has invalid data offsets")
        nbytes = end - start
    return TensorInfo(name, dtype, dims, nbytes, file_name), end


def load_shards(paths: list[Path]) -> tuple[list[TensorInfo], list[ShardInfo]]:
    tensors: list[TensorInfo] = []
    shards: list[ShardInfo] = []
    owners: dict[str, str] = {}
    for path in paths:
        header, header_size = read_header(path)
        meta = header.pop("__metadata__", None)
        if not isinstance(meta, dict):
            meta = {}
        data_end = 0
        for name, entry in header.items():
            if name.startswith("__"):
                continue
            if name in owners:
                raise ValueError(f"Tensor {name!r} appears in both {owners[name]} and {path.name}")
            owners[name] = path.name
            tensor, end = parse_tensor(name, entry, path.name)
            data_end = max(data_end, end)
            tensors.append(tensor)
        shards.append(ShardInfo(path, header_size, data_end, meta))
    return tensors, shards


def _natural_key(name: str) -> tuple:
    return tuple((0, int(part)) if part.isdigit() else (1, part) for part in re.split(r"(\d+)", name))


def tensor_type(tensor: TensorInfo, quant_method: str, nvfp4: bool, comfy: bool) -> str:
    leaf = tensor.name.rsplit(".", 1)[-1]
    if nvfp4 and leaf == "weight_packed":
        return "NVFP4"
    if comfy and leaf == "weight" and tensor.dtype == "U8":
        return "NVFP4"
    if quant_method.lower() in {"gptq", "awq"} and leaf == "qweight":
        return quant_method.upper()
    return tensor.dtype


def _numel(shape: list[int]) -> int:
    count = 1
    for dim in shape:
        count *= dim
    return count


def parameter_count(tensors: list[TensorInfo], type_of) -> int | None:
    """Logical weights and biases.

    Packed NVFP4 stores two FP4 values in each byte, so those element counts
    are doubled. Scale tensors are not parameters. Other packed layouts are
    left uncounted rather than reported as a partial total.
    """
    total = 0
    seen = False
    for tensor in tensors:
        leaf = tensor.name.rsplit(".", 1)[-1]
        packed = leaf in {"weight_packed", "qweight"}
        if packed and type_of(tensor) != "NVFP4":
            return None
        if leaf not in {"weight", "bias", "weight_packed"}:
            continue
        seen = True
        count = _numel(tensor.shape)
        if type_of(tensor) == "NVFP4":
            count *= 2
        total += count
    return total if seen else None


def _is_model_dir(path: Path) -> bool:
    return path.name in {"config.json", "adapter_config.json"} or path.name.endswith(".safetensors.index.json")


def _load(path: Path) -> Checkpoint:
    if not path.exists():
        raise SystemExit(f"Not found: {path}")
    if path.is_file() and path.suffix == ".gguf":
        raise SystemExit(f"{path} is a GGUF file. Use gguf_info.py for GGUF files.")
    if path.is_file() and _is_model_dir(path):
        root = path.parent
        selected: list[Path] | None = None
    elif path.is_file():
        if path.suffix != ".safetensors":
            raise SystemExit(f"{path} is not a model directory or a safetensors file")
        root = path.parent
        selected = [path]
    elif path.is_dir():
        root = path
        selected = None
    else:
        raise SystemExit(f"Not a file or directory: {path}")

    config_path = root / "config.json"
    config_name = "config.json"
    if not config_path.is_file() and (root / "adapter_config.json").is_file():
        config_path = root / "adapter_config.json"
        config_name = config_path.name
    config: dict = {}
    if config_path.is_file():
        loaded = read_json(config_path)
        if not isinstance(loaded, dict):
            raise ValueError(f"{config_path} is not a JSON object")
        config = loaded

    generation_path = root / "generation_config.json"
    generation: dict = {}
    if generation_path.is_file():
        loaded = read_json(generation_path)
        if isinstance(loaded, dict):
            generation = loaded

    tokenizer_path = root / "tokenizer_config.json"
    tokenizer: dict = {}
    if tokenizer_path.is_file():
        loaded = read_json(tokenizer_path)
        if isinstance(loaded, dict):
            tokenizer = loaded

    quant, hf_quant = load_quant(config, root)
    recipe = read_recipe_scheme(root / "recipe.yaml")
    base_model = read_base_model(root / "README.md")
    repo = config.get("_name_or_path")
    if not base_model and isinstance(repo, str) and "/" in repo and not repo.startswith(("/", ".")):
        base_model = repo

    index_path = find_index(root)
    index_payload: dict = {}
    weight_map: dict = {}
    indexed: list[Path] = []
    if index_path is not None:
        loaded = read_json(index_path)
        if not isinstance(loaded, dict):
            raise ValueError(f"{index_path} is not a JSON object")
        index_payload = loaded
        raw_map = loaded.get("weight_map")
        weight_map = raw_map if isinstance(raw_map, dict) else {}
        indexed = shards_from_index(index_path, loaded)
    if selected is None:
        shard_paths = indexed if indexed else discover_shards(root)
        missing = [path.name for path in shard_paths if not path.is_file()]
        if missing:
            raise ValueError(f"Index entries are missing from {root}: {', '.join(missing)}")
    else:
        shard_paths = selected

    if not config and not shard_paths:
        raise SystemExit(f"No config.json or safetensors files in {root}")

    tensors, shards = load_shards(shard_paths) if shard_paths else ([], [])
    opened = {shard.path.name for shard in shards}
    index_names = {path.name for path in indexed}
    partial = bool(index_names) and opened != index_names and opened.issubset(index_names)
    if index_names and opened == index_names:
        missing = [name for name in weight_map if isinstance(name, str) and name not in {tensor.name for tensor in tensors}]
        if missing:
            raise ValueError(f"Index lists {missing[0]!r}, which is not in the safetensors headers")

    comfy = any(tensor.name.endswith(".comfy_quant") for tensor in tensors) or any(
        "_quantization_metadata" in shard.metadata for shard in shards
    )
    packed = any(tensor.name.endswith("weight_packed") for tensor in tensors)
    recipe_nvfp4 = bool(recipe) and "NVFP4" in recipe.upper()
    nvfp4 = quant_is_nvfp4(quant) or comfy or (recipe_nvfp4 and packed)
    template, template_where, template_count = load_chat_template(root, tokenizer)
    index_meta = index_payload.get("metadata") if isinstance(index_payload.get("metadata"), dict) else {}

    return Checkpoint(
        path=path,
        root=root,
        config=config,
        config_name=config_name,
        generation=generation,
        tokenizer=tokenizer,
        quant=quant,
        hf_quant=hf_quant,
        recipe=recipe,
        base_model=base_model,
        tensors=tensors,
        shards=shards,
        shard_total=len(indexed) if indexed else len(shards),
        partial=partial,
        index_meta=index_meta,
        template=template,
        template_where=template_where,
        template_count=template_count,
        nvfp4=nvfp4,
        comfy=comfy,
    )


def load_checkpoint(path: Path) -> Checkpoint:
    try:
        return _load(path)
    except (OSError, ValueError, struct.error, UnicodeError) as exc:
        raise SystemExit(str(exc)) from exc


def _type_of(checkpoint: Checkpoint):
    method = str(checkpoint.quant.get("quant_method") or "")

    def kind(tensor: TensorInfo) -> str:
        return tensor_type(tensor, method, checkpoint.nvfp4, checkpoint.comfy)

    return kind


def _quant_title(checkpoint: Checkpoint) -> str:
    quant = checkpoint.quant
    algo = quant.get("quant_algo")
    if isinstance(algo, str) and algo:
        return algo
    if checkpoint.nvfp4:
        return "NVFP4"
    fmt = quant.get("format")
    if isinstance(fmt, str) and fmt:
        return fmt
    method = quant.get("quant_method")
    if isinstance(method, str) and method:
        return method
    if checkpoint.recipe:
        return checkpoint.recipe
    return ""


def print_identity(checkpoint: Checkpoint) -> None:
    _line("file", checkpoint.path)
    if checkpoint.shards:
        size = sum(shard.path.stat().st_size for shard in checkpoint.shards)
        _line("size", format_size(size))
    model_type = _text(_present(checkpoint.config, ("model_type",)))
    architectures = checkpoint.config.get("architectures")
    if isinstance(architectures, list):
        arch_text = ", ".join(str(item) for item in architectures if item not in (None, ""))
    else:
        arch_text = _text(architectures)
    if model_type:
        _line("architecture", model_type)
        if arch_text and arch_text != model_type:
            _line("class", arch_text)
    elif arch_text:
        _line("architecture", arch_text)
    if checkpoint.path.is_dir() or checkpoint.config:
        _line("name", checkpoint.root.name)
    else:
        _line("name", checkpoint.path.stem)
    if checkpoint.base_model and checkpoint.base_model != checkpoint.root.name:
        _line("base model", checkpoint.base_model)
    dtype = _text(_present(checkpoint.config, ("dtype", "torch_dtype")))
    if dtype:
        _line("dtype", dtype)


def print_quantization(checkpoint: Checkpoint) -> None:
    title = _quant_title(checkpoint)
    if checkpoint.nvfp4 or checkpoint.quant:
        if title:
            _line("quantization", title)
    elif checkpoint.recipe:
        _line("recipe", checkpoint.recipe)
        return
    else:
        return
    method = _text(checkpoint.quant.get("quant_method"))
    fmt = _text(checkpoint.quant.get("format"))
    if method and method != title:
        _line("quant method", method)
    if fmt and fmt != title:
        _line("quant format", fmt)
    elif checkpoint.comfy and not fmt:
        _line("quant format", "comfy")
    schemes = _weight_schemes(checkpoint.quant)
    bits = []
    groups = []
    act_bits = []
    for weight_bits, group_size, activation_bits in schemes:
        if isinstance(weight_bits, int) and not isinstance(weight_bits, bool) and weight_bits not in bits:
            bits.append(weight_bits)
        if isinstance(group_size, int) and not isinstance(group_size, bool) and group_size > 0 and group_size not in groups:
            groups.append(group_size)
        if isinstance(activation_bits, int) and not isinstance(activation_bits, bool) and activation_bits not in act_bits:
            act_bits.append(activation_bits)
    if bits:
        _line("bits", ", ".join(str(item) for item in bits))
    if groups:
        _line("group size", ", ".join(str(item) for item in groups))
    if act_bits:
        _line("act bits", ", ".join(str(item) for item in act_bits))
    status = _text(checkpoint.quant.get("quantization_status"))
    if status:
        _line("quant status", status)
    version = checkpoint.quant.get("version")
    if isinstance(version, (str, int, float)) and not isinstance(version, bool) and version != "":
        _line("quant version", format_number(version) if isinstance(version, (int, float)) else version)
    ignored = checkpoint.quant.get("ignore")
    if isinstance(ignored, list) and ignored:
        if all(isinstance(item, str) for item in ignored):
            text = ", ".join(ignored)
            if len(text) > 160:
                text = text[:160] + "..."
            _line("quant ignore", text)
        else:
            _line("quant ignore", _preview(ignored))
    elif isinstance(ignored, str) and ignored:
        _line("quant ignore", ignored)
    layers = _comfy_layers(checkpoint)
    if layers:
        _line("quant layers", layers)


def _comfy_layers(checkpoint: Checkpoint) -> int:
    total = 0
    found = False
    for shard in checkpoint.shards:
        raw = shard.metadata.get("_quantization_metadata")
        parsed = None
        if isinstance(raw, str):
            try:
                parsed = json.loads(raw)
            except json.JSONDecodeError:
                parsed = None
        elif isinstance(raw, dict):
            parsed = raw
        if not isinstance(parsed, dict):
            continue
        layers = parsed.get("layers")
        if isinstance(layers, dict):
            total += len(layers)
            found = True
    return total if found else 0


def _submodule_text(block: dict) -> str:
    parts = []
    for label, keys in (
        ("layers", ("num_hidden_layers", "num_layers", "depth")),
        ("width", ("hidden_size", "d_model", "embed_dim")),
        ("heads", ("num_attention_heads", "num_heads")),
    ):
        for key in keys:
            value = block.get(key)
            if isinstance(value, int) and not isinstance(value, bool) and value > 0:
                parts.append(f"{value} {label}")
                break
    return ", ".join(parts)


def print_architecture(checkpoint: Checkpoint) -> None:
    config = checkpoint.config
    context = _positive_int(config, (
        "max_position_embeddings",
        "max_seq_len",
        "n_ctx",
        "n_positions",
        "max_sequence_length",
        "seq_length",
        "model_max_length",
    ))
    if context:
        _line("context", context)
    tok_length = _positive_int(checkpoint.tokenizer, ("model_max_length",))
    if tok_length and tok_length != context:
        _line("tok length", tok_length)
    fields = (
        ("blocks", ("num_hidden_layers", "n_layer", "num_layers", "n_layers")),
        ("embedding", ("hidden_size", "n_embd", "d_model", "hidden_dim")),
        ("feed forward", ("intermediate_size", "n_inner", "ffn_dim", "intermediate_dim")),
        ("heads", ("num_attention_heads", "n_head", "num_heads")),
        ("kv heads", ("num_key_value_heads", "num_kv_heads", "n_head_kv")),
        ("experts", ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")),
        ("experts used", ("num_experts_per_tok", "num_experts_per_token")),
        ("expert ffn", ("moe_intermediate_size", "expert_intermediate_size")),
        ("vocab", ("vocab_size", "padded_vocab_size")),
    )
    values = {label: _positive_int(config, keys) for label, keys in fields}
    head_dim = _positive_int(config, ("head_dim", "attention_head_dim", "head_size"))
    if head_dim is None and values["embedding"] and values["heads"] and values["embedding"] % values["heads"] == 0:
        if _present(config, ("qk_nope_head_dim", "qk_rope_head_dim")) is None:
            head_dim = values["embedding"] // values["heads"]
    for label, _keys in fields:
        if values[label]:
            _line(label, values[label])
        if label == "kv heads" and head_dim:
            _line("head dim", head_dim)
    activation = _text(_present(config, ("hidden_act", "activation_function", "hidden_activation")))
    if activation:
        _line("activation", activation)
    rope = _present(config, ("rope_theta",))
    if isinstance(rope, (int, float)) and not isinstance(rope, bool) and rope != 0:
        _line("rope theta", format_number(rope))
    window = _positive_int(config, ("sliding_window", "sliding_window_size"))
    if window and _present(config, ("use_sliding_window",)) is not False:
        _line("sliding window", window)
    layer_types = _present(config, ("layer_types", "sliding_window_pattern"))
    if isinstance(layer_types, list) and layer_types:
        if all(isinstance(item, bool) for item in layer_types):
            sliding = sum(bool(item) for item in layer_types)
            _line("swa layers", f"{sliding} sliding, {len(layer_types) - sliding} full")
        elif all(isinstance(item, str) for item in layer_types):
            counts = Counter(layer_types)
            if not (len(counts) == 1 and next(iter(counts)) in {"full_attention", "full"}):
                text = ", ".join(f"{count} {name}" for name, count in counts.most_common())
                _line("layers", text)
    for key, label in (("vision_config", "vision"), ("audio_config", "audio"), ("vision_tower", "vision")):
        block = config.get(key)
        if isinstance(block, dict):
            text = _submodule_text(block)
            if text:
                _line(label, text)
        elif isinstance(block, str) and block:
            _line(label, block)
    if _present(config, ("tie_word_embeddings", "tie_embeddings")) is True:
        _line("tied embed", "yes")


def print_tokenizer(checkpoint: Checkpoint) -> None:
    tokenizer = checkpoint.tokenizer
    kind = _text(tokenizer.get("tokenizer_class"))
    if not kind and (checkpoint.root / "tokenizer.model").is_file():
        kind = "sentencepiece"
    elif not kind and (checkpoint.root / "tokenizer.json").is_file():
        kind = "tokenizer.json"
    if kind:
        _line("tokenizer", kind)
    bos = _token_id(checkpoint.config, "bos_token_id")
    if bos is None:
        bos = _token_id(checkpoint.generation, "bos_token_id")
    eos = _token_id(checkpoint.generation, "eos_token_id")
    if eos is None:
        eos = _token_id(checkpoint.config, "eos_token_id")
    pad = _token_id(checkpoint.generation, "pad_token_id")
    if pad is None:
        pad = _token_id(checkpoint.config, "pad_token_id")
    if bos is not None:
        _line("bos", token_label(bos, tokenizer))
    elif _token_text(tokenizer.get("bos_token")):
        _line("bos", _token_text(tokenizer.get("bos_token")))
    if eos is not None:
        _line("eos", token_label(eos, tokenizer))
    elif _token_text(tokenizer.get("eos_token")):
        _line("eos", _token_text(tokenizer.get("eos_token")))
    if pad is not None:
        _line("pad", token_label(pad, tokenizer))
    elif _token_text(tokenizer.get("pad_token")):
        _line("pad", _token_text(tokenizer.get("pad_token")))
    if checkpoint.template:
        detail = f"{len(checkpoint.template)} chars"
        if checkpoint.template_count > 1:
            detail = f"{checkpoint.template_count} templates, {detail}"
        if checkpoint.template_where:
            detail = f"{detail}, {checkpoint.template_where}"
        if chat_template_has_tools(checkpoint.template):
            detail = f"{detail}, tools"
        _line("chat template", f"yes ({detail})")
    else:
        _line("chat template", "no")


def print_tensor_summary(checkpoint: Checkpoint) -> None:
    if checkpoint.partial:
        _line("shards", f"{len(checkpoint.shards)} of {checkpoint.shard_total}")
    elif len(checkpoint.shards) > 1:
        _line("shards", str(len(checkpoint.shards)))
    truncated = [
        shard.path.name
        for shard in checkpoint.shards
        if shard.path.stat().st_size < 8 + shard.header_size + shard.data_end
    ]
    if truncated:
        shown = ", ".join(truncated[:4])
        if len(truncated) > 4:
            shown += f", ... ({len(truncated)})"
        _line("truncated", shown)
    kind = _type_of(checkpoint)
    params = None if checkpoint.partial else parameter_count(checkpoint.tensors, kind)
    if params:
        _line("parameters", format_params(params))
    _line("tensors", str(len(checkpoint.tensors)))
    if not checkpoint.tensors:
        return
    counts: Counter[str] = Counter()
    stored: dict[str, Counter[str]] = {}
    known = 0
    unknown = 0
    for tensor in checkpoint.tensors:
        name = kind(tensor)
        counts[name] += 1
        stored.setdefault(name, Counter())[tensor.dtype] += 1
        if tensor.nbytes is None:
            unknown += 1
        else:
            known += tensor.nbytes
    width = max([10, *(len(name) for name in counts)])
    for name, count in counts.most_common():
        extra = ""
        dtypes = stored[name]
        if list(dtypes) != [name]:
            if len(dtypes) == 1:
                extra = f"  (stored {next(iter(dtypes))})"
            else:
                detail = ", ".join(f"{dtype} x {number}" for dtype, number in dtypes.most_common())
                extra = f"  (stored {detail})"
        print(f"  {name:<{width}} {count}{extra}")
    _line("tensor bytes", format_size(known))
    if unknown:
        _line("unmeasured", f"{unknown} tensors have no data length")


def print_summary(checkpoint: Checkpoint) -> None:
    print_identity(checkpoint)
    print_quantization(checkpoint)
    print_architecture(checkpoint)
    print_tokenizer(checkpoint)
    print_tensor_summary(checkpoint)


def _metadata_tokenizer(tokenizer: dict) -> dict:
    view = {}
    for key, value in tokenizer.items():
        if key == "added_tokens_decoder" and isinstance(value, dict):
            view[key] = f"{len(value)} added tokens"
        else:
            view[key] = value
    return view


def _metadata_safetensors(metadata: dict) -> dict:
    view = {}
    for key, value in metadata.items():
        if key == "format" and value == "pt":
            continue
        if key == "_quantization_metadata":
            parsed = None
            if isinstance(value, str):
                try:
                    parsed = json.loads(value)
                except json.JSONDecodeError:
                    parsed = None
            elif isinstance(value, dict):
                parsed = value
            layers = parsed.get("layers") if isinstance(parsed, dict) else None
            if isinstance(parsed, dict) and isinstance(layers, dict):
                version = parsed.get("format_version", "?")
                view[key] = f"format {version}, {len(layers)} layers"
                continue
        view[key] = value
    return view


def _print_mapping(title: str, mapping: dict) -> None:
    if not mapping:
        return
    print()
    print(title)
    width = max(len(str(key)) for key in mapping)
    for key, value in mapping.items():
        print(f"  {key:<{width}}  {_preview(value)}")


def print_metadata(checkpoint: Checkpoint) -> None:
    if checkpoint.config:
        _print_mapping(checkpoint.config_name, checkpoint.config)
    if checkpoint.hf_quant:
        _print_mapping("hf_quant_config.json", checkpoint.hf_quant)
    if checkpoint.generation:
        _print_mapping("generation_config.json", checkpoint.generation)
    if checkpoint.tokenizer:
        _print_mapping("tokenizer_config.json", _metadata_tokenizer(checkpoint.tokenizer))
    if checkpoint.template and not _template_parts(checkpoint.tokenizer.get("chat_template") if checkpoint.tokenizer else None):
        _print_mapping(checkpoint.template_where or "chat_template", {"chat_template": checkpoint.template})
    if checkpoint.index_meta:
        _print_mapping("safetensors index", checkpoint.index_meta)
    metas = []
    for shard in checkpoint.shards:
        cleaned = _metadata_safetensors(shard.metadata)
        if cleaned and cleaned not in metas:
            metas.append(cleaned)
            title = "safetensors metadata" if len(checkpoint.shards) == 1 else f"safetensors metadata ({shard.path.name})"
            if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards):
                title = "safetensors metadata"
            _print_mapping(title, cleaned)
            if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards):
                break


def _shape_text(tensor: TensorInfo) -> str:
    return " x ".join(str(dim) for dim in tensor.shape) or "-"


def print_tensors(checkpoint: Checkpoint) -> None:
    print()
    if not checkpoint.tensors:
        print("tensors")
        print("  (none)")
        return
    kind = _type_of(checkpoint)
    tensors = sorted(checkpoint.tensors, key=lambda tensor: _natural_key(tensor.name))
    if any(kind(tensor) == "NVFP4" for tensor in tensors):
        print("NVFP4 shapes are packed storage, 2 FP4 values per byte.")
    name_width = max(len(tensor.name) for tensor in tensors)
    type_width = max(8, max(len(kind(tensor)) for tensor in tensors))
    shapes = [_shape_text(tensor) for tensor in tensors]
    show_shard = len({tensor.file for tensor in tensors}) > 1
    shape_width = max(len("SHAPE"), *(len(shape) for shape in shapes)) if show_shard else 0
    shape_header = f"{'SHAPE':<{shape_width}}" if show_shard else "SHAPE"
    header = f"{'NAME':<{name_width + 2}}  {'TYPE':<{type_width}}  {shape_header}"
    if show_shard:
        header += "  SHARD"
    print(header)
    for tensor, shape in zip(tensors, shapes):
        shown = f"{shape:<{shape_width}}" if show_shard else shape
        row = f"  {tensor.name:<{name_width}}  {kind(tensor):<{type_width}}  {shown}"
        if show_shard:
            row += f"  {tensor.file}"
        print(row)


def inspect(path: Path, metadata_dump: bool, tensor_dump: bool) -> None:
    checkpoint = load_checkpoint(path)
    print_summary(checkpoint)
    if metadata_dump:
        print_metadata(checkpoint)
    if tensor_dump:
        print_tensors(checkpoint)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Print architecture, quantization, and tensor info from a Hugging Face checkpoint."
    )
    parser.add_argument("model", nargs="+", type=Path, help="A model directory or a .safetensors file")
    parser.add_argument(
        "--metadata",
        action="store_true",
        help="Print config, tokenizer, and safetensors header metadata. Long values are shortened.",
    )
    parser.add_argument("--tensors", action="store_true", help="Print every tensor name, type, and shape.")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    for index, path in enumerate(args.model):
        if index:
            print()
        inspect(path, args.metadata, args.tensors)


if __name__ == "__main__":
    try:
        main()
    except BrokenPipeError:
        # head and less close the pipe early. Don't turn that into a traceback.
        try:
            sys.stdout.close()
        except BrokenPipeError:
            pass
        try:
            sys.stderr.close()
        except BrokenPipeError:
            pass
        os._exit(0)