#!/usr/bin/env python3
"""Fast 24-worker HTTP Range fetcher + packer for Qwen3.8-Flash-Next MTP (mtp-q2_0.gguf)."""
from __future__ import annotations
import concurrent.futures, hashlib, json, os, shutil, subprocess, sys, time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "vendor" / "strata" / "tools"))
import mtp_fetch as mf

OUT_DIR = ROOT / "models" / "K512_iq2_xs" / "mtp_src"
MTP_GGUF = ROOT / "models" / "K512_iq2_xs" / "mtp-q2_0.gguf"
RT_DIR = ROOT / "models" / "K512_iq2_xs" / "mtp_rt"

def main():
    t0 = time.monotonic()
    mf.resolve_repo()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    inv_path = OUT_DIR / "mtp-inventory.json"
    if inv_path.exists():
        rows = json.loads(inv_path.read_text(encoding="utf-8"))["tensors"]
    else:
        print("Fetching MTP inventory headers...", flush=True)
        rows = mf.inventory(str(OUT_DIR))
    tdir = OUT_DIR / "tensors"
    tdir.mkdir(parents=True, exist_ok=True)

    # Build 16 MiB chunk tasks for all missing/incomplete tensors
    chunk_sz = 16 << 20
    to_fetch = []
    for r in rows:
        p = tdir / (r["name"] + ".bin")
        want = mf.SHA256.get(r["name"]) if mf.pinned() else None
        if p.exists() and p.stat().st_size == r["bytes"]:
            if want is None or mf.sha256_of(str(p)) == want:
                continue
        # Pre-allocate file
        with open(p, "wb") as f:
            f.truncate(r["bytes"])
        pos = r["start"]
        while pos <= r["end"]:
            end = min(pos + chunk_sz - 1, r["end"])
            to_fetch.append((r["name"], r["shard"], p, pos - r["start"], pos, end))
            pos = end + 1

    total_dl = sum(e - s + 1 for _, _, _, _, s, e in to_fetch)
    print(f"Downloading {len(to_fetch)} chunks ({total_dl/1e9:.2f} GB) with 24 workers...", flush=True)

    def dl_chunk(task):
        name, shard, p, file_off, start, end = task
        data = mf.get(mf.REPO + shard, start, end)
        fd = os.open(str(p), os.O_WRONLY)
        try:
            os.pwrite(fd, data, file_off)
        finally:
            os.close(fd)
        return len(data)

    done_bytes = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=24) as ex:
        futs = [ex.submit(dl_chunk, t) for t in to_fetch]
        for idx, fut in enumerate(concurrent.futures.as_completed(futs), 1):
            done_bytes += fut.result()
            if idx % 20 == 0 or idx == len(futs):
                dt = time.monotonic() - t0
                print(f"  [{idx}/{len(futs)}] {done_bytes/1e9:.2f}/{total_dl/1e9:.2f} GB ({done_bytes/1e6/max(dt,1e-3):.1f} MB/s)", flush=True)

    # Verify SHA256 and write mtp-manifest.json
    manifest = []
    for r in rows:
        p = tdir / (r["name"] + ".bin")
        digest = mf.sha256_of(str(p))
        want = mf.SHA256.get(r["name"]) if mf.pinned() else None
        if want and digest != want:
            raise RuntimeError(f"SHA256 mismatch on {r['name']}: {digest} != {want}")
        manifest.append(dict(r, file=os.path.relpath(str(p), str(OUT_DIR)), sha256=digest))
    (OUT_DIR / "mtp-manifest.json").write_text(json.dumps(manifest, indent=1), encoding="utf-8")
    print(f"All 31 MTP tensors verified in {time.monotonic()-t0:.1f}s! Packing mtp-q2_0.gguf...", flush=True)

    # Run mtp_pack.py
    env = dict(os.environ, STRATA_GGUF_PY=str((ROOT / "vendor" / "llama.cpp" / "gguf-py").resolve()))
    subprocess.run(
        [sys.executable, str(ROOT / "vendor" / "strata" / "tools" / "mtp_pack.py"),
         "--src", str(OUT_DIR), "--experts", "q2_0", "--out", str(MTP_GGUF)],
        env=env,
        check=True,
    )
    RT_DIR.mkdir(parents=True, exist_ok=True)
    shutil.copy2(ROOT / "vendor" / "strata" / "data" / "draft_vocab_en.bin", RT_DIR / "draft_vocab.bin")
    # Clean up raw 4.9 GB BF16 tensor dir after packing
    shutil.rmtree(OUT_DIR, ignore_errors=True)
    print(f"COMPLETE in {time.monotonic()-t0:.1f}s: {MTP_GGUF} ({MTP_GGUF.stat().st_size/1e6:.1f} MB)", flush=True)

if __name__ == "__main__":
    main()