"""Stub node HTTP server for the Distributed Inference Network."""

import base64
import http.server
import json
import struct
import threading

# Activation wire format (contract for all shard nodes):
# JSON object: {"shape": [B, S, D], "dtype": "float32",
#               "data": "<base64-encoded little-endian IEEE 754 float32 bytes>",
#               "context": {"prompt": "<original user prompt>"}}
# shape: [batch_size, seq_len, hidden_dim] — all ints
# dtype: always "float32" in this network version
# data: base64(struct.pack(f"<{B*S*D}f", *values))
# context: opaque dict forwarded unchanged by every intermediate shard

_STUB_HIDDEN_DIM = 64
_STUB_SEQ_LEN = 4


def _make_stub_activations(prompt: str) -> dict:
    """Return zeroed stub activations carrying the original prompt as context."""
    shape = [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM]
    n = shape[0] * shape[1] * shape[2]
    data = base64.b64encode(struct.pack(f"<{n}f", *([0.0] * n))).decode()
    return {"shape": shape, "dtype": "float32", "data": data, "context": {"prompt": prompt}}


class _StubHTTPServer(http.server.HTTPServer):
    def __init__(self, addr, handler, shard_start: int, shard_end: int, is_last_shard: bool):
        super().__init__(addr, handler)
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.is_last_shard = is_last_shard
        self.received_activations: bool = False


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

    def do_POST(self):
        if self.path == "/v1/infer":
            self._handle_infer()
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_infer(self):
        server: _StubHTTPServer = self.server  # type: ignore[assignment]
        length = int(self.headers.get("Content-Length", 0))
        body = json.loads(self.rfile.read(length))

        if "activations" in body:
            # Shard-to-shard call: this node receives activations from the previous shard.
            server.received_activations = True
            prompt = body["activations"].get("context", {}).get("prompt", "")
            if server.is_last_shard:
                payload = json.dumps({"text": f"stub response to: {prompt}"}).encode()
            else:
                payload = json.dumps({"activations": _make_stub_activations(prompt)}).encode()
        else:
            # Gateway-to-first-shard call: this node receives the original messages.
            messages = body.get("messages", [])
            last_content = messages[-1]["content"] if messages else ""
            if server.is_last_shard:
                # Single-node shortcut: no activation hop needed.
                payload = json.dumps({"text": f"stub response to: {last_content}"}).encode()
            else:
                payload = json.dumps({"activations": _make_stub_activations(last_content)}).encode()

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


class StubNodeServer:
    """Single-process stub node that returns fixed inference responses.

    shard_start / shard_end define which transformer layer range this node owns.
    is_last_shard controls whether the node returns a text response (True) or
    activation tensors (False) after processing its shard.
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 0,
        shard_start: int = 0,
        shard_end: int = 31,
        is_last_shard: bool = True,
    ):
        self._host = host
        self._requested_port = port
        if shard_start > shard_end:
            raise ValueError("shard_start must be less than or equal to shard_end")
        self._shard_start = shard_start
        self._shard_end = shard_end
        self._is_last_shard = is_last_shard
        self._server: _StubHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    @property
    def received_activations(self) -> bool:
        """True if this node received an activation tensor since it was started."""
        return self._server.received_activations if self._server is not None else False

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

        self._server = _StubHTTPServer(
            (self._host, self._requested_port),
            _StubHandler,
            self._shard_start,
            self._shard_end,
            self._is_last_shard,
        )
        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