#!/usr/bin/env python3
"""Print a summary of a GGUF file from its header.

The weight bytes are not read. Metadata and the tensor directory are enough
for the architecture, quantization, tokenizer, and tensor list.
"""

from __future__ import annotations

import argparse
import struct
from collections import Counter
from pathlib import Path

# GGML type id -> name, and (values per block, bytes per block).
GGML_TYPES = {
    0: ("F32", 1, 4),
    1: ("F16", 1, 2),
    2: ("Q4_0", 32, 18),
    3: ("Q4_1", 32, 20),
    6: ("Q5_0", 32, 22),
    7: ("Q5_1", 32, 24),
    8: ("Q8_0", 32, 34),
    9: ("Q8_1", 32, 36),
    10: ("Q2_K", 256, 84),
    11: ("Q3_K", 256, 110),
    12: ("Q4_K", 256, 144),
    13: ("Q5_K", 256, 176),
    14: ("Q6_K", 256, 210),
    15: ("Q8_K", 256, 292),
    16: ("IQ2_XXS", 256, 66),
    17: ("IQ2_XS", 256, 74),
    18: ("IQ3_XXS", 256, 98),
    19: ("IQ1_S", 256, 50),
    20: ("IQ4_NL", 32, 18),
    21: ("IQ3_S", 256, 110),
    22: ("IQ2_S", 256, 82),
    23: ("IQ4_XS", 256, 136),
    24: ("I8", 1, 1),
    25: ("I16", 1, 2),
    26: ("I32", 1, 4),
    27: ("I64", 1, 8),
    28: ("F64", 1, 8),
    29: ("IQ1_M", 256, 56),
    30: ("BF16", 1, 2),
    34: ("TQ1_0", 256, 54),
    35: ("TQ2_0", 256, 66),
    39: ("MXFP4", 32, 17),
    40: ("NVFP4", 64, 36),
    41: ("Q1_0", 128, 18),
    42: ("Q2_0", 64, 18),
}

FILE_TYPES = {
    0: "ALL_F32",
    1: "MOSTLY_F16",
    2: "MOSTLY_Q4_0",
    3: "MOSTLY_Q4_1",
    7: "MOSTLY_Q8_0",
    8: "MOSTLY_Q5_0",
    9: "MOSTLY_Q5_1",
    10: "MOSTLY_Q2_K",
    11: "MOSTLY_Q3_K_S",
    12: "MOSTLY_Q3_K_M",
    13: "MOSTLY_Q3_K_L",
    14: "MOSTLY_Q4_K_S",
    15: "MOSTLY_Q4_K_M",
    16: "MOSTLY_Q5_K_S",
    17: "MOSTLY_Q5_K_M",
    18: "MOSTLY_Q6_K",
    19: "MOSTLY_IQ2_XXS",
    20: "MOSTLY_IQ2_XS",
    21: "MOSTLY_Q2_K_S",
    22: "MOSTLY_IQ3_XS",
    23: "MOSTLY_IQ3_XXS",
    24: "MOSTLY_IQ1_S",
    25: "MOSTLY_IQ4_NL",
    26: "MOSTLY_IQ3_S",
    27: "MOSTLY_IQ3_M",
    28: "MOSTLY_IQ2_S",
    29: "MOSTLY_IQ2_M",
    30: "MOSTLY_IQ4_XS",
    31: "MOSTLY_IQ1_M",
    32: "MOSTLY_BF16",
    36: "MOSTLY_TQ1_0",
    37: "MOSTLY_TQ2_0",
    38: "MOSTLY_MXFP4_MOE",
    39: "MOSTLY_NVFP4",
    40: "MOSTLY_Q1_0",
    41: "MOSTLY_Q2_0",
    1024: "GUESSED",
}

_SCALAR = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
_INTEGER = {0: "<B", 1: "<b", 2: "<H", 3: "<h", 4: "<I", 5: "<i", 7: "<?", 10: "<Q", 11: "<q"}
_FLOAT = {6: "<f", 12: "<d"}


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 _read(handle, n: int) -> bytes:
    data = handle.read(n)
    if len(data) != n:
        raise ValueError("truncated GGUF header")
    return data


def _read_string(handle) -> bytes:
    size = struct.unpack("<Q", _read(handle, 8))[0]
    if size > 64 * 1024 * 1024:
        raise ValueError("GGUF string is too long")
    return _read(handle, size)


def _read_scalar(handle, value_type: int):
    if value_type == 8:
        return _read_string(handle).decode("utf-8", "replace")
    if value_type not in _SCALAR:
        raise ValueError(f"unknown GGUF value type {value_type}")
    raw = _read(handle, _SCALAR[value_type])
    if value_type in _INTEGER:
        return struct.unpack(_INTEGER[value_type], raw)[0]
    if value_type in _FLOAT:
        return struct.unpack(_FLOAT[value_type], raw)[0]
    return None


def _read_array(handle, keep: bool):
    subtype = struct.unpack("<I", _read(handle, 4))[0]
    count = struct.unpack("<Q", _read(handle, 8))[0]
    if subtype in _SCALAR and not keep:
        handle.seek(_SCALAR[subtype] * count, 1)
        return f"{_type_name(subtype)}[{count}]"
    if subtype == 8 and not keep:
        for _ in range(count):
            _read_string(handle)
        return f"string[{count}]"
    values = []
    limit = 8 if subtype == 8 else 16
    for index in range(count):
        if subtype == 8:
            text = _read_string(handle).decode("utf-8", "replace")
            if index < limit:
                values.append(text)
            continue
        if subtype not in _SCALAR:
            raise ValueError(f"unsupported GGUF array type {subtype}")
        item = _read_scalar(handle, subtype)
        if index < limit:
            values.append(item)
    if count > limit:
        values.append(f"... {count} items")
    return values


def _type_name(value_type: int) -> str:
    names = {
        0: "u8", 1: "i8", 2: "u16", 3: "i16", 4: "u32", 5: "i32",
        6: "f32", 7: "bool", 8: "string", 10: "u64", 11: "i64", 12: "f64",
    }
    return names.get(value_type, str(value_type))


def load_gguf(path: Path) -> tuple[dict, list[dict], int]:
    metadata: dict = {}
    tensors: list[dict] = []
    with path.open("rb") as handle:
        magic, version = struct.unpack("<II", _read(handle, 8))
        if magic != 0x46554747:
            raise ValueError(f"{path} is not a GGUF file")
        if version < 2:
            raise ValueError(f"{path} uses GGUF version {version}, which this reader does not support")
        tensor_count, kv_count = struct.unpack("<QQ", _read(handle, 16))
        for _ in range(kv_count):
            key = _read_string(handle).decode("utf-8", "replace")
            value_type = struct.unpack("<I", _read(handle, 4))[0]
            if value_type == 9:
                keep = "token" not in key and "merges" not in key and not key.endswith(".scores")
                metadata[key] = _read_array(handle, keep)
            else:
                metadata[key] = _read_scalar(handle, value_type)
        for _ in range(tensor_count):
            name = _read_string(handle).decode("utf-8", "replace")
            n_dims = struct.unpack("<I", _read(handle, 4))[0]
            dims = list(struct.unpack("<" + "Q" * n_dims, _read(handle, 8 * n_dims))) if n_dims else []
            ggml_type = struct.unpack("<I", _read(handle, 4))[0]
            offset = struct.unpack("<Q", _read(handle, 8))[0]
            tensors.append({"name": name, "dims": dims, "type": ggml_type, "offset": offset})
    return metadata, tensors, version


def _meta(metadata: dict, suffix: str):
    if suffix in metadata:
        return metadata[suffix]
    for key, value in metadata.items():
        if key.endswith("." + suffix) and ".rope." not in key:
            return value
    return None


def _context(metadata: dict) -> str:
    for key, value in metadata.items():
        if key.endswith(".context_length") and ".rope." not in key and isinstance(value, int) and value > 0:
            return str(value)
    return ""


def _tensor_bytes(tensor: dict) -> int | None:
    spec = GGML_TYPES.get(tensor["type"])
    if spec is None or not tensor["dims"]:
        return None
    _name, block, type_size = spec
    elements = 1
    for dim in tensor["dims"]:
        elements *= dim
    if elements % block != 0:
        return None
    return elements // block * type_size


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, 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 f"{value:.6g}"
    return str(value)


def _filetype(value) -> str:
    if not isinstance(value, int):
        return _preview(value)
    name = FILE_TYPES.get(value)
    return f"{name} ({value})" if name else str(value)


def print_summary(path: Path, metadata: dict, tensors: list[dict], version: int) -> None:
    arch = str(metadata.get("general.architecture") or "")
    print(f"file          {path}")
    print(f"size          {format_size(path.stat().st_size)}")
    print(f"gguf          version {version}")
    print(f"architecture  {arch or '-'}")
    for label, key in (
        ("name", "general.name"),
        ("basename", "general.basename"),
        ("size label", "general.size_label"),
        ("finetune", "general.finetune"),
        ("type", "general.type"),
    ):
        if key in metadata and metadata[key] not in ("", None):
            print(f"{label:<14}{_preview(metadata[key])}")
    if "general.file_type" in metadata:
        print(f"filetype      {_filetype(metadata['general.file_type'])}")
    if "general.quantization_version" in metadata:
        print(f"quant version {metadata['general.quantization_version']}")
    context = _context(metadata)
    if context:
        print(f"context       {context}")
    fields = (
        ("blocks", "block_count"),
        ("embedding", "embedding_length"),
        ("feed forward", "feed_forward_length"),
        ("heads", "attention.head_count"),
        ("kv heads", "attention.head_count_kv"),
        ("head dim", "attention.key_length"),
        ("experts", "expert_count"),
        ("vocab", "vocab_size"),
        ("sliding window", "attention.sliding_window"),
    )
    for label, suffix in fields:
        value = _meta(metadata, suffix)
        if isinstance(value, int) and value > 0:
            print(f"{label:<14}{value}")
    swa = _meta(metadata, "attention.sliding_window_pattern")
    if isinstance(swa, list) and swa and all(isinstance(item, (bool, int)) for item in swa):
        flags = [bool(item) for item in swa]
        print(f"swa layers    {sum(flags)} sliding, {len(flags) - sum(flags)} full")
    tok_model = metadata.get("tokenizer.ggml.model")
    if tok_model:
        print(f"tokenizer     {tok_model}")
    for label, key in (
        ("bos", "tokenizer.ggml.bos_token_id"),
        ("eos", "tokenizer.ggml.eos_token_id"),
    ):
        if key in metadata:
            print(f"{label:<14}{metadata[key]}")
    template = metadata.get("tokenizer.chat_template")
    if isinstance(template, str) and template:
        print(f"chat template yes ({len(template)} chars)")
    elif isinstance(template, list) and template:
        print(f"chat template {len(template)} templates")
    else:
        print("chat template no")
    if isinstance(metadata.get("general.description"), str):
        print(f"description   {_preview(metadata['general.description'], 200)}")

    counts: Counter[str] = Counter()
    known_bytes = 0
    unknown = 0
    for tensor in tensors:
        spec = GGML_TYPES.get(tensor["type"])
        counts[spec[0] if spec else f"type {tensor['type']}"] += 1
        nbytes = _tensor_bytes(tensor)
        if nbytes is None:
            unknown += 1
        else:
            known_bytes += nbytes
    print(f"tensors       {len(tensors)}")
    for name, count in counts.most_common():
        print(f"  {name:<10} {count}")
    print(f"tensor bytes  {format_size(known_bytes)}")
    if unknown:
        print(f"unmeasured    {unknown} tensors have an unknown type or shape")


def print_metadata(metadata: dict) -> None:
    print()
    print("metadata")
    width = max(len(key) for key in metadata) if metadata else 8
    for key in metadata:
        print(f"  {key:<{width}}  {_preview(metadata[key])}")


def print_tensors(tensors: list[dict]) -> None:
    print()
    print(f"{'NAME':<48}  {'TYPE':<8}  {'SHAPE'}")
    for tensor in tensors:
        spec = GGML_TYPES.get(tensor["type"])
        kind = spec[0] if spec else str(tensor["type"])
        # File order: the first dimension is the contiguous one.
        shape = " x ".join(str(dim) for dim in tensor["dims"]) or "-"
        print(f"  {tensor['name']:<46}  {kind:<8}  {shape}")


def inspect(path: Path, metadata_dump: bool, tensor_dump: bool) -> None:
    if not path.is_file():
        raise SystemExit(f"Not a file: {path}")
    try:
        metadata, tensors, version = load_gguf(path)
    except (OSError, ValueError, struct.error, UnicodeError) as exc:
        raise SystemExit(f"{path}: {exc}") from exc
    print_summary(path, metadata, tensors, version)
    if metadata_dump:
        print_metadata(metadata)
    if tensor_dump:
        print_tensors(tensors)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Print architecture, quantization, and tensor info from a GGUF header.")
    parser.add_argument("gguf", nargs="+", type=Path, help="One or more .gguf files")
    parser.add_argument("--metadata", action="store_true", help="Print every metadata key. Long strings and the vocabulary 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.gguf):
        if index:
            print()
        inspect(path, args.metadata, args.tensors)


if __name__ == "__main__":
    main()