"""HTTP server for Lychee-27B.

    python -m lychee.server --model <repo dir> [--host 127.0.0.1] [--port 8000]

GET  /health            200 {"status": "ok"} once the model is loaded
GET  /v1/models         the served model
POST /v1/systemone      {"state": ..., "questions": {name: {"type", "instructions", "criteria"}}, "model": optional}
                        -> {"model", "answers": {name: answer}, "usage": {"input_tokens", "output_tokens": 0}}

One request runs at a time; others wait their turn. With LYCHEE_API_KEY set, requests must send
`Authorization: Bearer <key>`. Questions the model cannot answer as asked (too many options, a prompt over the
maximum context length, images) get HTTP 422 with the reason.
"""
from __future__ import annotations

import argparse
import hmac
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from lychee.model import Lychee, RequestError

MODEL_NAME = "lychee-27b"


def make_handler(model: Lychee):
    lock = threading.Lock()
    key = os.environ.get("LYCHEE_API_KEY", "")

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass

        def reply(self, status: int, body: dict):
            data = json.dumps(body).encode()
            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:
            if not key:
                return True
            sent = self.headers.get("Authorization", "")
            return hmac.compare_digest(sent.encode(), f"Bearer {key}".encode())

        def do_GET(self):
            path = self.path.rstrip("/")
            if path == "/health":
                return self.reply(200, {"status": "ok", "model": MODEL_NAME})
            if path == "/v1/models":
                if not self.authorized():
                    return self.reply(401, {"detail": "missing or invalid API key"})
                return self.reply(200, {"object": "list", "data": [{"id": MODEL_NAME, "object": "model"}]})
            self.reply(404, {"detail": "not found"})

        def do_POST(self):
            if self.path.rstrip("/") != "/v1/systemone":
                return self.reply(404, {"detail": "not found"})
            if not self.authorized():
                return self.reply(401, {"detail": "missing or invalid API key"})
            try:
                body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))))
            except (ValueError, json.JSONDecodeError):
                return self.reply(400, {"detail": "the body must be JSON"})
            if not isinstance(body, dict) or "state" not in body:
                return self.reply(422, {"detail": "a request needs a state and questions"})
            if body.get("images"):
                return self.reply(422, {"detail": "images are not supported by this server"})
            try:
                with lock:
                    answers, tokens = model.decide(body["state"], body.get("questions"))
            except RequestError as error:
                return self.reply(422, {"detail": str(error)})
            self.reply(200, {"model": MODEL_NAME, "answers": answers,
                             "usage": {"input_tokens": tokens, "output_tokens": 0}})

    return Handler


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True, help="the repository directory")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()
    model = Lychee(args.model, device=args.device)
    ThreadingHTTPServer((args.host, args.port), make_handler(model)).serve_forever()


if __name__ == "__main__":
    main()