#!/usr/bin/env python3
"""Parse GGUF headers from HuggingFace via HTTP range requests (no full download).

Produces a JSON tensor map per shard: name, shape, ggml type, absolute byte
range within the shard file. This lets us (a) recover unsloth's per-tensor
quant-type assignment (their importance map), and (b) later download exactly
the expert tensors we need for the hybrid checkpoint build.
"""
import json
import struct
import sys
import urllib.request

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


def fetch_range(url: str, start: int, length: int) -> bytes:
    req = urllib.request.Request(url)
    req.add_header("Range", f"bytes={start}-{start + length - 1}")
    with urllib.request.urlopen(req) as r:
        return r.read()


class Cursor:
    def __init__(self, buf: bytes):
        self.buf = buf
        self.pos = 0

    def need(self, n):
        if self.pos + n > len(self.buf):
            raise EOFError(f"header larger than fetched window ({len(self.buf)} bytes)")

    def u32(self):
        self.need(4); v = struct.unpack_from("<I", self.buf, self.pos)[0]; self.pos += 4; return v

    def u64(self):
        self.need(8); v = struct.unpack_from("<Q", self.buf, self.pos)[0]; self.pos += 8; return v

    def i64(self):
        self.need(8); v = struct.unpack_from("<q", self.buf, self.pos)[0]; self.pos += 8; return v

    def f32(self):
        self.need(4); v = struct.unpack_from("<f", self.buf, self.pos)[0]; self.pos += 4; return v

    def f64(self):
        self.need(8); v = struct.unpack_from("<d", self.buf, self.pos)[0]; self.pos += 8; return v

    def s(self):
        n = self.u64(); self.need(n)
        v = self.buf[self.pos:self.pos + n].decode("utf-8", errors="replace"); self.pos += n
        return v

    def value(self, t):
        if t == 0: self.need(1); v = self.buf[self.pos]; self.pos += 1; return v
        if t == 1: self.need(1); v = struct.unpack_from("<b", self.buf, self.pos)[0]; self.pos += 1; return v
        if t == 2: self.need(2); v = struct.unpack_from("<H", self.buf, self.pos)[0]; self.pos += 2; return v
        if t == 3: self.need(2); v = struct.unpack_from("<h", self.buf, self.pos)[0]; self.pos += 2; return v
        if t == 4: return self.u32()
        if t == 5: self.need(4); v = struct.unpack_from("<i", self.buf, self.pos)[0]; self.pos += 4; return v
        if t == 6: return self.f32()
        if t == 7: self.need(1); v = self.buf[self.pos] != 0; self.pos += 1; return v
        if t == 8: return self.s()
        if t == 9:  # array
            et = self.u32(); n = self.u64()
            return [self.value(et) for _ in range(n)]
        if t == 10: return self.u64()
        if t == 11: return self.i64()
        if t == 12: return self.f64()
        raise ValueError(f"bad gguf kv type {t}")


def parse_header(url: str, window: int = 4 << 20, want_kv: bool = False):
    buf = fetch_range(url, 0, window)
    while True:
        try:
            return _parse(buf, want_kv)
        except EOFError:
            window *= 4
            if window > (1 << 30):
                raise
            buf = fetch_range(url, 0, window)


def _parse(buf: bytes, want_kv: bool):
    c = Cursor(buf)
    magic = c.u32()
    assert magic == 0x46554747, f"not gguf: {magic:#x}"
    version = c.u32()
    assert version in (2, 3), version
    n_tensors = c.u64()
    n_kv = c.u64()
    kvs = {}
    alignment = 32
    for _ in range(n_kv):
        k = c.s()
        t = c.u32()
        v = c.value(t)
        if k == "general.alignment":
            alignment = v
        if want_kv and not (isinstance(v, list) and len(v) > 64):
            kvs[k] = v
    tensors = []
    for _ in range(n_tensors):
        name = c.s()
        nd = c.u32()
        dims = [c.u64() for _ in range(nd)]
        ty = GGML_TYPES.get(c.u32(), "?")
        off = c.u64()
        tensors.append({"name": name, "dims": dims, "type": ty, "off": off})
    data_start = (c.pos + alignment - 1) // alignment * alignment
    for t in tensors:
        bpb, epb = GGML_BLOCK[t["type"]]
        nelem = 1
        for d in t["dims"]:
            nelem *= d
        nbytes = nelem // epb * bpb
        t["abs_off"] = data_start + t["off"]
        t["nbytes"] = nbytes
    return {"n_tensors": n_tensors, "data_start": data_start, "tensors": tensors, "kv": kvs}


if __name__ == "__main__":
    base = sys.argv[1]      # e.g. https://huggingface.co/unsloth/GLM-5.2-GGUF/resolve/main/UD-Q2_K_XL/GLM-5.2-UD-Q2_K_XL
    nshards = int(sys.argv[2])
    out = sys.argv[3]
    all_shards = {}
    for i in range(1, nshards + 1):
        url = f"{base}-{i:05d}-of-{nshards:05d}.gguf"
        h = parse_header(url, want_kv=(i == 1))
        all_shards[i] = {"url": url, "data_start": h["data_start"], "tensors": h["tensors"], "kv": h.get("kv", {})}
        print(f"shard {i}: {h['n_tensors']} tensors", file=sys.stderr)
    json.dump(all_shards, open(out, "w"))
    print(f"wrote {out}", file=sys.stderr)