#!/usr/bin/env python3
"""Dump GGUF part headers: split metadata, tensor names/dims/types/bytes.

Reads ONLY the header region (metadata + tensor-info table); tensor data is
never touched, so a 50 GiB part scans in milliseconds. Split GGUFs are valid
GGUFs per part; part 1 of this payload is metadata-only (ntensors=0) and
parts 2-6 carry the tensors. Confirmed layout for this payload: GGUF v3 —
magic, u32 version, u64 ntensors, u64 nkv, then kvs (v3: u64-length strings),
then tensor infos (v3: u64-length names).
"""
import struct, sys, traceback

TYPE_SIZE = {0:4,1:2,2:1,3:1/2,4:4,5:8,6:2,7:1,8:0.4375,9:1.4375,10:1,11:0.5,
             12:2,13:4/3,14:2,15:2,16:4,17:2,18:2,19:1,20:2,21:1,22:2,23:2,
             24:2,25:1.0625,26:1,27:1.375,28:2,29:2,30:2,31:4,32:8,33:2.125,
             34:2.5,35:2,36:2.375,37:1.5625,38:2.3125,39:2.1875,40:0,41:1.75,
             42:1.75,43:1.8125,44:2,45:2.25,46:1.5,47:2.0625,48:2.625,49:1.0625,
             50:2.3125,51:2.875,52:2,53:2.3125,54:2.75,55:2.5,56:2.875,57:2.25,
             58:2,59:3.8125,60:2.8125,61:1.5625,62:1.5,63:0.5}

class P:
    def __init__(self, f):
        self.f = f
        self.ver = 2
    def rd(self, n):
        b = self.f.read(n)
        if len(b) != n:
            raise EOFError(f"short read {len(b)}/{n}")
        return b
    def u8(self): return self.rd(1)[0]
    def u32(self): return struct.unpack("<I", self.rd(4))[0]
    def u64(self): return struct.unpack("<Q", self.rd(8))[0]
    def f32(self): return struct.unpack("<f", self.rd(4))[0]
    def f64(self): return struct.unpack("<d", self.rd(8))[0]
    def s(self):
        # GGUF v3 strings carry a u64 length; v2 used u32.
        n = self.u64() if self.ver >= 3 else self.u32()
        return self.rd(n).decode("utf-8", "replace")
    def val(self, t):
        # GGUF kv value types: 0 u8, 1 i8, 2 u16, 3 i16, 4 u32, 5 i32,
        # 6 f32, 7 bool, 8 string, 9 array, 10 u64, 11 i64, 12 f64.
        if t == 8: return self.s()
        if t == 9:
            et = self.u32(); n = self.u64()
            return [self.val(et) for _ in range(n)]
        if t == 0: return self.u8()
        if t == 1: return struct.unpack("<b", self.rd(1))[0]
        if t == 2: return struct.unpack("<H", self.rd(2))[0]
        if t == 3: return struct.unpack("<h", self.rd(2))[0]
        if t == 4: return self.u32()
        if t == 5: return struct.unpack("<i", self.rd(4))[0]
        if t == 6: return self.f32()
        if t == 7: return self.u8()
        if t == 10: return self.u64()
        if t == 11: return struct.unpack("<q", self.rd(8))[0]
        if t == 12: return self.f64()
        raise ValueError(f"unknown kv type {t}")

def main(path):
    tensors = []
    kvs = {}
    with open(path, "rb") as f:
        p = P(f)
        magic = p.rd(4)
        assert magic == b"GGUF", magic
        p.ver = p.u32()
        ntens = p.u64()
        nkv = p.u64()
        for _ in range(nkv):
            k = p.s()
            kvs[k] = p.val(p.u32())
        for _ in range(ntens):
            name = p.s()
            nd = p.u32()
            dims = [p.u64() for _ in range(nd)]
            t = p.u32()
            off = p.u64()
            numel = 1
            for d in dims: numel *= d
            tensors.append((name, dims, t, numel))
    print(f"== {path}  (gguf v{p.ver}, {ntens} tensors, {nkv} kv)")
    for k in ("general.architecture", "general.name", "split.no", "split.count",
              "split.tensors", "qwen4exp.expert_count", "qwen4exp.expert_used_count",
              "qwen4exp.ngram.tensors"):
        if k in kvs: print(f"  kv {k} = {kvs[k]}")
    ngram_kvs = {k: v for k, v in kvs.items() if "ngram" in k.lower()}
    if ngram_kvs: print("  ngram kvs:", ngram_kvs)
    big = []
    for name, dims, t, numel in tensors:
        big.append((numel * TYPE_SIZE.get(t, 0), name, dims, t, numel))
    big.sort(reverse=True)
    print(f"  -- {min(12, len(big))} largest of {len(big)}:")
    for bytes_, name, dims, t, numel in big[:12]:
        print(f"    {bytes_/2**30:9.3f} GiB  {name}  dims={dims} type={t} ({numel:,} elems)")
    named = [(b, n, d, t) for b, n, d, t, e in big
             if any(s in n.lower() for s in ("ngram", "ple", "token_embd", "exps"))]
    if named:
        print("  -- ngram/ple/embed/experts-named:")
        for bytes_, name, dims, t in named[:20]:
            print(f"    {bytes_/2**30:9.3f} GiB  {name}  dims={dims} type={t}")

for pth in sys.argv[1:]:
    try:
        main(pth)
    except Exception:
        print(f"== {pth}: FAILED")
        traceback.print_exc()