#!/usr/bin/env python3
"""BF16 teacher streamer for GLM-5.2 (phase-2 full-AQLM teacher).

Backs `get_expert(layer, expert, proj) -> torch.bfloat16` with HTTP range
reads of zai-org/GLM-5.2 (the original BF16 weights). The safetensors index
is fetched once; per-shard headers are cached on first touch under
/data/bf16-cache/headers/; fetched tensor byte ranges live in an on-disk LRU
blob cache under /data/bf16-cache/blobs/, capped at 400 GB with mtime-based
eviction.

Range-fetch pattern reused from tools/gguf_remote.py and
tools/make_hot_manifest.py (safetensors: 8-byte little-endian header length,
JSON header of {name: {dtype, shape, data_offsets:[b,e]}}, tensor bytes at
8 + header_len + b .. 8 + header_len + e).

Run directly to VERIFY against the NVFP4 teacher:
    python tools/bf16_stream.py [--n 10]
prints cosine similarity for N random (layer, expert, proj) triples between
the streamed BF16 tensor and the NVFP4-dequantized teacher (regions in
/tmp/glm52-hot-dl2); all must exceed 0.98.
"""
import argparse
import hashlib
import json
import os
import struct
import time
import urllib.request

BASE = "https://huggingface.co/zai-org/GLM-5.2/resolve/main"
CACHE = "/data/bf16-cache"
CAP_BYTES = 400 * (1 << 30)          # 400 GB LRU cap on the blob cache
LOCAL_TEACHER_LAYERS = {3, 4, 5, 8, 74, 75, 76, 77}

# canonical projection names + convenient aliases
_PROJ = {
    "gate_proj": "gate_proj", "up_proj": "up_proj", "down_proj": "down_proj",
    "gate": "gate_proj", "up": "up_proj", "down": "down_proj",
    "w1": "gate_proj", "w3": "up_proj", "w2": "down_proj",
}

# safetensors dtype -> (torch dtype str, numpy dtype for raw view)
_ST_DT = {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}


def _fetch(url, start=None, length=None, retries=6, timeout=120):
    """HTTP GET (optionally a byte range) with exponential backoff."""
    last = None
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url)
            if start is not None:
                req.add_header("Range", f"bytes={start}-{start + length - 1}")
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as e:            # noqa: BLE001 - transient net errors
            last = e
            if attempt == retries - 1:
                break
            time.sleep(min(30, 2 ** attempt))
    raise RuntimeError(f"fetch failed {url} [{start}:{length}]: {last}")


class Bf16Teacher:
    def __init__(self, base=BASE, cache=CACHE, cap_bytes=CAP_BYTES):
        self.base = base.rstrip("/")
        self.cache = cache
        self.cap = cap_bytes
        self.hdr_dir = os.path.join(cache, "headers")
        self.blob_dir = os.path.join(cache, "blobs")
        os.makedirs(self.hdr_dir, exist_ok=True)
        os.makedirs(self.blob_dir, exist_ok=True)
        self._headers = {}               # shard -> {name: info}
        self._hdr_len = {}               # shard -> header_len (bytes offset)
        self._wm = None                  # weight_map (name -> shard)

    # -------------------------------------------------- index / headers
    def _index(self):
        if self._wm is None:
            p = os.path.join(self.cache, "index.json")
            if os.path.exists(p):
                idx = json.load(open(p))
            else:
                idx = json.loads(_fetch(f"{self.base}/model.safetensors.index.json"))
                json.dump(idx, open(p, "w"))
            self._wm = idx["weight_map"]
        return self._wm

    def _header(self, shard):
        if shard in self._headers:
            return self._headers[shard], self._hdr_len[shard]
        hp = os.path.join(self.hdr_dir, shard + ".json")
        if os.path.exists(hp):
            d = json.load(open(hp))
            hdr, hlen = d["header"], d["header_len"]
        else:
            url = f"{self.base}/{shard}"
            hlen = struct.unpack("<Q", _fetch(url, 0, 8))[0]
            hdr = json.loads(_fetch(url, 8, hlen))
            hdr.pop("__metadata__", None)
            json.dump({"header_len": hlen, "header": hdr}, open(hp, "w"))
        self._headers[shard] = hdr
        self._hdr_len[shard] = hlen
        return hdr, hlen

    # -------------------------------------------------- blob LRU cache
    def _blob_path(self, key):
        h = hashlib.sha1(key.encode()).hexdigest()
        return os.path.join(self.blob_dir, h + ".bin")

    def _evict_if_needed(self, incoming):
        files = []
        total = 0
        with os.scandir(self.blob_dir) as it:
            for e in it:
                if e.name.endswith(".bin"):
                    st = e.stat()
                    files.append((st.st_mtime, e.path, st.st_size))
                    total += st.st_size
        if total + incoming <= self.cap:
            return
        files.sort()                      # oldest mtime first
        for _, path, size in files:
            if total + incoming <= self.cap:
                break
            try:
                os.remove(path)
                total -= size
            except OSError:
                pass

    def _get_bytes(self, name):
        wm = self._index()
        if name not in wm:
            raise KeyError(name)
        shard = wm[name]
        hdr, hlen = self._header(shard)
        info = hdr[name]
        b, e = info["data_offsets"]
        key = f"{shard}:{name}"
        bp = self._blob_path(key)
        if os.path.exists(bp):
            os.utime(bp, None)            # LRU touch
            return open(bp, "rb").read(), info
        length = e - b
        buf = _fetch(f"{self.base}/{shard}", 8 + hlen + b, length)
        self._evict_if_needed(len(buf))
        tmp = bp + f".tmp{os.getpid()}"
        with open(tmp, "wb") as f:
            f.write(buf)
        os.replace(tmp, bp)
        return buf, info

    # -------------------------------------------------- public API
    def get_expert(self, layer: int, expert: int, proj: str):
        import torch

        pj = _PROJ.get(proj, proj)
        name = f"model.layers.{layer}.mlp.experts.{expert}.{pj}.weight"
        buf, info = self._get_bytes(name)
        dt = _ST_DT.get(info["dtype"])
        if dt is None:
            raise ValueError(f"unexpected dtype {info['dtype']} for {name}")
        t = torch.frombuffer(bytearray(buf), dtype=getattr(torch, dt))
        return t.reshape(info["shape"]).to(torch.bfloat16)


# ------------------------------------------------------------- verification
def _verify(n):
    import random

    import torch

    import sys
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from aqlm_converge import TeacherReader, dequant_teacher, FP4_LUT

    teacher = Bf16Teacher()
    tr = TeacherReader()
    lut = torch.tensor(FP4_LUT, dtype=torch.float32)

    # experts with actual NVFP4 region coverage, layers using the ranged path
    avail = {}
    for shard, h in tr.headers.items():
        for nm in h["header"]:
            if ".mlp.experts." in nm and nm.endswith("gate_proj.weight"):
                p = nm.split(".")
                li, e = int(p[2]), int(p[5])
                if li not in LOCAL_TEACHER_LAYERS:
                    avail.setdefault(li, set()).add(e)
    pool = [(li, e) for li, es in avail.items() for e in es]
    rng = random.Random(1234)
    rng.shuffle(pool)

    cosines = []
    tried = 0
    for li, e in pool:
        if len(cosines) >= n:
            break
        proj = rng.choice(["gate_proj", "up_proj", "down_proj"])
        tried += 1
        try:
            nv = dequant_teacher(tr, li, e, proj, "cpu", lut)      # fp32 [out,in]
        except KeyError:
            continue
        bf = teacher.get_expert(li, e, proj).float()
        a = nv.reshape(-1)
        b = bf.reshape(-1)
        cos = torch.dot(a, b) / (a.norm() * b.norm() + 1e-12)
        cosines.append(cos.item())
        print(f"  L{li:>2} e{e:<3} {proj:<10} cos={cos.item():.5f}")
    print(f"\n{len(cosines)} cosines (from {tried} tries); "
          f"min={min(cosines):.5f} mean={sum(cosines)/len(cosines):.5f}")
    assert all(c > 0.98 for c in cosines), "cosine <= 0.98 for some triple"
    print("VERIFY PASS: all cosines > 0.98")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=10)
    _verify(ap.parse_args().n)