# E21: long-context capacity + needle sweep on RTX 3080 10GB (supporting measurement, no novelty claim).
import json, os, subprocess, sys, time, random, urllib.request
from huggingface_hub import hf_hub_download
ROOT = r"C:\AGI\agi-research"; OUT = os.path.join(ROOT, r"results\quality\e21_longctx.jsonl")
LOG = os.path.join(ROOT, r"bench\E21.log"); MD = r"C:\AGI\gguf\q35"; SRV = r"C:\AGI\llama.cpp\llama-server.exe"
PORT = 8097; URL = f"http://127.0.0.1:{PORT}"
def log(m):
    line = f"{time.strftime('%Y-%m-%d %H:%M:%S')} {m}"; print(line, flush=True)
    open(LOG, "a", encoding="utf-8").write(line + "\n")
def http(path, body=None, timeout=60):
    req = urllib.request.Request(URL + path, data=None if body is None else json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r: return json.loads(r.read())
def vram():
    try: return int(subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"]).decode().split()[0])
    except Exception: return -1
def shared_mb(pid):
    try:
        o = subprocess.check_output(["powershell", "-NoProfile", "-Command", f"(Get-Counter '\\GPU Process Memory(pid_{pid}*)\\Shared Usage').CounterSamples | Measure-Object CookedValue -Sum | % Sum"], timeout=60).decode().strip()
        return round(float(o) / 2**20, 1) if o else 0.0
    except Exception: return -1
CONFIGS = [("2B", "unsloth/Qwen3.5-2B-GGUF", "Qwen3.5-2B-Q4_K_M.gguf", [131072, 262144, 524288, 1048576], ("q4_0", "q4_0")),
           ("4B", "unsloth/Qwen3.5-4B-GGUF", "Qwen3.5-4B-Q4_K_M.gguf", [131072, 262144, 524288], ("q4_0", "q4_0")),
           ("9B", "unsloth/Qwen3.5-9B-GGUF", "Qwen3.5-9B-Q4_K_M.gguf", [131072, 262144], ("q8_0", "q4_0"))]
NEEDLES = [("literal", "The secret passcode for the blue vault is 7392-ALPHA.", "What is the secret passcode for the blue vault? Answer with the passcode only.", "7392-ALPHA"),
           ("paraphrase", "Yuki actually lives right next to the Semperoper.", "Which character has been to Dresden? Answer with the name only.", "Yuki"),
           ("code", "In the helper module, def compute_flux_ratio(): return 4417", "What value does the function compute_flux_ratio return? Answer with the number only.", "4417")]
def filler_paragraphs():
    ps = []
    for l in open(os.path.join(ROOT, r"e20\data\sft.jsonl"), encoding="utf-8"):
        d = json.loads(l); t = " ".join(str(v) for v in d.values() if isinstance(v, str))
        ps += [p.strip() for p in t.split("\n\n") if len(p.strip()) > 80]
    return ps
def main():
    os.makedirs(MD, exist_ok=True); paras = filler_paragraphs(); log(f"E21 start, {len(paras)} filler paragraphs")
    done = set()
    if os.path.exists(OUT):
        for l in open(OUT, encoding="utf-8"):
            d = json.loads(l); done.add((d["model"], d["ctx"]))
    for name, repo, fn, ctxs, (ctk, ctv) in CONFIGS:
        path = hf_hub_download(repo, fn, local_dir=MD); log(f"{name} model at {path}")
        for ctx in ctxs:
            if (name, ctx) in done: continue
            rec = {"model": name, "gguf": fn, "ctx": ctx, "ctk": ctk, "ctv": ctv}
            args = [SRV, "--port", str(PORT), "-m", path, "-ngl", "99", "-fa", "on", "-c", str(ctx), "-np", "1", "-ctk", ctk, "-ctv", ctv, "-t", "8", "--no-mmap"]
            if ctx > 262144: f = ctx // 262144; args += ["--rope-scaling", "yarn", "--rope-scale", str(f), "--yarn-orig-ctx", "262144"]; rec["yarn"] = f
            errf = open(os.path.join(ROOT, rf"bench\e21_srv_{name}_{ctx}.err"), "w")
            p = subprocess.Popen(args, stdout=errf, stderr=errf); t0 = time.time(); ok = False
            while time.time() - t0 < 900 and p.poll() is None:
                try: ok = http("/health", timeout=5).get("status") == "ok"
                except Exception: ok = False
                if ok: break
                time.sleep(3)
            rec["vram_loaded_mb"] = vram(); rec["shared_loaded_mb"] = shared_mb(p.pid)
            if not ok:
                rec["status"] = "load_failed"; rec["exit"] = p.poll(); log(f"{name} {ctx} load failed"); p.kill(); errf.close()
                open(OUT, "a").write(json.dumps(rec) + "\n"); continue
            trials = []; peak = rec["vram_loaded_mb"]; shared_peak = rec["shared_loaded_mb"]
            for depth in (0.1, 0.5, 0.9):
                rnd = random.Random(int(ctx * depth)); target = ctx - 1200; text = []; ntok = 0
                while ntok < target:
                    chunk = [paras[rnd.randrange(len(paras))] for _ in range(400)]
                    ntok += len(http("/tokenize", {"content": "\n\n".join(chunk)}, timeout=600)["tokens"]); text += chunk
                while ntok > target and text:
                    rm = text.pop(); ntok -= max(1, len(rm) // 4)
                pos = int(len(text) * depth)
                for i, (_, s, _, _) in enumerate(NEEDLES): text.insert(pos + i * 3, s)
                hay = "\n\n".join(text)
                for kind, _, q, ans in NEEDLES:
                    msgs = [{"role": "user", "content": "Read the document below and answer the question at the end.\n\n<document>\n" + hay + "\n</document>\n\nQuestion: " + q}]
                    t1 = time.time()
                    try:
                        r = http("/v1/chat/completions", {"messages": msgs, "max_tokens": 48, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}, "cache_prompt": True}, timeout=4 * 3600)
                        out = r["choices"][0]["message"]["content"]; tm = r.get("timings", {})
                        tr = {"depth": depth, "kind": kind, "hit": ans.lower() in out.lower(), "answer": out[:200], "prompt_n": tm.get("prompt_n"), "prefill_tps": tm.get("prompt_per_second"), "decode_tps": tm.get("predicted_per_second"), "wall_s": round(time.time() - t1, 1), "prompt_tokens": r.get("usage", {}).get("prompt_tokens")}
                    except Exception as e:
                        tr = {"depth": depth, "kind": kind, "hit": False, "error": str(e)[:300], "wall_s": round(time.time() - t1, 1)}
                    peak = max(peak, vram()); shared_peak = max(shared_peak, shared_mb(p.pid)); trials.append(tr)
                    log(f"{name} {ctx} d={depth} {kind} hit={tr['hit']} {tr.get('prefill_tps')} tok/s wall {tr['wall_s']}s")
                    if p.poll() is not None: break
                if p.poll() is not None: break
            rec.update(status="ok" if p.poll() is None else "server_died", trials=trials, recall=round(sum(t["hit"] for t in trials) / 9, 3), peak_vram_mb=peak, peak_shared_mb=shared_peak)
            p.kill(); errf.close(); time.sleep(5)
            open(OUT, "a").write(json.dumps(rec) + "\n"); log(f"{name} {ctx} recall {rec['recall']} peak {peak} MB shared {shared_peak} MB")
    log("=== E21DONE")
if __name__ == "__main__": main()