"""Layer-streaming MLX runtime for Cloudflare/clef-flash (text path only) + its joint-schema head.

Backbone: mlx-lm's Qwen3.5 DecoderLayer, built one layer at a time from the release shards (np.memmap of the bf16 bytes,
so only the layer being run is resident). The same sanitisation mlx-lm applies to an HF Qwen3.5 checkpoint is done per
layer: conv1d [C,1,K] -> [C,K,1]; RMSNorm weights (input/post-attention/q/k/final) stored as w are shifted to 1 + w.
Records are run layer-major in token-budgeted chunks (each record alone, batch 1, no padding), so the activations of
one chunk + one layer are what is resident.

Head: the release's JointSchemaHead (joint_schema_model.py from the snapshot, unchanged), run in torch on the CPU in
fp32 on the backbone's final-norm hidden states. It reads lm_head rows for the option tokens; lm_head is never resident:
bf16 rows come from a memmap of the shard, 4-bit rows from mlx affine codes (group 64) dequantized per lookup.
"""
import glob, json, math, os, struct, sys, time
from pathlib import Path

import numpy as np
import torch
import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten
from mlx_lm.models.qwen3_5 import TextModelArgs, DecoderLayer
from mlx_lm.models.base import create_attention_mask, create_ssm_mask

_HERE = Path(__file__).resolve().parent
# release dir: CLEF_SNAPSHOT, else this file's directory when it is a packaged release (ships joint_schema_model.py), else the HF cache
SNAP = os.environ.get("CLEF_SNAPSHOT") or (str(_HERE) if (_HERE / "joint_schema_model.py").exists() else
                                           sorted(glob.glob("/Volumes/1Tb/hf-cache/hub/models--Cloudflare--clef-flash/snapshots/*"))[-1])
sys.path.insert(0, SNAP)
sys.dont_write_bytecode = True   # keep the HF snapshot clean (no __pycache__)
from joint_schema_model import JointSchemaHead, encode_record   # noqa: E402  (release code, unchanged)

mx.set_cache_limit(1 << 30)   # as kev/mlx_model.py: bound MLX's buffer cache (one buffer per new record shape otherwise)
LM = "model.language_model."
NORM_SHIFT = (".input_layernorm.weight", ".post_attention_layernorm.weight", ".q_norm.weight", ".k_norm.weight")
SKIP = ("in_proj_a", "in_proj_b")


def bf16_np_to_mx(a):
    return mx.view(mx.array(np.ascontiguousarray(a)), mx.bfloat16)


def bf16_np_to_torch(a):
    return torch.from_numpy(np.ascontiguousarray(a).view(np.int16)).view(torch.bfloat16).float()


class Shards:
    """Header index of the release safetensors shards; tensors are memmapped on demand (bf16 bytes, never converted)."""
    def __init__(self, snap=SNAP):
        self.snap = Path(snap)
        files = set(json.loads((self.snap / "model.safetensors.index.json").read_text())["weight_map"].values())
        self.hdr = {}
        for f in sorted(files):
            p = self.snap / f
            with open(p, "rb") as fh:
                n = struct.unpack("<Q", fh.read(8))[0]; h = json.loads(fh.read(n))
            for k, v in h.items():
                if k != "__metadata__":
                    self.hdr[k] = (str(p), 8 + n + v["data_offsets"][0], v["dtype"], tuple(v["shape"]))

    def memmap(self, k):
        p, off, dt, shape = self.hdr[k]
        assert dt == "BF16", (k, dt)
        return np.memmap(p, dtype=np.uint16, mode="r", offset=off, shape=shape)

    def mx(self, k): return bf16_np_to_mx(self.memmap(k))
    def torch(self, k): return bf16_np_to_torch(self.memmap(k))


def text_args(snap=SNAP):
    return TextModelArgs.from_dict(json.loads((Path(snap) / "config.json").read_text())["text_config"])


def sanitize(n, v):
    if n.endswith("conv1d.weight") and v.shape[-1] != 1: v = v.moveaxis(2, 1)
    if n == "norm.weight" or any(("." + n).endswith(s) for s in NORM_SHIFT): v = v + 1.0   # final norm too (Qwen3_5RMSNorm: 1 + w)
    return v


def layer_weights(sh, i):
    pre = f"{LM}layers.{i}."
    return {k[len(pre):]: sanitize(k[len(pre):], sh.mx(k)) for k in sh.hdr if k.startswith(pre)}


def build_layer(args, i, weights):
    layer = DecoderLayer(args, i)
    layer.eval()   # mlx Modules start in training mode, where GatedDeltaNet takes the slow reference ops instead of the Metal kernel
    layer.load_weights(list(weights.items()), strict=True)
    mx.eval(layer.parameters())
    return layer


def linears(layer, i):
    """(i, name, module, packed name) for every quantizable nn.Linear of a layer (in_proj_a / in_proj_b stay bf16)."""
    return [(i, n, m, f"{LM}layers.{i}.{n}.weight") for n, m in layer.named_modules()
            if isinstance(m, nn.Linear) and n.split(".")[-1] not in SKIP]


def set_module(layer, name, mod):
    parent = layer
    for p in name.split(".")[:-1]: parent = getattr(parent, p)
    setattr(parent, name.split(".")[-1], mod)


class QLoRA(nn.Module):
    """Quantized linear + bf16 low-rank compensation (CLoQ), kept unfolded: y = Q(x) + (x A^T) B^T."""
    def __init__(self, q, a, b):
        super().__init__(); self.q = q; self.lora_a = a; self.lora_b = b
    def __call__(self, x):
        return self.q(x) + (x @ self.lora_a.T) @ self.lora_b.T


def mlx_pack(codes, bits):
    """uint8 codes [N, K] -> mlx affine packed uint32 [N, K*bits/32] (element j of a word at bit j*bits)."""
    N, K = codes.shape; per = 32 // bits
    c = codes.to(torch.int64).reshape(N, K // per, per)
    w = (c << (torch.arange(per, dtype=torch.int64) * bits)).sum(-1)
    return mx.array(w.numpy().astype(np.uint32))


def qlinear(codes, scales, biases, bits, group=64):
    """Build an nn.QuantizedLinear from explicit codes / bf16 scales / bf16 biases (w = scale * code + bias)."""
    N, K = codes.shape
    q = nn.QuantizedLinear(K, N, bias=False, group_size=group, bits=bits)
    q.weight = mlx_pack(codes, bits)
    q.scales = mx.array(scales.numpy()).astype(mx.bfloat16)
    q.biases = mx.array(biases.numpy()).astype(mx.bfloat16)
    return q


def qembedding(w):
    V, D = w.shape
    e = nn.QuantizedEmbedding(V, D, group_size=64, bits=4)
    e.weight, e.scales, e.biases = mx.quantize(w, group_size=64, bits=4)
    return e


class Head:
    """JointSchemaHead (fp32, CPU) + lm_head row lookup in bf16 ('bf16') or 4-bit mlx affine ('q4')."""
    def __init__(self, snap=SNAP, sh=None, lm_head_q=None):
        from safetensors.torch import load_file
        torch.set_num_threads(8)
        cfg = json.loads((Path(snap) / "joint_head_config.json").read_text())
        self.head = JointSchemaHead(**cfg)
        self.head.load_state_dict({k: v.float() for k, v in load_file(str(Path(snap) / "joint_head.safetensors")).items()}, strict=True)
        self.head.eval()
        self.sh = sh; self.q = lm_head_q   # (wq, scales, biases) or None

    def rows_bf16(self, ids):
        return bf16_np_to_torch(self.sh.memmap("lm_head.weight")[np.asarray(ids)])

    def rows_q4(self, ids):
        ix = mx.array(np.asarray(ids, dtype=np.int32)); wq, s, b = self.q
        return torch.from_numpy(np.array(mx.dequantize(wq[ix], s[ix], b[ix], group_size=64, bits=4).astype(mx.float32)))

    @torch.no_grad()
    def logits(self, hidden, enc, mode):
        class Rows:
            def __getitem__(_, ids): return (self.rows_bf16 if mode == "bf16" else self.rows_q4)(ids.cpu().numpy())
        L = len(enc.input_ids)
        out = self.head(hidden.reshape(1, L, -1), torch.tensor([enc.input_ids]), torch.ones(1, L, dtype=torch.long), [enc], Rows())
        return [z.float() for z in out[0]]


def chunks_by_tokens(encs, budget):
    out, cur, n = [], [], 0
    for j, e in enumerate(encs):
        if cur and n + len(e.input_ids) > budget: out.append(cur); cur, n = [], 0
        cur.append(j); n += len(e.input_ids)
    if cur: out.append(cur)
    return out


def run_layer(layer, xs):
    ys = []
    for x in xs:
        m = create_ssm_mask(x) if layer.is_linear else create_attention_mask(x)
        y = layer(x, mask=m, cache=None); mx.eval(y); ys.append(y)
    return ys


def forward_logits(encs, embed, layer_at, final_norm, head, modes=("q4", "bf16"), budget=120_000, log=print):
    """Backbone (layer-major per chunk) + head. embed(ids)->[1,L,D]; layer_at(i)->DecoderLayer (built or resident).
    Returns {mode: [per-record list of per-question fp32 logit tensors]}."""
    nl = len_layers = layer_at.n
    res = {m: [None] * len(encs) for m in modes}
    for ci, idx in enumerate(chunks_by_tokens(encs, budget)):
        t0 = time.time()
        xs = [embed(encs[j].input_ids) for j in idx]
        for i in range(nl):
            layer = layer_at(i); xs = run_layer(layer, xs); del layer
            mx.clear_cache()
        for j, x in zip(idx, xs):
            h = final_norm(x); mx.eval(h)
            ht = torch.from_numpy(np.array(h.astype(mx.float32)))
            for m in modes: res[m][j] = head.logits(ht, encs[j], m)
        del xs; mx.clear_cache()
        log(f"chunk {ci}: {len(idx)} records in {time.time()-t0:.0f}s, peak {mx.get_peak_memory()/1e9:.2f} GB")
    return res


class FinalNorm:
    def __init__(self, w, eps):
        self.w = w; self.eps = eps
    def __call__(self, x): return mx.fast.rms_norm(x, self.w, self.eps)


class Bf16Backbone:
    """Streams bf16 layers from the shards; embeddings gathered from the memmap (never fully resident)."""
    def __init__(self, sh, args):
        self.sh, self.args, self.n = sh, args, args.num_hidden_layers
        self.norm = FinalNorm(sanitize("norm.weight", sh.mx(LM + "norm.weight")), args.rms_norm_eps)
        self.emb = sh.memmap(LM + "embed_tokens.weight")
    def __call__(self, i): return build_layer(self.args, i, layer_weights(self.sh, i))
    def embed(self, ids): return bf16_np_to_mx(self.emb[np.asarray(ids)])[None]


def predictions(records, encs, logits):
    """kev prediction dicts (probabilities + logits keyed like kev.api.question_keys) from Clef per-option logits."""
    out = []
    for r, e, zs in zip(records, encs, logits):
        probs, lg = {}, {}
        for q, z in zip(e.questions, zs):
            p = torch.softmax(z, -1)
            probs[q.question_id] = {o: float(v) for o, v in zip(q.option_ids, p.tolist())}
            lg[q.question_id] = {o: float(v) for o, v in zip(q.option_ids, z.tolist())}
        out.append({"probabilities": probs, "logits": lg, "inference_temperature": 1.0, "latency_ms": 0.0, "input_tokens": len(e.input_ids)})
    return out


class Cached:
    """A predictor for kev.benchmark.evaluate_records that returns precomputed predictions in record order."""
    def __init__(self, preds): self.preds = preds; self.k = 0; self.temperature = 1.0
    def __call__(self, record):
        p = self.preds[self.k]; self.k += 1; return p