#!/usr/bin/env python3
"""The front of the h2o-lightning-4b decision seat (fv-ml1). Stdlib only.

It sits in front of the UNMODIFIED shipped h2o_lightning_shim.py (which sits in front of stock vLLM 0.30.0) and adds
the three things the seat contract needs and the vendor stack does not have:

1. ADMISSION (from infra-hermes's tools/h2o-admission/admission_proxy.py @ 30f36f2, logic unchanged): the shim
   silently head-and-tail-trims state above ~32,000 tokens, so every decision request is counted first with h2o's own
   tokenizer (vLLM /tokenize over the shim's own per-question renders) and refused with 422 + the count when either the
   largest prompt or the shim-accounted total (shared head + tails) exceeds FRONT_CAP. Refused, never trimmed.
2. THE SEAT'S AUTH: a bearer token on everything but GET /health, as intern-decision-serve has (OPEN_PATHS /health).
3. CONTINUITY for /decide and /decide/shared (semif's surface, albok-nemi's sensitivity judge): each decision becomes
   one Jev `choice` question exactly as intern-decision-serve builds it ({type: choice, instructions: question,
   criteria: {option id: description}}), asked through /v1/systemone in calls of at most 16 questions, and re-keyed
   into the semif result shape the judge reads (results[i].native.probabilities[option id]). orderings other than
   "none" are refused (422): the judge does its own forward/reverse. /health carries {model: {name, revision}}, which
   the judge stamps into its cache key.

Env: FRONT_TOKEN (required), FRONT_BIND, FRONT_PORT, FRONT_SHIM (shim URL), FRONT_VLLM (vLLM URL), FRONT_CAP (32000),
FRONT_SHIM_FILE (vendor shim path, imported for its renderers), FRONT_SERVE_CONFIG, FRONT_MODEL_NAME,
FRONT_MODEL_SOURCE, FRONT_MODEL_REVISION, FRONT_MAX_DECISIONS (64).
"""
import hmac, importlib.util, json, os, sys, time, urllib.error, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

TOKEN = os.environ["FRONT_TOKEN"]
BIND = os.environ.get("FRONT_BIND", "127.0.0.1")
PORT = int(os.environ.get("FRONT_PORT", "8034"))
UP = os.environ.get("FRONT_SHIM", "http://127.0.0.1:18191")
VLLM = os.environ.get("FRONT_VLLM", "http://127.0.0.1:18199")
CAP = int(os.environ.get("FRONT_CAP", "32000"))
SHIM_FILE = os.environ.get("FRONT_SHIM_FILE", "/model/h2o_lightning_shim.py")
SERVE_CONFIG = os.environ.get("FRONT_SERVE_CONFIG", "/model/serve_config.json")
MODEL = {"name": os.environ.get("FRONT_MODEL_NAME", "h2o-lightning-4b"),
         "source": os.environ.get("FRONT_MODEL_SOURCE", "h2oai/h2o-lightning-4b"),
         "revision": os.environ.get("FRONT_MODEL_REVISION", "b43f759b6385d21182f87c3829931ec3d4dfa98a")}
MAX_DECISIONS = int(os.environ.get("FRONT_MAX_DECISIONS", "64"))
MAX_QUESTIONS_PER_CALL = 16
MIN_OPTIONS, MAX_OPTIONS = 2, 16

spec = importlib.util.spec_from_file_location("h2o_shim", SHIM_FILE)
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(SERVE_CONFIG)))
CFG.max_state_tokens = 10**9   # count the UNTRUNCATED render; the shim's own trim is what admission polices


class Refused(Exception):
    def __init__(self, code, obj):
        super().__init__(obj)
        self.code, self.obj = code, obj


def sent_tokens(body):
    """(max tokens of one question's prompt, shim-accounted total = shared head + tails), counted by vLLM's own
    tokenizer over the shim's own per-question renders (admission_proxy.py @ 30f36f2, unchanged)."""
    state = body.get("state") or ""
    if isinstance(state, (dict, list)):
        state = json.dumps(state)
    state_txt = shim_mod.render_state(state, CFG.state_compact)
    all_toks = []
    for qid, q in body.get("questions", {}).items():
        prompt_text = CFG.question_prompt(state_txt, q, qid)[0]
        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)
    return mx, shared + sum(len(t) - shared for t in all_toks)


def admit(body):
    try:
        mx, total = sent_tokens(body)
    except Refused:
        raise
    except Exception as e:  # noqa: BLE001
        raise Refused(502, {"error": f"admission: tokenizer count failed: {type(e).__name__}: {e}"})
    if mx > CAP or total > CAP:
        raise Refused(422, {"error": f"admission: {mx} max-prompt / {total} total tokens against the {CAP}-token seat "
                                     f"limit; refused, not trimmed", "max_prompt_tokens": mx,
                            "shim_accounted_total": total, "cap": CAP})
    return mx, total


def post_up(path, obj):
    req = urllib.request.Request(UP + path, data=json.dumps(obj).encode(), headers={"Content-Type": "application/json"},
                                 method="POST")
    try:
        with urllib.request.urlopen(req, timeout=300) as r:
            return r.status, json.loads(r.read())
    except urllib.error.HTTPError as e:
        raw = e.read() or b"{}"
        try:
            return e.code, json.loads(raw)
        except ValueError:
            return e.code, {"error": raw[:300].decode(errors="replace")}


def systemone(body):
    mx, total = admit(body)
    code, obj = post_up("/v1/systemone", body)
    if isinstance(obj, dict):
        obj.setdefault("_admission", {}).update({"max_prompt_tokens": mx, "shim_accounted_total": total, "cap": CAP})
    return code, obj


def check_decision(d, i):
    if not isinstance(d, dict) or not isinstance(d.get("id"), str) or not isinstance(d.get("question"), str):
        raise Refused(422, {"error": f"decision {i}: needs string id and question"})
    opts = d.get("options")
    if not isinstance(opts, list) or not MIN_OPTIONS <= len(opts) <= MAX_OPTIONS:
        raise Refused(422, {"error": f"decision {d['id']!r}: options must contain {MIN_OPTIONS}-{MAX_OPTIONS} entries"})
    if any(not isinstance(o, dict) or not isinstance(o.get("id"), str) or not isinstance(o.get("description"), str)
           for o in opts):
        raise Refused(422, {"error": f"decision {d['id']!r}: every option needs string id and description"})
    if len({o["id"] for o in opts}) != len(opts):
        raise Refused(422, {"error": f"decision {d['id']!r}: option ids must be unique"})
    if d.get("orderings", "none") != "none":
        raise Refused(422, {"error": f"decision {d['id']!r}: orderings {d.get('orderings')!r} is not served by this "
                                     "seat (h2o); send orderings 'none' and combine on the client"})


def decide(state, decisions):
    """semif /decide/shared on top of /v1/systemone: intern-decision-serve's mapping (one Jev choice question per
    decision, calls of at most 16, fields q / q1..qN), re-keyed into its result shape."""
    if not isinstance(decisions, list) or not 1 <= len(decisions) <= MAX_DECISIONS:
        raise Refused(422, {"error": f"decisions must be a list of 1..{MAX_DECISIONS}"})
    for i, d in enumerate(decisions):
        check_decision(d, i)
    started, results, sizes, tokens = time.perf_counter(), [], [], []
    for k in range(0, len(decisions), MAX_QUESTIONS_PER_CALL):
        chunk = decisions[k:k + MAX_QUESTIONS_PER_CALL]
        fields = ["q"] if len(chunk) == 1 else [f"q{j + 1}" for j in range(len(chunk))]
        questions = {f: {"type": "choice", "instructions": d["question"],
                         "criteria": {o["id"]: o["description"] for o in d["options"]}} for f, d in zip(fields, chunk)}
        code, resp = systemone({"state": state, "questions": questions})
        if code != 200:
            raise Refused(code, resp if isinstance(resp, dict) else {"error": str(resp)})
        sizes.append(len(chunk))
        tokens.append((resp.get("usage") or {}).get("input_tokens"))
        for f, d in zip(fields, chunk):
            answer = (resp.get("answers") or {}).get(f)
            ids = [o["id"] for o in d["options"]]
            probs = (answer or {}).get("probabilities") or {}
            if answer is None or any(i not in probs for i in ids):
                raise Refused(500, {"error": f"scoring_failed: the model's answer for field {f!r} lacks option ids"})
            results.append({"id": d["id"], "option_ids": ids, "probabilities": [probs[i] for i in ids],
                            "top": answer.get("choice"), "confidence": answer.get("confidence"),
                            "calibration": resp.get("calibration"), "native": answer,
                            "input_tokens": tokens[-1], "model": MODEL,
                            "call": {"index": len(sizes) - 1, "field": f, "questions": len(chunk)}})
    return {"results": results, "timing": {"total_seconds": time.perf_counter() - started, "calls": len(sizes),
                                           "questions_per_call": sizes, "input_tokens": tokens}}


class H(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args))

    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 _authorized(self):
        scheme, _, tok = (self.headers.get("Authorization") or "").partition(" ")
        return scheme.lower() == "bearer" and hmac.compare_digest(tok.strip().encode(), TOKEN.encode())

    def do_GET(self):
        path = self.path.split("?")[0]
        if path == "/health":
            try:
                with urllib.request.urlopen(UP + "/health", timeout=10) as r:
                    shim = json.loads(r.read())
            except Exception as e:  # noqa: BLE001
                return self._send(503, {"status": "unavailable", "model": MODEL, "error": f"shim: {type(e).__name__}"})
            ok = bool(shim.get("ok"))
            return self._send(200 if ok else 503, {
                "status": "ok" if ok else "starting", "model": MODEL, "cap_tokens": CAP, "admission": "refuse, never trim",
                "endpoints": ["/decide", "/decide/shared", "/v1/systemone"], "max_decisions": MAX_DECISIONS,
                "max_questions_per_call": MAX_QUESTIONS_PER_CALL, "orderings": ["none"], "backend": shim})
        if not self._authorized():
            return self._send(401, {"error": "unauthorized"})
        if path == "/v1/models":
            try:
                with urllib.request.urlopen(UP + "/v1/models", timeout=10) as r:
                    return self._send(r.status, json.loads(r.read()))
            except urllib.error.HTTPError as e:
                return self._send(e.code, {"error": "upstream"})
        return self._send(404, {"error": "not found"})

    def do_POST(self):
        path = self.path.split("?")[0]
        if not self._authorized():
            return self._send(401, {"error": "unauthorized"})
        n = int(self.headers.get("Content-Length", 0) or 0)
        if n > 4 * 1024 * 1024:
            return self._send(413, {"error": "body over 4 MiB"})
        try:
            body = json.loads(self.rfile.read(n) or b"null")
        except ValueError:
            return self._send(422, {"error": "body is not JSON"})
        if not isinstance(body, dict):
            return self._send(422, {"error": "body must be a JSON object"})
        try:
            if path == "/v1/systemone":
                return self._send(*systemone(body))
            if path == "/api/alpha/decisions":
                admit(body)
                return self._send(*post_up(path, body))
            if path == "/decide/shared":
                return self._send(200, decide(body.get("state"), body.get("decisions")))
            if path == "/decide":
                d = {k: body.get(k) for k in ("id", "question", "options", "orderings") if k in body}
                out = decide(body.get("state"), [d])
                return self._send(200, {**out["results"][0], "total_seconds": out["timing"]["total_seconds"]})
        except Refused as r:
            return self._send(r.code, r.obj)
        return self._send(404, {"error": "not found"})


ThreadingHTTPServer((BIND, PORT), H).serve_forever()