"""A local HTTP server for a DecisionNode Edge model, on the standard library only.

    python -m decisionnode_edge serve path/to/model [--host 127.0.0.1] [--port 8000] [--api-key FILE] [--backend ...]

    POST /v1/systemone   the /v1/systemone format used by public decision-model leaderboards:
                         {"model", "state", "questions", "assets"}; noul, choice and score questions; a noul is
                         answered {"type": "noul", "noul": p}, a choice and a score as on /v1/decide
    POST /v1/decide      the DecisionNode API's call: the runtime's own request and answers (EdgeModel.decide)
    GET  /v1/models      {"models": [{"name", "description", "release_date"}]}
    GET  /healthz        {"status": "ok", "model", "backend"} (no key needed)

Pictures: `images` as the runtime takes them, base64 `data` only (a local `path` is refused over HTTP: a caller must
never make the server read its files), and on /v1/systemone the `assets` map of that format, {"<ref>": {"type": "image",
"format": "png" | "jpg" | "jpeg", "data_base64": ...}}, each asset a picture whose id is its reference (the state names
it), in the map's order.

Errors carry a `detail` body: {"detail": {"error_type", "message"}} (a 422 as a list of field errors). A
request over the model's length limit says "maximum context length", a choice over 255 options "options per choice",
the words leaderboard harnesses look for to count a request as beyond the model rather than as a failure.

The model answers one request at a time (a lock); the server listens on the loopback address unless told otherwise.
`model` is required on /v1/systemone, as on the hosted API, and any name is accepted: the answer names the model that
answered.
"""

from __future__ import annotations

import argparse
import hmac
import json
import logging
import sys
import threading
import time
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

from .request import RequestError

MAX_BODY_BYTES = 64 << 20  # 16 pictures as base64 with room to spare
SYSTEMONE_TYPES = ("noul", "choice", "score")
ASSET_MEDIA_TYPES = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg"}
AUTH_MESSAGE = "Cannot authenticate with the server. Please check your API key and try again."
ROUTES = {"/v1/systemone": "POST", "/v1/decide": "POST", "/v1/models": "GET", "/healthz": "GET"}

log = logging.getLogger("decisionnode_edge.server")


class HTTPError(Exception):
    def __init__(self, status: int, error_type: str, message: str, headers: dict | None = None):
        super().__init__(message)
        self.status, self.error_type, self.message, self.headers = status, error_type, message, headers or {}


def error_body(status: int, error_type: str, message: str) -> dict:
    if status == 422:
        return {"detail": [{"type": "value_error", "loc": ["body"], "msg": message, "input": None}]}
    return {"detail": {"error_type": error_type, "message": message}}


def _images_over_http(images: Any) -> list:
    if images is None:
        return []
    if not isinstance(images, list):
        raise RequestError("images must be a list", 422)
    for image in images:
        if isinstance(image, dict) and image.get("path") is not None:
            raise RequestError("Over HTTP a picture is sent as base64 `data` (or an asset), never as a local `path`.")
    return list(images)


def _asset_images(assets: Any) -> list[dict]:
    """The /v1/systemone `assets` map as the runtime's images: the reference is the picture's id."""
    if assets is None:
        return []
    if not isinstance(assets, dict):
        raise RequestError("assets must be an object of references", 422)
    out = []
    for ref, asset in assets.items():
        if not isinstance(asset, dict) or asset.get("type", "image") not in ("image", None):
            raise RequestError(f"Asset {ref!r}: only image assets are taken.")
        media_type = ASSET_MEDIA_TYPES.get(str(asset.get("format", "")).lower())
        if media_type is None:
            raise RequestError(f"Asset {ref!r}: format must be one of png, jpg, jpeg.")
        if not isinstance(asset.get("data_base64"), str) or not asset["data_base64"]:
            raise RequestError(f"Asset {ref!r} needs data_base64.", 422)
        out.append({"id": ref, "data": asset["data_base64"], "media_type": media_type})
    return out


class Service:
    """The routes, without the socket: handle(method, path, headers, body) gives (status, JSON body, extra headers)."""

    def __init__(self, model: Any, api_key: str | None = None, description: str | None = None, release_date: str | None = None):
        self.model = model
        self.api_key = api_key
        self.name = model.config["name"]
        self.description = description or model.config.get("description") or f"{self.name}: typed, calibrated decisions, run locally"
        self.release_date = release_date or model.config.get("release_date")
        self.lock = threading.Lock()

    def handle(self, method: str, path: str, headers: dict[str, str], body: bytes) -> tuple[int, dict, dict]:
        try:
            return 200, self._route(method, path.split("?", 1)[0], headers, body), {}
        except HTTPError as error:
            return error.status, error_body(error.status, error.error_type, error.message), error.headers
        except RequestError as error:
            return error.status, error_body(error.status, error.error_type, str(error)), {}
        except Exception:  # a fault of the model or the backend, never the caller's
            log.error("the model could not answer:\n%s", traceback.format_exc())
            return 500, error_body(500, "internal_error", "The model could not answer this request."), {}

    def _route(self, method: str, path: str, headers: dict[str, str], body: bytes) -> dict:
        if path not in ROUTES:
            raise HTTPError(404, "not_found_error", "Not found.")
        if method != ROUTES[path]:
            raise HTTPError(405, "api_usage_error", "Method not allowed.", {"Allow": ROUTES[path]})
        if path == "/healthz":
            return {"status": "ok", "model": self.name, "backend": self.model.backend}
        self._authenticate(headers)
        if path == "/v1/models":
            return {"models": [{"name": self.name, "description": self.description, "release_date": self.release_date}]}
        try:
            request = json.loads(body)
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise RequestError(f"the body is not JSON: {error}", 422) from None
        if path == "/v1/decide":
            if isinstance(request, dict) and request.get("images") is not None:
                request = {**request, "images": _images_over_http(request["images"])}
            with self.lock:
                return self.model.decide(request)
        return self._systemone(request)

    def _authenticate(self, headers: dict[str, str]) -> None:
        if self.api_key is None:
            return
        given = headers.get("authorization", "")
        scheme, _, token = given.partition(" ")
        if scheme.lower() != "bearer" or not hmac.compare_digest(token.strip().encode(), self.api_key.encode()):
            raise HTTPError(401, "authentication_error", AUTH_MESSAGE, {"WWW-Authenticate": "Bearer"})

    def _systemone(self, request: Any) -> dict:
        if not isinstance(request, dict):
            raise RequestError("the request must be a JSON object", 422)
        if not isinstance(request.get("model"), str):
            raise RequestError("model: Field required, a string", 422)
        questions = request.get("questions")
        if isinstance(questions, dict):
            for qid, question in questions.items():
                kind = question.get("type") if isinstance(question, dict) else None
                if kind == "number":
                    raise RequestError(f"question {qid!r}: the number type is not available on /v1/systemone; send it to /v1/decide")
                if isinstance(question, dict) and "type" in question and kind not in SYSTEMONE_TYPES:
                    raise RequestError(f"question {qid!r}: unknown question type {kind!r}")
        images = _images_over_http(request.get("images")) + _asset_images(request.get("assets"))
        body = {"state": request.get("state"), "questions": questions}
        if images:
            body["images"] = images
        with self.lock:
            response = self.model.decide(body)
        answers = {}
        for qid, answer in response["answers"].items():  # a noul is answered under its own name, as the format has it
            answers[qid] = {"type": "noul", "noul": answer["truth"]} if answer["type"] == "truth" else answer
        return {"model": response["model"], "answers": answers, "usage": response["usage"]}


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"  # keep-alive: every response carries Content-Length
    disable_nagle_algorithm = True  # headers and body are two writes: without it a kept-alive client waits for a delayed ACK
    server_version = "decisionnode-edge"
    service: Service = None  # set by make_server

    def _answer(self, method: str) -> None:
        started = time.perf_counter()
        body = b""
        if method == "POST":
            length = self.headers.get("Content-Length")
            if length is None or not length.isdigit():
                return self._send(411, error_body(411, "api_usage_error", "Content-Length is required."), {"Connection": "close"}, started)
            if int(length) > MAX_BODY_BYTES:
                return self._send(413, error_body(413, "api_usage_error", "Request body too large."), {"Connection": "close"}, started)
            body = self.rfile.read(int(length))
        headers = {k.lower(): v for k, v in self.headers.items()}
        status, payload, extra = self.service.handle(method, self.path, headers, body)
        self._send(status, payload, extra, started)

    def _send(self, status: int, payload: dict, extra: dict, started: float) -> None:
        data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("X-Latency-Ms", f"{(time.perf_counter() - started) * 1000:.2f}")
        for name, value in extra.items():
            self.send_header(name, value)
        self.end_headers()
        self.wfile.write(data)
        if extra.get("Connection") == "close":
            self.close_connection = True

    def do_GET(self) -> None:  # noqa: N802 (the standard library's names)
        self._answer("GET")

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

    def do_PUT(self) -> None:  # noqa: N802
        self._answer("PUT")

    def do_DELETE(self) -> None:  # noqa: N802
        self._answer("DELETE")

    def log_message(self, format: str, *args: Any) -> None:  # noqa: A002
        log.info("%s %s", self.address_string(), format % args)


def make_server(model: Any, host: str = "127.0.0.1", port: int = 8000, api_key: str | None = None, **service: Any) -> ThreadingHTTPServer:
    handler = type("EdgeHandler", (Handler,), {"service": Service(model, api_key=api_key, **service)})
    server = ThreadingHTTPServer((host, port), handler)
    server.daemon_threads = True
    return server


def main(argv: list[str] | None = None) -> int:
    from . import EdgeModel

    parser = argparse.ArgumentParser(prog="python -m decisionnode_edge serve", description="Serve a DecisionNode Edge model over HTTP: /v1/systemone, /v1/decide, /v1/models.")
    parser.add_argument("model", help="the model directory")
    parser.add_argument("--host", default="127.0.0.1", help="the address to listen on (default 127.0.0.1, this machine only)")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--api-key", default=None, help="a file holding the key callers must send as `Authorization: Bearer <key>`")
    parser.add_argument("--device", default=None)
    parser.add_argument("--dtype", default=None)
    parser.add_argument("--backend", default=None, help="auto, torch, vllm, openvino, mlx or llamacpp (default: auto)")
    parser.add_argument("--artifacts", default=None, help="the directory python -m decisionnode_edge.build wrote")
    parser.add_argument("--description", default=None, help="the description GET /v1/models gives")
    parser.add_argument("--release-date", default=None, help="the release date GET /v1/models gives (YYYY-MM-DD)")
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
    api_key = None
    if args.api_key:
        with open(args.api_key, encoding="utf-8") as f:
            api_key = f.read().strip()
        if not api_key:
            raise SystemExit(f"{args.api_key} is empty")
    model = EdgeModel.load(args.model, device=args.device, dtype=args.dtype, backend=args.backend, artifacts=args.artifacts)
    model.decide({"state": "The color is red.", "questions": {"warmup": {"type": "choice", "instructions": "Which color is named?", "criteria": {"red": "red", "blue": "blue"}}}})
    server = make_server(model, args.host, args.port, api_key, description=args.description, release_date=args.release_date)
    log.info("serving %s (%s backend) on http://%s:%d%s", model.config["name"], model.backend, args.host, server.server_address[1], ", key required" if api_key else "")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    sys.exit(main())