#!/usr/bin/env python3
"""Manifest of NVFP4 per-expert tensors for the 2-tier hot sets.

Fetches the HF repo's safetensors index + shard headers via range requests
and emits (url, offset, length, out) entries for every tensor of every hot
expert in the assignment, excluding layers whose NVFP4 bytes are already in
the local checkpoint's compacted arrays (LOCAL_LAYERS).
"""
import json
import os
import struct
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor

BASE = "https://huggingface.co/lukealonso/GLM-5.2-NVFP4/resolve/main"
DL = "/tmp/glm52-hot-dl2"
ASSIGN = "/data/glm52-need-experts.json"
LOCAL_LAYERS = set()
COALESCE_GAP = 4 << 20


def fetch(url, start=None, length=None):
    req = urllib.request.Request(url)
    if start is not None:
        req.add_header("Range", f"bytes={start+0}-{start+length-1}")
    for attempt in range(6):
        try:
            with urllib.request.urlopen(req, timeout=90) as r:
                return r.read()
        except Exception:
            if attempt == 5:
                raise
    return None


def main():
    os.makedirs(f"{DL}/regions", exist_ok=True)
    assign = {int(k): v for k, v in json.load(open(ASSIGN)).items()}
    want = set()
    for li, a in assign.items():
        if li in LOCAL_LAYERS:
            continue
        for e in a:
            want.add((li, e))
    print(f"{len(want)} hot experts to download", file=sys.stderr)

    idx = json.loads(fetch(f"{BASE}/model.safetensors.index.json"))
    wm = idx["weight_map"]
    import re
    by_shard: dict[str, list[str]] = {}
    pat = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.")
    for name, shard in wm.items():
        m = pat.match(name)
        if m and (int(m.group(1)), int(m.group(2))) in want:
            by_shard.setdefault(shard, []).append(name)

    def one(shard):
        url = f"{BASE}/{shard}"
        hlen = struct.unpack("<Q", fetch(url, 0, 8))[0]
        hdr = json.loads(fetch(url, 8, hlen))
        base = 8 + hlen
        ranges = sorted(
            hdr[n]["data_offsets"] for n in by_shard[shard] if n in hdr
        )
        merged = []
        for b, e in ranges:
            if merged and b - merged[-1][1] <= COALESCE_GAP:
                merged[-1][1] = max(merged[-1][1], e)
            else:
                merged.append([b, e])
        ents = [{"url": url, "start": base + b, "length": e - b,
                 "out": f"{DL}/regions/{shard}/{b}.bin", "rel_start": b}
                for b, e in merged]
        return shard, hlen, hdr, ents

    entries, headers = [], {}
    with ThreadPoolExecutor(16) as ex:
        for shard, hlen, hdr, ents in ex.map(one, sorted(by_shard)):
            headers[shard] = {"header_len": hlen, "header": hdr}
            entries.extend(ents)
    json.dump(headers, open(f"{DL}/headers.json", "w"))
    json.dump(entries, open(f"{DL}/manifest.json", "w"))
    total = sum(e["length"] for e in entries)
    print(f"{len(entries)} regions, {total/1e9:.1f} GB", file=sys.stderr)


if __name__ == "__main__":
    main()