#!/usr/bin/env python3
"""Thin admission wrapper for the h2o confirmation run (Amendment 1, brokkr 63d7a1b).

Sits IN FRONT of the unmodified shipped h2o_lightning_shim.py and refuses over-cap
input loudly, mirroring INC's refusal. Counts with h2o's OWN tokenizer: renders the
state exactly as the shim does (imported verbatim from the shipped file) and asks
vLLM's /tokenize for the count. Refuses > CAP tokens with an explicit 422; forwards
everything else untouched. No trim. The shim is not modified.
"""
import importlib.util, json, os, sys, time, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

CAP = int(os.environ.get("ADMIT_CAP", "32000"))
UP = os.environ.get("ADMIT_UP", "http://127.0.0.1:18091")   # the shim
VLLM = os.environ.get("ADMIT_VLLM", "http://127.0.0.1:18099")
PORT = int(os.environ.get("ADMIT_PORT", "18092"))
SHIM_PATH = os.environ.get("ADMIT_SHIM", "/tank/bench/pkg-h2o/h2o_lightning_shim.py")

spec = importlib.util.spec_from_file_location("h2o_shim", SHIM_PATH)
shim_mod = importlib.util.module_from_spec(spec)
sys.modules["h2o_shim"] = shim_mod
spec.loader.exec_module(shim_mod)

CFG = shim_mod.Contract(json.load(open("/tank/bench/pkg-h2o/serve_config.json")))
CFG.max_state_tokens = 10**9   # wrapper must count the UNTRUNCATED render; the shim's own trim is what we police

def sent_tokens(body):
    """Tokens per prompt as vLLM will actually receive them, counted by vLLM's own tokenizer via
    the shim's own renderers (chat path: tokenize over messages, matching the image-free prompt
    the shim sends per question). Returns (max_per_prompt, shim_accounted_total).
    The shim accounts usage as shared_prefix + sum(question tails); we mirror that exactly."""
    import concurrent.futures
    state = body.get("state") or ""
    if isinstance(state, dict):
        state = json.dumps(state)
    state_txt = shim_mod.render_state(state, CFG.state_compact)   # UNTRUNCATED: what admission guards
    all_toks = []
    for qid, q in body.get("questions", {}).items():
        prompt_text = CFG.question_prompt(state_txt, q, qid)[0]   # EXACTLY what the shim sends per question
        req = urllib.request.Request(VLLM + "/tokenize",
            data=json.dumps({"model": CFG.model, "prompt": prompt_text, "add_special_tokens": False}).encode(),
            headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=30) as r:
            all_toks.append(json.loads(r.read())["tokens"])
    mx = max(len(t) for t in all_toks)
    shared = shim_mod.common_prefix_len(all_toks)
    total = shared + sum(len(t) - shared for t in all_toks)   # the shim's own accounting: head + tails
    return mx, total

class H(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    def log_message(self, *a): pass
    def _send(self, code, obj):
        b = json.dumps(obj).encode()
        self.send_response(code); self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b)
    def do_GET(self):
        try:
            req = urllib.request.Request(UP + self.path)
            with urllib.request.urlopen(req, timeout=30) as r:
                self._send(r.status, json.loads(r.read()))
        except urllib.error.HTTPError as e:
            self._send(e.code, json.loads(e.read() or b"{}"))
    def do_POST(self):
        n = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(n)
        mx = total = None
        if self.path.split("?")[0] in ("/v1/systemone", "/api/alpha/decisions"):
            try:
                body = json.loads(raw)
            except ValueError:
                return self._send(422, {"error": "admission: body is not JSON"})
            try:
                mx, total = sent_tokens(body)
            except Exception as e:
                return self._send(502, {"error": f"admission: tokenizer count failed: {e}"})
            if mx > CAP or total > CAP:
                return self._send(422, {"error": f"admission: {mx} max-prompt / {total} total tokens against the "
                                                f"{CAP}-token seat limit; refused, not trimmed",
                                        "max_prompt_tokens": mx, "shim_accounted_total": total, "cap": CAP})
        req = urllib.request.Request(UP + self.path, data=raw,
            headers={"Content-Type": "application/json"}, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=300) as r:
                obj = json.loads(r.read())
                obj.setdefault("_admission", {})["max_prompt_tokens"] = mx if "mx" in dir() else None
                obj["_admission"]["shim_accounted_total"] = total if "total" in dir() else None
                self._send(r.status, obj)
        except urllib.error.HTTPError as e:
            self._send(e.code, json.loads(e.read() or b"{}"))

ThreadingHTTPServer(("127.0.0.1", PORT), H).serve_forever()