"""Local HTTP server speaking TypeSafe's System One contract (POST /v1/systemone) on top of a System1 backend.
Run: python system1_server.py --model <path> [--adapter <safetensors|dir>] [--port 8020] [--model-name system1-9b]
Probabilities come from the same path as scripts/evaluate_system1.py predict (render_messages -> backend.score ->
calibrated_probabilities at T=1). Confidence formulas mirror the Kev reference (Apache-2.0, kev/api.py: choice_confidence,
score_confidence, round_prob), themselves mirroring TypeSafe system-one-adapter 0.2.1 confidence_metrics."""
import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from kev_convert import kev_record_to_rows
from system1_client import (
    WIDE_CHOICE_KEYS, MAX_PAYLOAD_CHARS, OverlengthError, calibrated_probabilities, render_messages, request_from_row)

MAX_BODY_BYTES = 4 * MAX_PAYLOAD_CHARS  # generous: JSON escaping and multi-question requests; rows are re-bounded per question


class ContractError(ValueError):
    pass


def round_prob(x):
    return round(float(x), 4)


def _normalize(p):
    t = sum(p)
    return [1 / len(p)] * len(p) if t == 0 else [x / t for x in p]


def choice_confidence(p):
    K = len(p)
    return 1.0 if K == 1 else (max(_normalize(p)) - 1 / K) / (1 - 1 / K)


def score_confidence(p):
    L = len(p)
    if L == 1:
        return 1.0
    p = _normalize(p)
    mode = max(range(L), key=p.__getitem__)
    D = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
    return max(0.0, 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / D)


def request_to_rows(req):
    """Validate a System One request -> ({qid: row}, {qid: candidate names}). Rows come from kev_convert with a dummy label
    (never rendered), so option text equals what the evaluator and Jev conversion produce."""
    if not isinstance(req, dict):
        raise ContractError("request: expected object")
    if "state" not in req or not isinstance(req["state"], (str, dict, list)):
        raise ContractError("state: required, string, object or array")
    qs = req.get("questions")
    if not isinstance(qs, dict) or not qs:
        raise ContractError("questions: expected nonempty object")
    dummy, rows, names = {}, {}, {}
    for qid, q in qs.items():
        if not isinstance(q, dict):
            raise ContractError(f"{qid}: expected object")
        kind, crit = q.get("type"), q.get("criteria")
        if kind == "choice":
            if not isinstance(crit, (dict, list)) or not 2 <= len(crit) <= len(WIDE_CHOICE_KEYS):
                raise ContractError(f"{qid}: choice requires 2-{len(WIDE_CHOICE_KEYS)} options per choice")
            label = next(iter(crit))
        elif kind == "score":
            if not isinstance(crit, list):
                raise ContractError(f"{qid}: score criteria must be an ordered list")
            label = 0
        elif kind == "noul":
            if crit is not None and not isinstance(crit, dict):
                raise ContractError(f"{qid}: noul criteria must be an object")
            label = True
        else:
            raise ContractError(f"{qid}: unknown question type {kind!r}")
        dummy[qid] = {"type": kind, "instructions": q.get("instructions"), "criteria": crit, "label": label}
    try:
        for row in kev_record_to_rows({"state": req["state"], "questions": dummy}, "serve", "none", "serve", keys=WIDE_CHOICE_KEYS):
            qid = row["meta"]["kev_question_id"]
            rows[qid] = row
            names[qid] = list(row["meta"].get("candidate_ids", {}).values())
    except (ValueError, KeyError, TypeError) as e:
        raise ContractError(str(e))
    return rows, names


def answer_for(row, names, probs):
    kind = row["type"]
    if kind == "noul":
        return {"type": "noul", "noul": round_prob(probs["Y"])}
    if kind == "choice":
        p = [probs[k] for k in row["options"]]
        return {"type": "choice", "choice": names[max(range(len(p)), key=p.__getitem__)], "confidence": round_prob(choice_confidence(p)),
                "probabilities": {n: round_prob(v) for n, v in zip(names, p)}}
    p = [probs[k] for k in row["options"]]
    return {"type": "score", "score": round_prob(sum(i * v for i, v in enumerate(p))),
            "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)}, "confidence": round_prob(score_confidence(p))}


class Service:
    def __init__(self, backend, model_name):
        self.backend, self.model_name, self.lock = backend, model_name, threading.Lock()  # ponytail: one global lock, GPU is serial anyway

    def _tokens(self, text_or_messages):
        tok = getattr(self.backend, "tokenizer", None)
        if tok is None:
            return 0
        if isinstance(text_or_messages, str):
            return len(tok.encode(text_or_messages, add_special_tokens=False))
        prompt = tok.apply_chat_template(text_or_messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
        return len(tok.encode(prompt, add_special_tokens=False))

    def systemone(self, req):
        start = time.perf_counter()
        rows, names = request_to_rows(req)
        answers, in_tokens = {}, 0
        for qid, row in rows.items():
            try:
                messages = render_messages(request_from_row(row))
            except OverlengthError as e:
                raise ContractError(f"{qid}: payload too large, over the maximum context length: {e}")
            except ValueError as e:
                raise ContractError(f"{qid}: {e}")
            with self.lock:
                try:
                    scores = self.backend.score(messages, list(row["options"]))
                except OverlengthError as e:
                    raise ContractError(f"{qid}: context too large, over the maximum context length: {e}")
                in_tokens += self._tokens(messages)
            if not isinstance(scores, dict) or set(scores) != set(row["options"]):
                raise RuntimeError(f"{qid}: backend returned incomplete scores")
            answers[qid] = answer_for(row, names[qid], calibrated_probabilities(scores, 1.0))
        return {"model": self.model_name, "answers": answers,
                "usage": {"input_tokens": in_tokens, "output_tokens": self._tokens(json.dumps(answers))},
                "latency_ms": round((time.perf_counter() - start) * 1000, 1)}


def make_server(backend, model_name="system1", host="127.0.0.1", port=8020):
    svc = Service(backend, model_name)

    class Handler(BaseHTTPRequestHandler):
        def _send(self, code, body):
            data = json.dumps(body).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self):
            if self.path == "/health":
                self._send(200, {"status": "ok"})
            elif self.path == "/v1/models":
                self._send(200, {"object": "list", "data": [{"id": svc.model_name, "object": "model", "owned_by": "local"}]})
            else:
                self._send(404, {"error": "not found"})

        def do_POST(self):
            if self.path != "/v1/systemone":
                return self._send(404, {"error": "not found"})
            n = int(self.headers.get("Content-Length") or 0)
            if n > MAX_BODY_BYTES:
                return self._send(422, {"error": f"payload too large: {n} bytes"})
            try:
                req = json.loads(self.rfile.read(n))
            except (ValueError, UnicodeDecodeError) as e:
                return self._send(400, {"error": f"malformed JSON: {e}"})
            try:
                self._send(200, svc.systemone(req))
            except ContractError as e:
                self._send(422, {"error": str(e)})
            except Exception as e:
                self._send(500, {"error": f"{type(e).__name__}: {e}"})

        def log_message(self, *a):
            pass

    return ThreadingHTTPServer((host, port), Handler)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--model", required=True)
    ap.add_argument("--adapter")
    ap.add_argument("--port", type=int, default=8020)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--model-name", default="system1-9b")
    ap.add_argument("--backend", choices=("mlx", "torch"), default="mlx", help="torch: CUDA via transformers")
    a = ap.parse_args()
    import system1_client
    backend = system1_client.TorchBackend if a.backend == "torch" else system1_client.MlxBackend
    server = make_server(backend(a.model, adapter_path=a.adapter), a.model_name, a.host, a.port)
    print(f"serving {a.model_name} on http://{a.host}:{a.port}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()