"""Local System One server for Drex v2.0.

    python serve.py --port 8000

POST /v1/systemone accepts one decision, or ``requests``: a list of decisions.
Several decisions that fit share one forward. This process scores the pointer head
of the checkpoint.
"""
import argparse
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from inference import decide, decide_batch, load


def reject_nonfinite_json(value):
    raise ValueError(f"non-finite JSON number: {value}")


def make_handler(tok, model, model_name):
    lock = threading.Lock()

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

        def log_message(self, fmt, *args):
            print("%s %s" % (self.address_string(), fmt % args))

        def _send(self, status, payload):
            body = json.dumps(payload).encode()
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            if self.close_connection:
                self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path.rstrip("/") == "/health":
                self._send(200, {"status": "ok", "model": model_name})
                return
            self._send(404, {"error": "not found"})

        def do_POST(self):
            if self.path.rstrip("/") != "/v1/systemone":
                self._send(404, {"error": "not found"})
                return
            try:
                length = int(self.headers.get("Content-Length", "0"))
            except ValueError:
                self.close_connection = True
                self._send(400, {"error": "Content-Length must be an integer"})
                return
            if length <= 0 or length > 1_000_000:
                self.close_connection = True
                self._send(400, {"error": "request body must be between 1 byte and 1 MB"})
                return
            try:
                request = json.loads(self.rfile.read(length), parse_constant=reject_nonfinite_json)
                if not isinstance(request, dict):
                    raise ValueError("request must be an object")
                if "model" in request and not isinstance(request["model"], str):
                    raise ValueError("model must be a string")
                with lock:
                    if "requests" in request:
                        if "state" in request or "questions" in request:
                            raise ValueError("a batch cannot also contain state or questions")
                        items = request["requests"]
                        if not isinstance(items, list) or not items:
                            raise ValueError("requests must be a non-empty list")
                        if len(items) > 1024:
                            raise ValueError("requests exceeds 1024 items")
                        results = decide_batch(tok, model, items)
                        payload = {"model": request.get("model") or model_name, "results": results}
                    else:
                        payload = {"model": request.get("model") or model_name, **decide(tok, model, request)}
            except ValueError as exc:
                self._send(400, {"error": str(exc)})
                return
            except Exception as exc:
                self.log_error("System One scoring failed: %r", exc)
                self._send(500, {"error": "internal scoring error"})
                return
            self._send(200, payload)

    return Handler


def main():
    parser = argparse.ArgumentParser(description="Serve Drex v2.0 at POST /v1/systemone.")
    parser.add_argument("--model", default=str(Path(__file__).resolve().parent), help="Checkpoint directory (default: this directory)")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--name", default="drex-v2.0")
    args = parser.parse_args()
    name = args.name
    tok, model = load(args.model)
    server = ThreadingHTTPServer((args.host, args.port), make_handler(tok, model, name))
    print(f"{name} listening on http://{args.host}:{args.port}/v1/systemone")
    server.serve_forever()


if __name__ == "__main__":
    main()