"""E6/E2: genuine two-node cache transfer over the LAN (laptop -> M3U96b worker) vs text transfer vs local recompute.
Usage: python scripts/e6_network.py --host 192.168.2.85 --port 9430 --model qwen3b --docs 6 --lengths 2000 8000 16000 --repeats 3
Requires the worker: on the remote node `python -m wp30 worker --port 9430 --model <repo>` (same checkpoint revision).
Loopback mode (--host 127.0.0.1) is labelled 'loopback' and is NOT a network measurement."""
import argparse, json, time, statistics, subprocess, sys, socket
import yaml, numpy as np, mlx.core as mx
from wp30.util import ROOT, read_json, utc_now, append_jsonl, sha256_text
from wp30.data.prepare import load_doc, load_key
from wp30.backends.mlx_backend import MLXModel
from wp30.cache.sparse import quantize_cache, cache_tensor_bytes, to_sparse_all
from wp30.cache.serialize import build_manifest, serialize_cache, codec_compress
from wp30.transport import Client
from wp30.prompts import SYSTEM_RECEIVER, receiver_user_text
from wp30.scoring import score_synthetic

ap = argparse.ArgumentParser()
ap.add_argument("--host", required=True); ap.add_argument("--port", type=int, default=9430); ap.add_argument("--model", default="qwen3b")
ap.add_argument("--docs", type=int, default=6); ap.add_argument("--lengths", type=int, nargs="+", default=[2000, 8000, 16000]); ap.add_argument("--repeats", type=int, default=3)
ap.add_argument("--label", default=None); ap.add_argument("--warmup", type=int, default=1); ap.add_argument("--split", default="dev"); ap.add_argument("--out", default="runs/raw/E6_network.jsonl")
a = ap.parse_args()
label = a.label or ("loopback" if a.host.startswith("127.") else "lan")
mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]
m = MLXModel(mcfg[a.model]["repo"])
cl = Client(a.host, a.port)
h, _, meta = cl.call({"cmd": "ping"})
assert h["config_hash"] == m.info.config_hash, "remote model differs"
print("remote ok", h["model_revision"], "rtt_ms", round(meta["roundtrip_ms"], 2))
# raw link bandwidth probe with an incompressible 64 MiB payload (frame echo is not implemented; use run_text no-op? -> use a dedicated ping with payload)
man = read_json(ROOT / "data/manifests/splits.json")
items = [it for it in man["items"] if it["family"] == "synthetic" and it["split"] == a.split and it.get("target_tokens") in a.lengths]
by_len = {L: [it for it in items if it["target_tokens"] == L][: a.docs] for L in a.lengths}
out = ROOT / a.out
# warm-up: compile kernels on both ends (not recorded)
if a.warmup:
    _w = m.chat_prompt_tokens([{"role": "user", "content": "warm up " * 200}])
    for _ in range(a.warmup):
        _c, _ = m.prefill(_w); m.generate_from_cache(_c, _w[-3:], max_tokens=4); cl.call({"cmd": "run_text", "token_ids": _w, "max_tokens": 4})
        _man = build_manifest(m.info, _w, to_sparse_all(_c), "warm", "B1", {}); cl.call({"cmd": "run_cache", "codec": "none", "prefix_hash": _man["prefix_token_hash"], "query_ids": _w[-3:], "max_tokens": 4}, serialize_cache(to_sparse_all(_c), _man))
for L, its in by_len.items():
    for it in its:
        doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")
        q = doc["questions"][0]; qk = key["questions"][0]
        full = m.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Document", doc["text"], q["question"], "synthetic")}])
        pre_txt = m.tokenizer.apply_chat_template([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + doc["text"] + "\n\n"}], add_generation_prompt=False, tokenize=False)
        pre = m.encode(pre_txt); n = 0
        while n < min(len(pre), len(full)) and pre[n] == full[n]: n += 1
        pre, query = full[:n], full[n:]
        for rep in range(a.repeats):
            # (1) local sender prefill + serialize + transfer cache + remote load + remote query  [warm remote cache scenario: cache moved]
            for variant in [("B1", 8 if False else None, "none"), ("B1", None, "zstd"), ("B2", 8, "none"), ("B2", 4, "none")]:
                meth, bits, codec = variant
                t0 = time.perf_counter(); c, tp = m.prefill(pre)
                c2 = to_sparse_all(c) if bits is None else quantize_cache(c, bits)
                mx.eval([x for cc in c2 for x in (cc.keys if isinstance(cc.keys, tuple) else (cc.keys,))]); mx.synchronize()
                t1 = time.perf_counter()
                manf = build_manifest(m.info, pre, c2, it["document_id"], meth, {"bits": bits})
                payload = serialize_cache(c2, manf); t2 = time.perf_counter()
                comp, cms = codec_compress(payload, codec) if codec != "none" else (payload, 0.0)
                hdr = {"cmd": "run_cache", "codec": codec, "prefix_hash": manf["prefix_token_hash"], "query_ids": query, "max_tokens": 32}
                rh, _, tm = cl.call(hdr, comp)
                rec = {"experiment_id": "E6_network", "label": label, "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,
                       "method": meth, "bits": bits, "codec": codec, "status": rh.get("status"), "answer": rh.get("text"), "correct": score_synthetic(rh.get("text", ""), qk)["correct"] if rh.get("text") else None,
                       "cache_tensor_bytes": cache_tensor_bytes(c2), "serialized_bytes": len(payload), "transport_bytes": tm["transport_bytes"], "sender_prefill_ms": tp["prefill_ms"], "compress_ms": (t1 - t0) * 1000 - tp["prefill_ms"],
                       "serialize_ms": (t2 - t1) * 1000, "codec_ms": cms, "send_ms": tm["send_ms"], "roundtrip_ms": tm["roundtrip_ms"], "remote_recv_ms": rh.get("recv_ms"), "remote_decompress_ms": rh.get("decompress_ms"),
                       "remote_load_ms": rh.get("load_ms"), "remote_ttft_ms": rh.get("ttft_ms"), "remote_decode_ms": rh.get("decode_ms"), "remote_server_ms": rh.get("server_ms"), "remote_peak_bytes": rh.get("peak_bytes"),
                       "effective_gbps": (tm["transport_bytes"] * 8 / 1e9) / max(1e-9, (tm["roundtrip_ms"] - (rh.get("server_ms") or 0)) / 1000)}
                append_jsonl(out, rec)
                print(f"{label} L={L} {meth} bits={bits} codec={codec} rep={rep} bytes={len(comp):>10d} send={tm['send_ms']:7.1f}ms rt={tm['roundtrip_ms']:7.1f}ms remote_load={rh.get('load_ms',0):6.1f} ttft={rh.get('ttft_ms',0):6.1f} eff={rec['effective_gbps']:.2f}Gb/s corr={rec['correct']}")
            # (2) text transfer: send token ids, remote does full prefill [cold remote scenario]
            rh, _, tm = cl.call({"cmd": "run_text", "token_ids": full, "max_tokens": 32})
            rec = {"experiment_id": "E6_network", "label": label, "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,
                   "method": "text_remote_prefill", "bits": None, "codec": "none", "status": "success", "answer": rh.get("text"), "correct": score_synthetic(rh.get("text", ""), qk)["correct"],
                   "serialized_bytes": len(doc["text"].encode()), "transport_bytes": tm["transport_bytes"], "send_ms": tm["send_ms"], "roundtrip_ms": tm["roundtrip_ms"], "remote_prefill_ms": rh.get("prefill_ms"),
                   "remote_ttft_ms": rh.get("ttft_ms"), "remote_decode_ms": rh.get("decode_ms"), "remote_server_ms": rh.get("server_ms"), "remote_peak_bytes": rh.get("peak_bytes")}
            append_jsonl(out, rec)
            print(f"{label} L={L} text_remote rep={rep} bytes={tm['transport_bytes']:>10d} rt={tm['roundtrip_ms']:7.1f}ms remote_prefill={rh.get('prefill_ms',0):7.1f} ttft={rh.get('ttft_ms',0):6.1f}")
            # (3) local processing (no transfer): run where the data is
            c, tp = m.prefill(pre); g = m.generate_from_cache(c, query, max_tokens=32)
            append_jsonl(out, {"experiment_id": "E6_network", "label": "local", "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,
                               "method": "local_prefill", "status": "success", "answer": g["text"], "correct": score_synthetic(g["text"], qk)["correct"], "local_prefill_ms": tp["prefill_ms"], "local_ttft_ms": g["ttft_ms"], "local_decode_ms": g["decode_ms"]})
cl.close()
print("done ->", out)