"""Gateway HTTP server — accepts OpenAI-format requests and routes to inference nodes."""

import http.server
import json
import threading
import time
import urllib.request


class _GatewayHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def do_POST(self):
        if self.path == "/v1/chat/completions":
            self._handle_chat_completions()
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_chat_completions(self):
        length = int(self.headers.get("Content-Length", 0))
        body = json.loads(self.rfile.read(length))

        node_url = self.server.node_url  # type: ignore[attr-defined]
        payload = json.dumps({"messages": body.get("messages", [])}).encode()
        req = urllib.request.Request(
            f"{node_url}/v1/infer",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            node_resp = json.loads(r.read())

        response_body = json.dumps({
            "id": "chatcmpl-stub",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.get("model", "stub-model"),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": node_resp["text"],
                },
                "finish_reason": "stop",
            }],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
            },
        }).encode()

        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(response_body)))
        self.end_headers()
        self.wfile.write(response_body)


class _GatewayHTTPServer(http.server.HTTPServer):
    def __init__(self, addr, handler, node_url: str):
        super().__init__(addr, handler)
        self.node_url = node_url


class GatewayServer:
    """HTTP gateway that forwards /v1/chat/completions to a registered node."""

    def __init__(self, node_url: str, host: str = "127.0.0.1", port: int = 0):
        self._node_url = node_url
        self._host = host
        self._requested_port = port
        self._server: _GatewayHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("GatewayServer is already running")

        self._server = _GatewayHTTPServer(
            (self._host, self._requested_port),
            _GatewayHandler,
            self._node_url,
        )
        self.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self.port

    def stop(self) -> None:
        if self._server is None:
            return

        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        self._server = None
        self._thread = None
        self.port = None