"""deg HTTP server: OpenAI Decisions, TypeSafe Jev SystemOne and deg fill on one model.

Routes (shapes in :mod:`deg.api`):

* ``POST /v1/decisions``: OpenAI Decisions API (stock ``openai`` >= 3.26 works with
  ``OpenAI(base_url="http://host:8790/v1")``); deg extensions ``extract`` / ``generate``.
* ``POST /v1/systemone``: TypeSafe Jev / SystemOne (and Clef); same extensions.
* ``POST /v1/fill``: deg's schema fill (SPEC.md §9).
* ``GET /v1/models``, ``GET /health``.

Every request becomes one :func:`deg.infer.fill` job (all of a request's questions are
fields of one schema: the input is read once). Jobs go through a :class:`Batcher`: one GPU
worker; concurrent jobs that arrive within ``window_s`` of the first are taken together (up
to ``max_batch``) and handed to ``fill_many`` when the runtime provides one, else run back
to back on the worker. A failing job fails only itself. More than ``max_pending`` queued
jobs -> 429.

Auth: when ``DEG_API_KEY`` is set, requests need ``Authorization: Bearer <key>`` (401 in
OpenAI shape on ``/v1/decisions`` and ``/v1/fill``; 403 in Jev shape on ``/v1/systemone``,
as api.typesafe.ai answers).

Telemetry (OpenTelemetry API; exporters are the deployer's choice). Cardinality:

* span ``deg.request``: ``deg.route`` in {decisions, systemone, fill, models, health,
  other} (6), ``deg.status`` in {ok, client_error, unauthorized, overloaded,
  server_error} (5), ``deg.questions`` and ``deg.batch_size`` (integers on spans only,
  never metric dimensions);
* histogram ``deg.server.duration`` (s) by route x status: <= 30 series;
* histogram ``deg.server.batch_size``: no attributes;
* counter ``deg.server.questions`` by route (2) x question_type (5): <= 10 series;
* counter ``deg.server.tokens`` by route (3) x direction {input, output}: <= 6 series.

No request content, question names, model names or client identifiers are recorded.

``python -m deg.server --ckpt /data/deg/runs/run-1/step-7000 --port 8790``
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import math
import os
import queue
import sys
import threading
import time
from collections.abc import Callable
from concurrent.futures import Future
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

from opentelemetry import metrics, trace

from deg.api import (
    JEV_COMPAT_MODELS,
    OPENAI_COMPAT_MODELS,
    ApiError,
    Job,
    decisions_response,
    parse_decisions,
    parse_systemone,
    systemone_response,
)
from deg.request import (
    MODEL_ALIASES,
    FillOptions,
    RequestError,
    generate_arg,
    parse_fill_request,
)
from deg.response import MODEL_ID, FillError

READ_PIECE_BYTES = 1 << 20  # socket read granularity, not a size limit
ROUTES = {"/v1/decisions": "decisions", "/v1/systemone": "systemone", "/v1/fill": "fill",
          "/v1/models": "models", "/health": "health"}

FillFn = Callable[[Any, dict[str, Any], FillOptions], dict[str, Any]]
FillManyFn = Callable[[list[tuple[Any, dict[str, Any], FillOptions]]], list[Any]]

_tracer = trace.get_tracer("deg.server")
_meter = metrics.get_meter("deg.server")
_duration = _meter.create_histogram("deg.server.duration", unit="s",
                                    description="request latency by route and status")
_batch_size = _meter.create_histogram("deg.server.batch_size", unit="{job}",
                                      description="fill jobs taken together by the GPU worker")
_questions = _meter.create_counter("deg.server.questions", unit="{question}",
                                   description="questions answered, by route and type")
_tokens = _meter.create_counter("deg.server.tokens", unit="{token}",
                                description="model tokens, by route and direction")


class Overloaded(Exception):
    """Too many jobs queued (429) or the model ran out of memory (503 / 529)."""

    def __init__(self, message: str, status: int) -> None:
        super().__init__(message)
        self.status = status


class Batcher:
    """One GPU worker; jobs arriving within ``window_s`` of the first are taken together."""

    def __init__(self, fill: FillFn, *, fill_many: FillManyFn | None = None,
                 window_s: float = 0.005, max_batch: int = 16, max_pending: int = 256) -> None:
        self.fill, self.fill_many = fill, fill_many
        self.window, self.max_batch, self.max_pending = window_s, max_batch, max_pending
        self.q: queue.Queue = queue.Queue()
        self.lock = threading.Lock()
        self.counts = {"jobs": 0, "batches": 0, "max_seen_batch": 0, "rejected": 0}
        threading.Thread(target=self._loop, daemon=True, name="deg-gpu").start()

    def stats(self) -> dict[str, Any]:
        with self.lock:
            out: dict[str, Any] = dict(self.counts)
        out["mean_batch"] = round(out["jobs"] / out["batches"], 2) if out["batches"] else None
        out["queued"] = self.q.qsize()
        return out

    def submit(self, context: Any, schema: dict[str, Any], options: FillOptions) -> Future:
        if self.q.qsize() >= self.max_pending:
            with self.lock:
                self.counts["rejected"] += 1
            raise Overloaded("too many requests queued", 429)
        fut: Future = Future()
        self.q.put(((context, schema, options), fut))
        return fut

    def _take(self) -> list:
        batch = [self.q.get()]
        deadline = time.perf_counter() + self.window
        while len(batch) < self.max_batch:
            left = deadline - time.perf_counter()
            if left <= 0:
                break
            try:
                batch.append(self.q.get(timeout=left))
            except queue.Empty:
                break
        return batch

    def _run(self, batch: list) -> None:
        jobs = [job for job, _ in batch]
        if self.fill_many is not None and len(jobs) > 1:
            try:
                results = self.fill_many(jobs)
            except Exception:  # noqa: BLE001 -- isolated: rerun each job alone below
                results = None
            if results is not None and len(results) == len(jobs):
                for (_, fut), res in zip(batch, results):
                    if isinstance(res, BaseException):
                        fut.set_exception(res)
                    else:
                        fut.set_result(res)
                return
        for job, fut in batch:
            try:
                fut.set_result(self.fill(*job))
            except Exception as exc:  # noqa: BLE001 -- delivered to that request only
                fut.set_exception(exc)

    def _loop(self) -> None:
        while True:
            batch = self._take()
            with self.lock:
                self.counts["jobs"] += len(batch)
                self.counts["batches"] += 1
                self.counts["max_seen_batch"] = max(self.counts["max_seen_batch"], len(batch))
            _batch_size.record(len(batch))
            self._run(batch)


def _trace_null_margin(options: FillOptions) -> None:
    """A request-set null margin on the request span (a span attribute only: no metric
    dimension, so its unbounded float values add no metric cardinality)."""
    if options.null_margin is not None:
        trace.get_current_span().set_attribute("deg.request.null_margin",
                                               float(options.null_margin))


class Engine:
    """Routes -> parsed jobs -> batcher -> API responses."""

    def __init__(self, batcher: Batcher, *, model_id: str = MODEL_ID,
                 openai_models: frozenset[str] = MODEL_ALIASES | OPENAI_COMPAT_MODELS,
                 jev_models: frozenset[str] = MODEL_ALIASES | JEV_COMPAT_MODELS) -> None:
        self.batcher, self.model_id = batcher, model_id
        self.openai_models, self.jev_models = openai_models, jev_models

    def _fill(self, job: Job) -> dict[str, Any] | None:
        if job.schema is None:
            return None
        _trace_null_margin(job.options)
        return self.batcher.submit(job.context, job.schema, job.options).result()

    def decisions(self, payload: Any) -> tuple[dict[str, Any], Job]:
        job = parse_decisions(payload, self.openai_models)
        return decisions_response(job, self._fill(job), self.model_id), job

    def systemone(self, payload: Any) -> tuple[dict[str, Any], Job]:
        job = parse_systemone(payload, self.jev_models)
        return systemone_response(job, self._fill(job), self.model_id), job

    def fill(self, payload: Any) -> dict[str, Any]:
        req = parse_fill_request(payload)
        _trace_null_margin(req.options)
        return self.batcher.submit(req.context, payload["schema"], req.options).result()

    def models(self) -> dict[str, Any]:
        names = sorted(MODEL_ALIASES | self.openai_models | self.jev_models)
        return {"object": "list", "data": [
            {"id": n, "object": "model", "created": 0, "owned_by": "deg",
             **({"root": self.model_id} if n != self.model_id else {})} for n in names]}


class BodyError(ValueError):
    pass


def read_body(stream: Any, headers: Any) -> bytes:
    """Whole body, Content-Length or chunked, read in bounded pieces (no size ceiling)."""
    body = bytearray()
    if "chunked" in (headers.get("transfer-encoding") or "").lower():
        while True:
            line = stream.readline()
            try:
                size = int(line.split(b";", 1)[0].strip(), 16)
            except ValueError as error:
                raise BodyError("bad chunk size") from error
            if size == 0:
                while stream.readline() not in (b"\r\n", b"\n", b""):
                    pass
                return bytes(body)
            while size:
                piece = stream.read(min(size, READ_PIECE_BYTES))
                if not piece:
                    raise BodyError("truncated chunk")
                body += piece
                size -= len(piece)
            stream.readline()
    try:
        remaining = int(headers.get("content-length") or 0)
    except ValueError as error:
        raise BodyError("bad content-length") from error
    while remaining > 0:
        piece = stream.read(min(remaining, READ_PIECE_BYTES))
        if not piece:
            raise BodyError("truncated body")
        body += piece
        remaining -= len(piece)
    return bytes(body)


def _route_error(route: str, status: int, message: str) -> tuple[int, dict[str, Any]]:
    """Error body in the route's own convention."""
    if route == "systemone":
        kind = {401: "authentication_error", 403: "authentication_error",
                429: "rate_limit_error", 529: "overloaded_error"}.get(status, "internal_error")
        if status < 500 and status not in (401, 403, 429):
            kind = "api_usage_error"
        return status, {"detail": {"error_type": kind, "message": message}}
    if route == "fill":
        kind = {401: "unauthorized", 404: "not_found", 422: "invalid_request",
                429: "rate_limited", 503: "overloaded"}.get(status, "server_error")
        return status, {"error": {"type": kind, "status": status, "message": message}}
    kind = {401: "invalid_request_error", 404: "invalid_request_error",
            400: "invalid_request_error", 429: "requests", 503: "server_error"}.get(
        status, "server_error")
    code = {401: "invalid_api_key", 429: "rate_limit_exceeded"}.get(status)
    return status, {"error": {"message": message, "type": kind, "param": None, "code": code}}


def make_handler(engine: Engine, *, api_key: str | None = None):
    class Handler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def log_message(self, *args: Any) -> None:  # no request content in logs
            return

        def _send(self, status: int, body: Any) -> None:
            data = json.dumps(body, ensure_ascii=False).encode("utf-8")
            self.send_response(status)
            self.send_header("content-type", "application/json")
            self.send_header("content-length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def _authorized(self) -> bool:
            return not api_key or self.headers.get("authorization", "") == f"Bearer {api_key}"

        def _handle(self, method: str) -> None:
            path = self.path.split("?", 1)[0].rstrip("/")
            route = ROUTES.get(path, "other")
            t0 = time.perf_counter()
            with _tracer.start_as_current_span("deg.request") as span:
                span.set_attribute("deg.route", route)
                status, body, outcome = self._dispatch(method, route, span)
                span.set_attribute("deg.status", outcome)
                _duration.record(time.perf_counter() - t0, {"route": route, "status": outcome})
                self._send(status, body)

        def _dispatch(self, method: str, route: str, span: Any
                      ) -> tuple[int, Any, str]:
            if not self._authorized():
                if route == "systemone":
                    return 403, {"detail": {
                        "error_type": "authentication_error",
                        "message": "Must supply an API key! Check your request and try "
                                   "again."}}, "unauthorized"
                return (*_route_error(route, 401, "Incorrect API key provided."),
                        "unauthorized")
            if method == "GET":
                if route == "health":
                    return 200, {"ok": True, "model": engine.model_id,
                                 "batcher": engine.batcher.stats()}, "ok"
                if route == "models":
                    return 200, engine.models(), "ok"
                return (*_route_error("fill", 404, "unknown path"), "client_error")
            if route not in ("decisions", "systemone", "fill"):
                return (*_route_error("fill", 404, "unknown path"), "client_error")
            try:
                raw = read_body(self.rfile, self.headers)
                payload = json.loads(raw) if raw else None
            except (BodyError, json.JSONDecodeError, UnicodeDecodeError) as exc:
                self.close_connection = isinstance(exc, BodyError)
                status = 400 if route != "fill" else 422
                if route == "systemone":
                    status = 422
                return (*_route_error(route, status, "body is not valid JSON"),
                        "client_error")
            try:
                if route == "fill":
                    body = engine.fill(payload)
                    _tokens.add(body["usage"]["prompt_tokens"],
                                {"route": route, "direction": "input"})
                    _tokens.add(body["usage"]["generated_tokens"],
                                {"route": route, "direction": "output"})
                    return 200, body, "ok"
                body, job = (engine.decisions if route == "decisions"
                             else engine.systemone)(payload)
                span.set_attribute("deg.questions", len(job.questions))
                for kind in job.question_types:
                    _questions.add(1, {"route": route, "question_type": kind})
                usage = body["usage"]
                _tokens.add(usage["input_tokens"], {"route": route, "direction": "input"})
                _tokens.add(usage["output_tokens"], {"route": route, "direction": "output"})
                return 200, body, "ok"
            except ApiError as exc:
                return exc.status, exc.body, "client_error"
            except RequestError as exc:
                return exc.status, exc.body(), "client_error"
            except Overloaded as exc:
                return (*_route_error(route, exc.status, str(exc)), "overloaded")
            except FillError as exc:
                span.record_exception(exc)
                if route == "fill":
                    return 500, {"error": {"type": "fill_error", "status": 500,
                                           "detail": exc.errors}}, "server_error"
                return (*_route_error(route, 500, f"fill failed: {exc}"), "server_error")
            except RuntimeError as exc:
                span.record_exception(exc)
                if "out of memory" in str(exc).lower():
                    return (*_route_error(route, 529 if route == "systemone" else 503,
                                          "model overloaded"), "overloaded")
                return (*_route_error(route, 500, f"{type(exc).__name__}: {exc}"),
                        "server_error")
            except Exception as exc:  # noqa: BLE001 -- reported, never a silent answer
                span.record_exception(exc)
                return (*_route_error(route, 500, f"{type(exc).__name__}: {exc}"),
                        "server_error")

        def do_GET(self) -> None:  # noqa: N802
            self._handle("GET")

        def do_POST(self) -> None:  # noqa: N802
            self._handle("POST")

    return Handler


class DegHTTPServer(ThreadingHTTPServer):
    request_queue_size = 256
    daemon_threads = True


def make_server(engine: Engine, host: str, port: int, *, api_key: str | None = None
                ) -> DegHTTPServer:
    return DegHTTPServer((host, port), make_handler(engine, api_key=api_key))


def model_fill(ckpt: str | None, device: str, generate: str | None = None,
               null_margin: float = 0.0) -> tuple[FillFn, FillManyFn | None]:
    """``fill`` (and ``fill_many`` when the runtime exposes one) bound to a loaded model.
    ``generate="base"``: x-generate fields decode with the frozen base weights (hybrid).
    ``null_margin``: the extract null margin of requests that do not set one."""
    from deg.model import load_model

    return bind_model(load_model(ckpt, device=device, generate=generate),
                      null_margin=null_margin)


def with_null_margin(options: FillOptions, default: float) -> FillOptions:
    """``options`` with the server's default null margin when the request set none."""
    if options.null_margin is not None:
        return options
    return dataclasses.replace(options, null_margin=float(default))


def bind_model(model: Any, null_margin: float = 0.0) -> tuple[FillFn, FillManyFn | None]:
    """``fill`` and ``fill_many`` (auto-detected on :mod:`deg.infer`) bound to ``model``;
    requests without ``options.null_margin`` get ``null_margin``."""
    from deg import infer

    def fill(context: Any, schema: dict[str, Any], options: FillOptions) -> dict[str, Any]:
        return infer.fill(model, context, schema, with_null_margin(options, null_margin))

    many = getattr(infer, "fill_many", None)
    fill_many = (lambda jobs: many(model, [(c, s, with_null_margin(o, null_margin))
                                           for c, s, o in jobs])) if callable(many) else None
    return fill, fill_many


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description="deg HTTP server (decisions, systemone, fill)")
    ap.add_argument("--ckpt", help="deg checkpoint directory (omit: base model, untrained "
                                   "heads -- for smoke tests only)")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8790)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--batch-window-ms", type=float, default=5.0)
    ap.add_argument("--max-batch", type=int, default=16)
    ap.add_argument("--max-pending", type=int, default=256)
    ap.add_argument("--generate", type=generate_arg, default=None,
                    help="x-generate decoder: deg's own weights or the frozen base "
                         "Qwen3.5 (hybrid; +~1 GB), or base+lora:<dir> = that hybrid with a "
                         "GENERATE LoRA (deg.gentrain) merged into the base copy. "
                         "Default: the checkpoint's generate_backbone, else self")
    ap.add_argument("--null-margin", type=float, default=0.0,
                    help="extract null margin in nats for requests that do not set "
                         "options.null_margin: a nullable/optional extract leaf is null when "
                         "best span log-prob < null log-prob + margin (> 0 = fewer made-up "
                         "answers, more omissions; 0 = argmax). See SPEC.md 12.1")
    ap.add_argument("--no-compat-models", action="store_true",
                    help="accept only deg model names (not gpt-6-luna / jev-latest)")
    a = ap.parse_args(argv)
    if a.ckpt is None:
        print("warning: no --ckpt; serving the base model with untrained heads",
              file=sys.stderr, flush=True)
    if not math.isfinite(a.null_margin):
        ap.error("--null-margin must be finite")
    fill, fill_many = model_fill(a.ckpt, a.device, a.generate, a.null_margin)
    batcher = Batcher(fill, fill_many=fill_many, window_s=a.batch_window_ms / 1000,
                      max_batch=a.max_batch, max_pending=a.max_pending)
    engine = (Engine(batcher, openai_models=MODEL_ALIASES, jev_models=MODEL_ALIASES)
              if a.no_compat_models else Engine(batcher))
    server = make_server(engine, a.host, a.port, api_key=os.environ.get("DEG_API_KEY") or None)
    print(json.dumps({"serving": f"http://{a.host}:{a.port}", "model": MODEL_ID,
                      "ckpt": a.ckpt, "generate": a.generate, "null_margin": a.null_margin,
                      "fill_many": fill_many is not None}), flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()