#!/usr/bin/env python3
"""Prepare pruned K=256 Qwen3.8-Flash-Next for Strata Metal engine:
1. Generates K256-expert-profile.bin (STRP v1, 48 layers x 256 experts) ranked by
   multi-domain calibration frequency & REAP saliency (results/reap_saliency_v2.json).
2. Writes K256-metal.gguf (4.22 GiB = 4.53 GB dense tensors: 2.90 GB active GPU dense +
   1.62 GB skipped BF16 HC / CPU token_embd, 16 KiB aligned, F_NOCACHE) so Strata wraps
   dense weights zero-copy as clean file-backed pages (0 GiB anonymous RAM).
3. Writes K256-metal.gguf.experts (23.42 GiB STRX v1 contiguous [gate | up | down]
   per-expert slices, F_NOCACHE) so each missed expert is read in 1 contiguous pread.
"""

import fcntl
import json
import os
import struct
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "vendor" / "strata" / "metal" / "tools"))
import metal_repack as R

ALIGN = 16384
CHUNK = 64 << 20
SYNC_EVERY = 1 << 30
NL = 48
NE = 256


def nocache(f) -> None:
    try:
        fcntl.fcntl(f.fileno(), getattr(fcntl, "F_NOCACHE", 48), 1)
    except Exception:
        pass


def build_expert_profile(saliency_path: Path, keep_map_path: Path, out_path: Path) -> None:
    sal = json.loads(saliency_path.read_text())
    km = json.loads(keep_map_path.read_text())
    scored = []
    for l in range(NL):
        l_str = str(l)
        kept_orig = km[l_str]
        assert len(kept_orig) == NE, f"Layer {l} has {len(kept_orig)} experts != {NE}"
        counts = sal["layers"][l_str]["count"]
        probs = sal["layers"][l_str]["prob_sum"]
        for new_e, orig_e in enumerate(kept_orig):
            # Rank by empirical token selection count + prob_sum tie-breaker
            score = float(counts[orig_e]) + 1e-3 * float(probs[orig_e])
            scored.append((score, l, new_e))
    scored.sort(key=lambda x: (-x[0], x[1], x[2]))
    assert len(scored) == NL * NE

    buf = bytearray()
    buf += struct.pack("<4s5I", b"STRP", 1, NL, NE, len(scored), len(scored))
    for _, l, new_e in scored:
        buf += struct.pack("<2H", l, new_e)
    out_path.write_bytes(bytes(buf))
    print(f"[1/3] Wrote STRP profile {out_path} ({len(buf)} bytes, top pair=L{scored[0][1]}:E{scored[0][2]} hits={scored[0][0]:.0f})", flush=True)


def build_metal_gguf_and_experts(shard1: Path, out_gguf: Path, out_exp: Path) -> None:
    t0 = time.monotonic()
    shards = [R.Reader(p) for p in R.shard_paths(shard1)]
    every = [t for s in shards for t in s.tensors]
    keep = [t for t in every if t["name"] != R.PLE]

    def layer_of(t):
        import re
        m = re.match(r"blk\.(\d+)\.", t["name"])
        return int(m.group(1)) if m else -1

    dense = [t for t in keep if "_exps." not in t["name"]]
    role = {"ffn_gate_exps": 0, "ffn_up_exps": 1, "ffn_down_exps": 2}
    exps = sorted(
        (t for t in keep if "_exps." in t["name"]),
        key=lambda t: (layer_of(t), role.get(t["name"].split(".")[2], 9)),
    )
    order = dense + exps

    kv = [raw for key, raw in shards[0].kv if not key.startswith("split.") and key != "general.alignment"]
    kv.append(R.gguf_str("general.alignment") + struct.pack("<II", 4, ALIGN))

    off = 0
    infos = []
    dense_end_rel = 0
    for t in order:
        t["new"] = off
        infos.append(
            R.gguf_str(t["name"])
            + struct.pack("<I", len(t["dims"]))
            + struct.pack(f"<{len(t['dims'])}Q", *t["dims"])
            + struct.pack("<IQ", t["type"], off)
        )
        off = (off + t["bytes"] + ALIGN - 1) // ALIGN * ALIGN
        if "_exps." not in t["name"]:
            dense_end_rel = off

    head = struct.pack("<4sIQQ", b"GGUF", 3, len(order), len(kv)) + b"".join(kv) + b"".join(infos)
    data0 = (len(head) + ALIGN - 1) // ALIGN * ALIGN
    dense_total = data0 + dense_end_rel

    if out_gguf.exists() and out_gguf.stat().st_size == dense_total:
        print(f"[2/3] {out_gguf} already exists ({dense_total / 2**30:.2f} GiB)", flush=True)
    else:
        tmp_g = out_gguf.with_name(out_gguf.name + ".part")
        with open(tmp_g, "wb", buffering=0) as o:
            nocache(o)
            R.write_all(o, head)
            R.write_all(o, b"\0" * (data0 - len(head)))
            done = 0
            since = 0
            for t in dense:
                o.seek(data0 + t["new"])
                src = t["src"].f
                src.seek(t["abs"])
                left = t["bytes"]
                while left:
                    buf = src.read(min(left, CHUNK))
                    if not buf:
                        raise RuntimeError(f"Short read in {t['name']}")
                    R.write_all(o, buf)
                    left -= len(buf)
                    since += len(buf)
                    if since >= SYNC_EVERY:
                        os.fsync(o.fileno())
                        since = 0
                done += t["bytes"]
            o.truncate(dense_total)
            os.fsync(o.fileno())
        os.replace(tmp_g, out_gguf)
        print(f"[2/3] Wrote dense-only {out_gguf} ({dense_total / 2**30:.2f} GiB in {time.monotonic() - t0:.1f}s)", flush=True)

    # 3. Write STRX .experts file (contiguous [gate_e | up_e | down_e] per expert)
    t1 = time.monotonic()
    tab = {t["name"]: t for t in exps}
    lay = []
    for l in range(NL):
        g, u, d = (tab[f"blk.{l}.ffn_{k}_exps.weight"] for k in ("gate", "up", "down"))
        gs, us, ds = g["bytes"] // NE, u["bytes"] // NE, d["bytes"] // NE
        lay.append((g, u, d, gs, us, ds))

    exp_head_len = 64 + NL * 32
    exp_data0 = (exp_head_len + ALIGN - 1) // ALIGN * ALIGN
    offs, at = [], exp_data0
    for g, u, d, gs, us, ds in lay:
        offs.append(at)
        at += NE * (gs + us + ds)
        at = (at + ALIGN - 1) // ALIGN * ALIGN

    if out_exp.exists() and out_exp.stat().st_size == at:
        print(f"[3/3] {out_exp} already exists ({at / 2**30:.2f} GiB)", flush=True)
    else:
        exp_head = struct.pack("<4sIIIQ", b"STRX", 1, NL, NE, dense_total).ljust(64, b"\0")
        for (g, u, d, gs, us, ds), o_off in zip(lay, offs):
            exp_head += struct.pack("<QQQQ", o_off, gs, us, ds)

        tmp_e = out_exp.with_name(out_exp.name + ".part")
        with open(tmp_e, "wb", buffering=0) as f:
            nocache(f)
            R.write_all(f, exp_head.ljust(exp_data0, b"\0"))
            for l, ((g, u, d, gs, us, ds), o_off) in enumerate(zip(lay, offs)):
                g["src"].f.seek(g["abs"]); G = g["src"].f.read(g["bytes"])
                u["src"].f.seek(u["abs"]); U = u["src"].f.read(u["bytes"])
                d["src"].f.seek(d["abs"]); D = d["src"].f.read(d["bytes"])
                f.seek(o_off)
                layer_blob = b"".join(
                    G[e * gs : (e + 1) * gs] + U[e * us : (e + 1) * us] + D[e * ds : (e + 1) * ds]
                    for e in range(NE)
                )
                R.write_all(f, layer_blob)
                if (l + 1) % 8 == 0 or l == NL - 1:
                    os.fsync(f.fileno())
                    print(f"  STRX layer {l + 1}/{NL} ({f.tell() / 2**30:.2f} GiB)", flush=True)
            f.truncate(at)
            os.fsync(f.fileno())
        os.replace(tmp_e, out_exp)
        print(f"[3/3] Wrote STRX {out_exp} ({at / 2**30:.2f} GiB in {time.monotonic() - t1:.1f}s)", flush=True)

    for s in shards:
        s.f.close()


def main() -> int:
    root = Path(__file__).resolve().parent.parent
    model_dir = root / "models" / "K256_reap_l2_sum"
    shard1 = model_dir / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf"
    out_prof = model_dir / "K256-expert-profile.bin"
    out_gguf = model_dir / "K256-metal.gguf"
    out_exp = model_dir / "K256-metal.gguf.experts"

    build_expert_profile(
        root / "results" / "reap_saliency_v2.json",
        root / "results" / "keep_maps" / "keep_reap_l2_sum_k256.json",
        out_prof,
    )
    build_metal_gguf_and_experts(shard1, out_gguf, out_exp)

    tok_dir = model_dir / "tokenizer"
    if not (tok_dir / "vocab.json").exists():
        sys.path.insert(0, str(root / "vendor" / "strata" / "tools"))
        import strata_tokenizer as ST
        ST.extract(str(out_gguf), str(model_dir))
    jinja_p = tok_dir / "chat_template.jinja"
    if jinja_p.exists():
        jtxt = jinja_p.read_text(encoding="utf-8")
        if "reasoning_effort|default('xhigh')" in jtxt:
            jinja_p.write_text(jtxt.replace("reasoning_effort|default('xhigh')", "reasoning_effort|default('low')"), encoding="utf-8")

    healed_gguf = model_dir / "K256-q2_0-healed-metal.gguf"
    active_gguf = healed_gguf if healed_gguf.exists() else out_gguf
    vocab_out = model_dir / "vocab.gguf"
    if not vocab_out.exists():
        rd_v = R.Reader(active_gguf)
        rd_v.f.close()
        vocab_out.write_bytes(struct.pack("<4sIQQ", b"GGUF", rd_v.version, 0, len(rd_v.kv)) + b"".join(raw for _, raw in rd_v.kv))
        del rd_v
    vexe = root / "vendor" / "strata" / "build-vision" / "bin" / "strata-vision"
    mmproj_p = root / "models" / "K512_iq2_xs" / "mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf"
    if not mmproj_p.exists():
        mmproj_p = model_dir / "mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf"
    has_vision = vexe.exists() and mmproj_p.exists()
    cfg = {
        "exe": str(root / "vendor" / "strata" / "build-metal" / "strata_window_cli"),
        "args": [
            "--metal-gguf", str(active_gguf),
            "--ple-gguf", str(shard1),
            "--n-expert", "256",
            "--max-context", "131072",
            "--cache-disk", str(model_dir / "cache_disk"),
            "--cache-disk-size", "16384",
        ] + (["--vision"] if has_vision else []),
        "cwd": str(root / "vendor" / "strata"),
        "tokenizer": str(tok_dir),
        "model_name": "swift-1.5-k256-healed" if healed_gguf.exists() else "swift-1.5-k256",
        "log": "/tmp/strata-k256-serve.log",
        "lib_dirs": [],
        "port": 8080,
        "backend": "metal",
        "fit_max_tokens": True,
        "env": {
            "STRATA_METAL_ARENA_GB": "6.0",
            "STRATA_METAL_MLOCK_DENSE": "1",
            "STRATA_METAL_WIRE_ARENA": "1",
            "STRATA_METAL_EXPERT_PROFILE": str(out_prof),
            "STRATA_METAL_MOCE_BONUS": "0.85",
            "STRATA_METAL_EXPERT_SPLIT": "1",
            "STRATA_METAL_SHEXP_OVERLAP": "1",
            "STRATA_METAL_SPLIT_TURN": "1",
            "STRATA_METAL_PRED_Q8": "1",
            "STRATA_METAL_ROUTER_BF16": "1",
            "STRATA_METAL_MOE_DN_TILE": "8,2,4",
            "STRATA_METAL_PREFETCH": "6",
            "STRATA_METAL_PREFETCH2": "2",
            "STRATA_METAL_IO_THREADS": "8",
            "STRATA_METAL_IO_URGENT": "2",
            "STRATA_ANTI_SLOP_PENALTY": "6.0",
        },
        "sampling": {
            "temperature": 0.6,
            "top_p": 0.95,
            "top_k": 40,
            "min_p": 0.05,
            "presence_penalty": 0.4,
            "repeat_penalty": 1.08,
        },
        "reasoning_effort": "low",
    }
    if has_vision:
        cfg["vram_elastic"] = True
        cfg["vision"] = {
            "exe": str(vexe),
            "mmproj": str(mmproj_p),
            "model": str(vocab_out),
            "gpu": True,
            "max_tokens": 1024,
            "elastic": True,
            "elastic_idle_s": 60,
        }
    (model_dir / "strata-k256.json").write_text(json.dumps(cfg, indent=2) + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())