#!/usr/bin/env python3
# Amendment 1 confirmation driver for h2o (brokkr 63d7a1b). Stack assumed up:
#   vLLM :18099, shipped shim :18091, admission wrapper :18092.
# Shape sizes are pinned by the WRAPPER's own count (max per-prompt prompt tokens, pre-trim):
# the shim's usage.input_tokens reports post-trim and is not trustworthy at the cap.
# The gate's filler budget was written for INC's ~3.6 chars/token; h2o's tokenizer runs ~5.1,
# so every builder target is iteratively rescaled against the wrapper count.
import json, os, sys, time, statistics, subprocess, threading, random, urllib.request, concurrent.futures
sys.path.insert(0, "/tank/bench")
from gate_longctx_s1 import build_state
CAP = 32000
MODEL = "h2oai/h2o-lightning-4b"
URL = "http://127.0.0.1:18092"   # through the admission wrapper
OUT = "/tank/bench/arm-h2o"
TAG = os.environ.get("CONFIRM_TAG","A")
rng = random.Random("h2o-confirm-2026-10-12")

def post(body, timeout=300):
    req = urllib.request.Request(URL+"/v1/systemone", data=json.dumps(body).encode(),
                                 headers={"Content-Type":"application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())
def topper(r, qid):
    a = (r.get("answers") or {}).get(qid) or {}
    p = a.get("probabilities") or {}
    return max(p, key=lambda k: float(p[k])) if p else None
def questions(nq):
    return {f"q{i}": {"type":"choice","instructions":f"Does the state satisfy criterion {i}?",
                      "criteria":{"yes":f"topic-{i} mentioned","no":"not mentioned"}} for i in range(nq)}
def adm(body):
    r = post(body)
    return r, (r.get("_admission") or {})
EVID=("DECISIVE: the underwriter's signed memo of 2026-03-14 states the claim was established "
      "by documents filed before the deadline; the renewal decision must rest on this memo.")
GQ={"type":"choice","instructions":"Based only on the state above, does the evidence establish the claim, "
    "contradict it, or is it insufficient?",
    "criteria":{"supported":"The evidence establishes the claim","contradicted":"The evidence establishes the opposite",
                "insufficient":"The evidence does not establish either"}}
def count_gate(st):
    # wrapper counts then forwards; over-cap it REFUSES with the count in the body — either way we learn the count
    try:
        r,a = adm({"model":MODEL,"state":st,"questions":{"q1":GQ}})
        return a.get("max_prompt_tokens")
    except urllib.error.HTTPError as e:
        if e.code == 422:
            try: return json.loads(e.read()).get("max_prompt_tokens")
            except Exception: return None
        raise

peak=0; stop=False
def poll():
    global peak, stop
    while not stop:
        try: peak=max(peak,int(subprocess.run(["nvidia-smi","-i","3","--query-gpu=memory.used",
             "--format=csv,noheader,nounits"],capture_output=True,text=True).stdout.split()[0]))
        except Exception: pass
        time.sleep(0.05)

res={"tag":TAG}
# ---- largest-admissible state: iterate until wrapper max-prompt is within 160 tok under CAP
STATE_FILE = os.environ.get("CONFIRM_STATE","")
if os.environ.get("CONFIRM_LOAD") == "1" and STATE_FILE and os.path.exists(STATE_FILE):
    state32 = open(STATE_FILE).read()
    hist=[{"note":"state loaded from file; first engine request is the first shape cell"}]
    target=None
else:
  target=31900; hist=[]
  for it in range(8):
    st = build_state(target, None, 0, rng)
    try:
        r, a = adm({"model":MODEL,"state":st,"questions":questions(16)})   # pin the TOTAL at 16q, per clause 3
        mx = a.get("max_prompt_tokens"); tot = a.get("shim_accounted_total")
    except urllib.error.HTTPError as e:
        if e.code != 422: raise
        b = json.loads(e.read()); mx = b.get("max_prompt_tokens"); tot = b.get("shim_accounted_total")
    hist.append({"iter":it,"builder_tok_target":int(target),"wrapper_max_prompt":mx,"wrapper_total":tot})
    if tot is None: raise SystemExit(str(mx))
    if tot <= CAP and (CAP-tot) <= 160: break
    target = int(target * max(0.5, min(1.3, (CAP-80)/max(tot,1))))
if target is not None: state32 = st
res["calibration"]=hist

def shape(nq, cold_expected=False):
    body={"model":MODEL,"state":state32,"questions":questions(nq)}
    t=time.perf_counter(); r=post(body); first=(time.perf_counter()-t)*1000
    xs=[first]
    for _ in range(4):
        t=time.perf_counter(); post(body); xs.append((time.perf_counter()-t)*1000)
    r2, a2 = adm(body)
    return {"nq":nq,"first_ms":round(first),"warm_median_ms":round(statistics.median(xs)),
            "cold_expected":cold_expected,
            "max_prompt_tokens":a2.get("max_prompt_tokens"),"shim_total":a2.get("shim_accounted_total"),
            "usage_input_tokens":(r2.get("usage") or {}).get("input_tokens")}

STATE_FILE=os.environ.get("CONFIRM_STATE","")
if STATE_FILE: open(STATE_FILE,"w").write(state32)
cold_first=int(os.environ.get("COLD_FIRST","0"))   # which nq gets the truly-cold first call (0 = none)
th=threading.Thread(target=poll); th.start()
res["cells"]={}
order=([cold_first] if cold_first in (1,16) else [1,16])
order=order+[n for n in (1,16) if n not in order]
for nq in order:
    res["cells"][f"{nq}q32k"]=shape(nq, cold_expected=(nq==cold_first))
# ---- concurrency: 4x concurrent 16q at the same state
body={"model":MODEL,"state":state32,"questions":questions(16)}
t=time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(4) as ex:
    futs=[ex.submit(post,body) for _ in range(4)]
    oks=[]
    for f in futs:
        try: oks.append(topper(f.result(),"q0") is not None)
        except Exception: oks.append(False)
res["concurrency_4x16q"]={"wall_ms":round((time.perf_counter()-t)*1000),"all_ok":all(oks),"per_call_ok":oks}
# ---- gate cells: evidence at absolute h2o-token depths ~1k/8k/16k/30k, whole state pinned ~31,900
gate=[]
for depth_tok in (1000,8000,16000,30000):
    tgt=31900
    g=tgt*1.42; hist2=[]
    for it in range(8):
        d_scaled = int(depth_tok * (g/tgt))   # relative depth scales with the char budget
        stg = build_state(int(g/3.6), EVID, d_scaled, rng)
        mx = count_gate(stg)
        hist2.append({"chars_target":int(g),"count":mx})
        if mx is None: break
        if abs(mx-tgt)<=200: break
        g = int(g * tgt/mx)
    try:
        r,a = adm({"model":MODEL,"state":stg,"questions":{"q1":GQ}})
        cell={"depth_tok_target":depth_tok,"max_prompt_tokens":a.get("max_prompt_tokens"),
              "top":topper(r,"q1"),"usage_input_tokens":(r.get("usage") or {}).get("input_tokens"),
              "refused":False,"cal":hist2[-1]}
    except urllib.error.HTTPError as e:
        cell={"depth_tok_target":depth_tok,"refused":True,"detail":json.loads(e.read()),"cal":hist2[-1]}
    gate.append(cell)
    time.sleep(1)
res["gate"]=gate
stop=True; th.join()
res["peak_MiB"]=peak
# ---- over-cap REFUSAL (the whole point): take state32 and grow it past the cap
st_over = state32 + ("Policy renewal review cycle X: coverage schedules, endorsements and rider terms are "
                     "reconciled against the master binder. " * 400)
refused=False; detail=None; mx_over=count_gate(st_over)
try:
    r_ok = post({"model":MODEL,"state":st_over,"questions":{"q1":GQ}})
    detail={"note":"not refused","max_prompt_tokens":(r_ok.get("_admission") or {}).get("max_prompt_tokens")}
except urllib.error.HTTPError as e:
    refused=True; detail={"status":e.code,"body":json.loads(e.read())}
res["over_cap_refusal"]={"refused":refused,"counted_tokens":mx_over,"detail":detail}
json.dump(res,open(f"{OUT}/confirm-{TAG}.json","w"),indent=1)
print("CONFIRM-DONE", json.dumps(res)[:2400])