#!/usr/bin/env python3
"""Generate download manifests (url, byte range, output path) for the hybrid build.

- NVFP4 repo: parse each shard's safetensors header remotely; select all
  non-routed-expert tensors plus routed experts of plan.nvfp4_layers;
  coalesce adjacent ranges.
- GGUF repos: use the tensor maps produced by gguf_remote.py to select
  blk.N.ffn_{gate,up,down}_exps.weight for the planned layers.
"""
import json
import re
import struct
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor

ROOT = "/home/coder/git/glm52"
DL = "/tmp/glm52-dl"
plan = json.load(open(f"{ROOT}/hybrid_plan.json"))

NVFP4_LAYERS = set(plan["nvfp4_layers"])
GGUF_LAYERS = {**{l: "ud_q3" for l in plan["ud_q3_layers"]},
               **{l: "ud_q2" for l in plan["ud_q2_layers"]}}
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}-{start+length-1}")
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                return r.read()
        except Exception as e:
            if attempt == 4:
                raise
    return None


def needed_nvfp4(name):
    m = re.match(r"model\.layers\.(\d+)\.mlp\.experts\.", name)
    if m:
        return int(m.group(1)) in NVFP4_LAYERS
    return True


def st_header(url):
    n = struct.unpack("<Q", fetch(url, 0, 8))[0]
    hdr = json.loads(fetch(url, 8, n))
    return n, hdr


def nvfp4_manifest():
    idx = json.load(open("/tmp/glm52-recon/nvfp4_model.safetensors.index.json"))
    shards = sorted(set(idx["weight_map"].values()))
    entries, headers = [], {}

    def one(shard):
        url = f'{plan["sources"]["nvfp4"]}/{shard}'
        hlen, hdr = st_header(url)
        data_base = 8 + hlen
        ranges = []
        for name, info in hdr.items():
            if name == "__metadata__":
                continue
            if needed_nvfp4(name):
                b, e = info["data_offsets"]
                ranges.append((b, e))
        ranges.sort()
        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": data_base + b,
                 "length": e - b,
                 "out": f"{DL}/nvfp4/regions/{shard}/{b}.bin",
                 "rel_start": b} for b, e in merged]
        return shard, hlen, hdr, ents

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


def gguf_manifest(tag):
    m = json.load(open(f"/tmp/glm52-recon/{'q3kxl' if tag=='ud_q3' else 'q2kxl'}_map.json"))
    layers = {l for l, t in GGUF_LAYERS.items() if t == tag}
    entries = []
    tinfo = {}
    for sh in m.values():
        for t in sh["tensors"]:
            mm = re.match(r"blk\.(\d+)\.(ffn_(?:gate|up|down)_exps)\.weight", t["name"])
            if mm and int(mm.group(1)) in layers:
                out = f"{DL}/gguf/{tag}/{t['name']}.{t['type']}.bin"
                entries.append({"url": sh["url"], "start": t["abs_off"],
                                "length": t["nbytes"], "out": out})
                tinfo[t["name"]] = {"type": t["type"], "dims": t["dims"],
                                    "nbytes": t["nbytes"], "file": out}
    json.dump(tinfo, open(f"{DL}/gguf/{tag}_tensors.json", "w"))
    return entries


def main():
    import os
    for d in ["nvfp4/regions", "gguf/ud_q3", "gguf/ud_q2"]:
        os.makedirs(f"{DL}/{d}", exist_ok=True)
    ents = []
    ents += gguf_manifest("ud_q3")
    ents += gguf_manifest("ud_q2")
    ents += nvfp4_manifest()
    total = sum(e["length"] for e in ents)
    json.dump(ents, open(f"{DL}/manifest.json", "w"))
    print(f"manifest: {len(ents)} entries, {total/1e9:.1f} GB total", file=sys.stderr)


if __name__ == "__main__":
    main()