#!/usr/bin/env python3
"""
Fast pipelined parallel HTTP Range builder for unpruned K=512 IQ2_XS:
  - models/K512_iq2_xs/K512-iq2_xs-metal.gguf         (3.90 GB dense-only GGUF)
  - models/K512_iq2_xs/K512-iq2_xs-metal.gguf.experts (36.23 GB STRX contiguous 512-expert file)
Uses 24 concurrent workers, overlapping all dense groups and pipelining G+U+D + 2-layer lookahead.
"""
from __future__ import annotations

import concurrent.futures
import fcntl
import json
import os
import struct
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "vendor" / "strata" / "metal" / "tools"))
import metal_repack as MR  # noqa: E402

MR.TYPE_SIZE[29] = (56, 256)

ALIGN = 16384
NL, NE = 48, 512

U_IQ2_S1 = "https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ2_XS-00001-of-00002.gguf"
U_IQ2_S2 = "https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ2_XS-00002-of-00002.gguf"
U_XXS_S1 = "https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00001-of-00002.gguf"
U_XXS_S2 = "https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00002-of-00002.gguf"


def set_nocache(fd: int) -> None:
    try:
        fcntl.fcntl(fd, getattr(fcntl, "F_NOCACHE", 48), 1)
    except OSError:
        pass


def fetch_range(url: str, start: int, length: int, retries: int = 8) -> bytes:
    end = start + length - 1
    for attempt in range(retries):
        try:
            req = urllib.request.Request(
                url,
                headers={"User-Agent": "Mozilla/5.0", "Range": f"bytes={start}-{end}"},
            )
            with urllib.request.urlopen(req, timeout=60) as r:
                data = r.read()
            if len(data) != length:
                raise IOError(f"Short read: got {len(data)} != {length}")
            return data
        except Exception:
            if attempt + 1 == retries:
                raise
            time.sleep(0.8 * (attempt + 1))
    raise RuntimeError("unreachable")


def start_blob(ex: concurrent.futures.ThreadPoolExecutor, url: str, abs_off: int, total_bytes: int, chunk_mb: int = 16):
    chunk = chunk_mb * 1024 * 1024
    futs = []
    for off in range(0, total_bytes, chunk):
        ln = min(chunk, total_bytes - off)
        futs.append((off, ln, ex.submit(fetch_range, url, abs_off + off, ln)))
    return total_bytes, futs


def collect_blob(job) -> bytearray:
    total_bytes, futs = job
    out = bytearray(total_bytes)
    for off, ln, fut in futs:
        out[off : off + ln] = fut.result()
    return out


def ensure_hdr(path: Path, url: str, size_mb: int) -> MR.Reader:
    if not path.exists() or path.stat().st_size < size_mb * 1024 * 1024:
        path.write_bytes(fetch_range(url, 0, size_mb * 1024 * 1024))
    rd = MR.Reader(path)
    rd.f.close()
    return rd


def pack_str(s: str) -> bytes:
    b = s.encode("utf-8")
    return struct.pack("<Q", len(b)) + b


def main() -> int:
    t0 = time.monotonic()
    out_dir = ROOT / "models" / "K512_iq2_xs"
    out_dir.mkdir(parents=True, exist_ok=True)
    gguf_out = out_dir / "K512-iq2_xs-metal.gguf"
    exp_out = out_dir / "K512-iq2_xs-metal.gguf.experts"
    state_path = out_dir / "build_state.json"

    r1 = ensure_hdr(Path("/tmp/iq2xs_s1.hdr"), U_IQ2_S1, 12)
    r2 = ensure_hdr(Path("/tmp/iq2xs_s2.hdr"), U_IQ2_S2, 2)
    rx1 = ensure_hdr(Path("/tmp/iq3xxs_s1.hdr"), U_XXS_S1, 12)
    rx2 = ensure_hdr(Path("/tmp/iq3xxs_s2.hdr"), U_XXS_S2, 2)

    tx_map = {t["name"]: (t, U_XXS_S1 if t in rx1.tensors else U_XXS_S2) for t in rx1.tensors + rx2.tensors}
    iq2_map = {t["name"]: (t, U_IQ2_S1 if t in r1.tensors else U_IQ2_S2) for t in r1.tensors + r2.tensors}

    dense = []
    for src_r, url in ((r1, U_IQ2_S1), (r2, U_IQ2_S2)):
        for t in src_r.tensors:
            if t["name"] == MR.PLE or "_exps." in t["name"]:
                continue
            tc = dict(t)
            tc["url"] = url
            dense.append(tc)

    experts = []
    for l in range(NL):
        for k in ("gate", "up", "down"):
            nm = f"blk.{l}.ffn_{k}_exps.weight"
            if l in (8, 13, 37) and k in ("gate", "up"):
                t, u = tx_map[nm]
            else:
                t, u = iq2_map[nm]
            tc = dict(t)
            tc["url"] = u
            experts.append(tc)

    order = dense + experts
    kvs = [(k, raw) for k, raw in r1.kv if not k.startswith("split.") and k != "general.alignment"]
    kvs.append(("general.alignment", pack_str("general.alignment") + struct.pack("<II", 4, ALIGN)))

    hdr = bytearray(struct.pack("<4sIQQ", b"GGUF", r1.version, len(order), len(kvs)))
    for _, raw in kvs:
        hdr += raw

    off = 0
    for t in order:
        t["out_off"] = off
        off = (off + t["bytes"] + ALIGN - 1) // ALIGN * ALIGN

    dense_end_off = (dense[-1]["out_off"] + dense[-1]["bytes"] + ALIGN - 1) // ALIGN * ALIGN
    for t in order:
        hdr += pack_str(t["name"])
        hdr += struct.pack(f"<I{len(t['dims'])}QIQ", len(t["dims"]), *t["dims"], t["type"], t["out_off"])

    data0 = (len(hdr) + ALIGN - 1) // ALIGN * ALIGN
    hdr += b"\0" * (data0 - len(hdr))
    gguf_final_size = data0 + dense_end_off

    with concurrent.futures.ThreadPoolExecutor(max_workers=24) as ex:
        if not gguf_out.exists() or gguf_out.stat().st_size != gguf_final_size:
            tmp_gguf = gguf_out.with_name(gguf_out.name + ".part")
            fd = os.open(str(tmp_gguf), os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o644)
            set_nocache(fd)
            try:
                os.pwrite(fd, bytes(hdr), 0)
                os.ftruncate(fd, gguf_final_size)
                groups = []
                for t in dense:
                    if (
                        groups
                        and groups[-1]["url"] == t["url"]
                        and t["abs"] - (groups[-1]["start"] + groups[-1]["len"]) <= 1024 * 1024
                    ):
                        groups[-1]["len"] = (t["abs"] + t["bytes"]) - groups[-1]["start"]
                        groups[-1]["tensors"].append(t)
                    else:
                        groups.append({"url": t["url"], "start": t["abs"], "len": t["bytes"], "tensors": [t]})

                print(f"Downloading {len(dense)} dense tensors across {len(groups)} ranges (24 workers)...", flush=True)
                t_d0 = time.monotonic()
                done_b = 0
                tot_dense_b = sum(t["bytes"] for t in dense)
                # Submit in sliding window of 16 groups so all 24 workers stay busy without excess RAM
                q = []
                nxt = 0
                while nxt < min(16, len(groups)):
                    g = groups[nxt]
                    q.append((nxt, g, start_blob(ex, g["url"], g["start"], g["len"], chunk_mb=16)))
                    nxt += 1
                while q:
                    gi, g, job = q.pop(0)
                    if nxt < len(groups):
                        gn = groups[nxt]
                        q.append((nxt, gn, start_blob(ex, gn["url"], gn["start"], gn["len"], chunk_mb=16)))
                        nxt += 1
                    blob = collect_blob(job)
                    for t in g["tensors"]:
                        rel = t["abs"] - g["start"]
                        os.pwrite(fd, memoryview(blob)[rel : rel + t["bytes"]], data0 + t["out_off"])
                        done_b += t["bytes"]
                    del blob
                    if (gi + 1) % 20 == 0 or gi + 1 == len(groups):
                        el = time.monotonic() - t_d0
                        print(
                            f"  [dense {gi+1:3d}/{len(groups)}] {done_b/1e9:.2f}/{tot_dense_b/1e9:.2f} GB "
                            f"({done_b/1e6/max(el,1e-3):.1f} MB/s)",
                            flush=True,
                        )
                os.fsync(fd)
            finally:
                os.close(fd)
            os.replace(tmp_gguf, gguf_out)
            print(f"Saved {gguf_out.name} ({gguf_out.stat().st_size} bytes)!", flush=True)
        else:
            print(f"{gguf_out.name} already complete ({gguf_final_size} bytes).", flush=True)

        lay = []
        for l in range(NL):
            g, u, d = experts[l * 3], experts[l * 3 + 1], experts[l * 3 + 2]
            lay.append((g, u, d, g["bytes"] // NE, u["bytes"] // NE, d["bytes"] // NE))
        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
        exp_final_size = at

        exp_head = struct.pack("<4sIIIQ", b"STRX", 1, NL, NE, gguf_final_size).ljust(64, b"\0")
        for (g, u, d, gs, us, ds), o in zip(lay, offs):
            exp_head += struct.pack("<QQQQ", o, gs, us, ds)
        exp_head = exp_head.ljust(exp_data0, b"\0")

        done_layers = set()
        if state_path.exists() and exp_out.exists():
            try:
                done_layers = set(json.loads(state_path.read_text()).get("done_layers", []))
            except Exception:
                done_layers = set()

        todo = [l for l in range(NL) if l not in done_layers]
        fd_e = os.open(str(exp_out), os.O_CREAT | os.O_RDWR, 0o644)
        set_nocache(fd_e)
        try:
            os.pwrite(fd_e, exp_head, 0)
            os.ftruncate(fd_e, exp_final_size)
            t_e0 = time.monotonic()
            bytes_streamed = 0

            # Pipeline 2 layers ahead (G, U, D all submitted together)
            active = {}
            def submit_layer(li: int):
                g, u, d, gs, us, ds = lay[li]
                active[li] = (
                    start_blob(ex, g["url"], g["abs"], g["bytes"], chunk_mb=16),
                    start_blob(ex, u["url"], u["abs"], u["bytes"], chunk_mb=16),
                    start_blob(ex, d["url"], d["abs"], d["bytes"], chunk_mb=16),
                )

            nxt_idx = 0
            while nxt_idx < min(2, len(todo)):
                submit_layer(todo[nxt_idx])
                nxt_idx += 1

            for l in todo:
                tl0 = time.monotonic()
                jg, ju, jd = active.pop(l)
                if nxt_idx < len(todo):
                    submit_layer(todo[nxt_idx])
                    nxt_idx += 1
                g, u, d, gs, us, ds = lay[l]
                o = offs[l]
                G = collect_blob(jg)
                U = collect_blob(ju)
                D = collect_blob(jd)

                layer_bytes = 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)
                )
                os.pwrite(fd_e, layer_bytes, o)
                nb = len(layer_bytes)
                del G, U, D, layer_bytes
                bytes_streamed += nb
                done_layers.add(l)
                state_path.write_text(json.dumps({"done_layers": sorted(done_layers)}))
                el = time.monotonic() - t_e0
                print(
                    f"  [experts layer {l+1:2d}/{NL}] type=({g['type']},{u['type']},{d['type']}) "
                    f"size={nb/1e6:.1f} MB in {time.monotonic()-tl0:.1f}s | "
                    f"avg={bytes_streamed/1e6/max(el,1e-3):.1f} MB/s ({len(done_layers)}/{NL} done)",
                    flush=True,
                )
            os.fsync(fd_e)
        finally:
            os.close(fd_e)

        # Optional standalone 1-tensor PLE GGUF (per_layer_token_embd.weight, 28.80 GB IQ4_NL)
        ple_k256 = ROOT / "models" / "K256_reap_l2_sum" / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf"
        ple_out = out_dir / "ple-iq4_nl.gguf"
        ple_path = ple_k256 if ple_k256.exists() else ple_out
        if ("--with-ple" in sys.argv or not ple_path.exists()):
            ple_t = next(t for t in r1.tensors if t["name"] == MR.PLE)
            ple_hdr = bytearray(struct.pack("<4sIQQ", b"GGUF", 3, 1, 0))
            ple_hdr += pack_str(ple_t["name"])
            ple_hdr += struct.pack(f"<I{len(ple_t['dims'])}QIQ", len(ple_t["dims"]), *ple_t["dims"], ple_t["type"], 0)
            ple_data0 = (len(ple_hdr) + 31) // 32 * 32
            ple_hdr += b"\0" * (ple_data0 - len(ple_hdr))
            ple_final_size = ple_data0 + ple_t["bytes"]
            if not ple_out.exists() or ple_out.stat().st_size != ple_final_size:
                print(f"Downloading standalone PLE table ({ple_t['bytes']/1e9:.2f} GB) to {ple_out.name}...", flush=True)
                tmp_ple = ple_out.with_name(ple_out.name + ".part")
                fd_p = os.open(str(tmp_ple), os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o644)
                set_nocache(fd_p)
                try:
                    os.pwrite(fd_p, bytes(ple_hdr), 0)
                    os.ftruncate(fd_p, ple_final_size)
                    win = 64 * 1024 * 1024
                    offs_p = list(range(0, ple_t["bytes"], win))
                    q_p = []
                    nxt_p = 0
                    while nxt_p < min(8, len(offs_p)):
                        off_p = offs_p[nxt_p]
                        ln_p = min(win, ple_t["bytes"] - off_p)
                        q_p.append((off_p, start_blob(ex, U_IQ2_S1, ple_t["abs"] + off_p, ln_p, chunk_mb=16)))
                        nxt_p += 1
                    while q_p:
                        off_p, job_p = q_p.pop(0)
                        if nxt_p < len(offs_p):
                            off2 = offs_p[nxt_p]
                            ln2 = min(win, ple_t["bytes"] - off2)
                            q_p.append((off2, start_blob(ex, U_IQ2_S1, ple_t["abs"] + off2, ln2, chunk_mb=16)))
                            nxt_p += 1
                        blob_p = collect_blob(job_p)
                        os.pwrite(fd_p, blob_p, ple_data0 + off_p)
                        del blob_p
                    os.fsync(fd_p)
                finally:
                    os.close(fd_p)
                os.replace(tmp_ple, ple_out)
            ple_path = ple_out

    # Extract tokenizer from GGUF metadata and set default reasoning_effort to 'low'
    tok_dir = out_dir / "tokenizer"
    if "--no-tokenizer" not in sys.argv and not (tok_dir / "vocab.json").exists():
        sys.path.insert(0, str(ROOT / "vendor" / "strata" / "tools"))
        import strata_tokenizer as ST
        ST.extract(str(gguf_out), str(out_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")

    # Write 0-tensor vocab.gguf for strata-vision (mtmd vocab_only header)
    vocab_out = out_dir / "vocab.gguf"
    if not vocab_out.exists():
        rd_v = MR.Reader(gguf_out)
        rd_v.f.close()
        vocab_hdr = struct.pack("<4sIQQ", b"GGUF", rd_v.version, 0, len(rd_v.kv)) + b"".join(raw for _, raw in rd_v.kv)
        vocab_out.write_bytes(vocab_hdr)
        del rd_v, vocab_hdr

    # Optional vision projector (mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf, 907.54 MB)
    mmproj_out = out_dir / "mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf"
    mmproj_url = "https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf"
    if "--no-vision" not in sys.argv and (not mmproj_out.exists() or mmproj_out.stat().st_size < 900_000_000):
        print("Downloading mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf (907.5 MB)...", flush=True)
        mm_sz = head_size(mmproj_url)
        tmp_mm = mmproj_out.with_name(mmproj_out.name + ".part")
        fd_m = os.open(str(tmp_mm), os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o644)
        try:
            os.ftruncate(fd_m, mm_sz)
            with ThreadPoolExecutor(max_workers=16) as ex_m:
                jb_m = start_blob(ex_m, mmproj_url, 0, mm_sz, chunk_mb=16)
                os.pwrite(fd_m, collect_blob(jb_m), 0)
            os.fsync(fd_m)
        finally:
            os.close(fd_m)
        os.replace(tmp_mm, mmproj_out)

    vexe = ROOT / "vendor" / "strata" / "build-vision" / "bin" / "strata-vision"
    has_vision = vexe.exists() and mmproj_out.exists()

    # Write strata-k512-iq2xs.json using checkout ROOT
    cfg_path = out_dir / "strata-k512-iq2xs.json"
    cfg = {
        "exe": str(ROOT / "vendor" / "strata" / "build-metal" / "strata_window_cli"),
        "args": [
            "--metal-gguf", str(gguf_out),
            "--ple-gguf", str(ple_path),
            "--mtp-gguf", str(out_dir / "mtp-q2_0.gguf"),
            "--n-expert", "512",
            "--max-context", "131072",
            "--spec", "3",
            "--spec-min-p", "0.35",
            "--suffix-draft", "3",
            "--cache-disk", str(out_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-k512-iq2xs",
        "log": "/tmp/strata-k512-serve.log",
        "lib_dirs": [],
        "port": 8080,
        "backend": "metal",
        "fit_max_tokens": True,
        "sampling": {
            "temperature": 0.6,
            "top_p": 0.95,
            "top_k": 40,
            "min_p": 0.05,
            "presence_penalty": 0.4,
            "repeat_penalty": 1.08,
        },
        "env": {
            "STRATA_METAL_ARENA_GB": "6.5",
            "STRATA_METAL_MLOCK_DENSE": "1",
            "STRATA_METAL_WIRE_ARENA": "1",
            "STRATA_METAL_EXPERT_PROFILE": str(ROOT / "vendor" / "strata" / "data" / "expert-profile.bin"),
            "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": "9.0",
            "STRATA_ANTI_SLOP_HONEST": "6.0",
            "STRATA_ANTI_SLOP_MD": "4.0",
            "STRATA_ANTI_SLOP_COLON": "3.0",
            "STRATA_METAL_PROMPT_CHUNK": "2048",
            "STRATA_METAL_PREFETCH_DRAFT": "0",
            "STRATA_METAL_PUBLISH_PREFETCH": "1",
            "STRATA_METAL_MTP_SLOP": "1",
            "STRATA_METAL_DRAFT_MOCE": "1",
        },
        "reasoning_effort": "low",
    }
    if has_vision:
        cfg["vram_elastic"] = True
        cfg["vision"] = {
            "exe": str(vexe),
            "mmproj": str(mmproj_out),
            "model": str(vocab_out),
            "gpu": True,
            "max_tokens": 1024,
            "elastic": True,
            "elastic_idle_s": 60,
        }
    cfg_path.write_text(json.dumps(cfg, indent=2) + "\n", encoding="utf-8")

    print(
        f"COMPLETE in {time.monotonic()-t0:.1f}s! "
        f"GGUF={gguf_out.stat().st_size/1e9:.3f} GB, Experts={exp_out.stat().st_size/1e9:.3f} GB",
        flush=True,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())