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

import http.server
import json
import time
import threading
import urllib.parse
from pathlib import Path

from .downloader import compute_shard_checksum, write_shard_archive

# Binary activation wire format (contract for all shard nodes):
# POST /forward with raw tensor bytes in the body and tensor/session/chunk
# metadata in X-Meshnet-* headers. Bodies may be zstd-compressed.

_STUB_HIDDEN_DIM = 64
_STUB_SEQ_LEN = 4
_WIRE_VERSION = "2"
_DTYPE_SIZES = {
    "float16": 2,
    "bfloat16": 2,
    "float32": 4,
}


def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
    """Return a zeroed stub tensor body for the requested shape and dtype."""
    element_size = _DTYPE_SIZES[dtype]
    element_count = 1
    for dim in shape:
        element_count *= dim
    return b"\x00" * (element_count * element_size)


def _parse_shape(value: str | None) -> list[int]:
    if not value:
        raise ValueError("missing X-Meshnet-Shape")
    shape = [int(part) for part in value.split(",")]
    if not shape or any(dim <= 0 for dim in shape):
        raise ValueError("invalid X-Meshnet-Shape")
    return shape


def _decompress_body(body: bytes, encoding: str | None) -> bytes:
    if not encoding:
        return body
    if encoding != "zstd":
        raise ValueError("unsupported X-Meshnet-Encoding")
    import zstandard as zstd

    try:
        return zstd.ZstdDecompressor().decompress(body)
    except zstd.ZstdError as exc:
        raise ValueError("invalid zstd activation body") from exc


def _compress_body(body: bytes, encoding: str | None) -> bytes:
    if not encoding:
        return body
    if encoding != "zstd":
        raise ValueError("unsupported X-Meshnet-Encoding")
    import zstandard as zstd

    try:
        return zstd.ZstdCompressor(level=1).compress(body)
    except zstd.ZstdError as exc:
        raise ValueError("could not zstd-compress activation body") from exc


def _validate_activation_body(body: bytes, shape: list[int], dtype: str) -> None:
    if dtype not in _DTYPE_SIZES:
        raise ValueError("unsupported X-Meshnet-Dtype")
    expected_len = len(_make_stub_binary_activation(shape, dtype))
    if len(body) != expected_len:
        raise ValueError("activation byte length does not match shape and dtype")
    try:
        import numpy as np
    except ModuleNotFoundError:
        return
    if dtype == "bfloat16" and not hasattr(np, "bfloat16"):
        return
    np_dtype = getattr(np, "bfloat16") if dtype == "bfloat16" else np.dtype(dtype)
    np.frombuffer(body, dtype=np_dtype).reshape(shape)


class _StubHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr,
        handler,
        shard_start: int,
        shard_end: int,
        is_last_shard: bool,
        response_prefix: str,
        model: str,
        shard_path: Path | None,
        tracker_mode: bool = False,
        tracker_url: str | None = None,
    ):
        super().__init__(addr, handler)
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.is_last_shard = is_last_shard
        self.response_prefix = response_prefix
        self.model = model
        self.shard_path = shard_path
        self.received_activations: bool = False
        self.forward_chunk_count: int = 0
        self.tracker_mode: bool = tracker_mode
        self.tracker_url: str | None = tracker_url
        self.chat_completion_count: int = 0


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

    def do_POST(self):
        server: _StubHTTPServer = self.server  # type: ignore[assignment]
        if self.path == "/v1/infer":
            self._handle_infer()
        elif self.path == "/forward":
            self._handle_forward()
        elif self.path == "/v1/chat/completions" and server.tracker_mode:
            self._handle_chat_completions()
        else:
            self.send_response(404)
            self.end_headers()

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/v1/shards/download":
            self._handle_shard_download(parsed)
        else:
            self.send_response(404)
            self.end_headers()

    def _send_json(self, status: int, data: dict) -> None:
        payload = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _handle_chat_completions(self) -> None:
        server: _StubHTTPServer = self.server  # type: ignore[assignment]
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length) or b"{}")
        except (json.JSONDecodeError, ValueError):
            self._send_json(400, {"error": "invalid JSON body"})
            return
        if not isinstance(body, dict):
            self._send_json(400, {"error": "JSON body must be an object"})
            return
        server.chat_completion_count += 1
        streaming = bool(body.get("stream", False))
        model = str(body.get("model", server.model))
        messages = body.get("messages", [])
        last_content = ""
        if isinstance(messages, list) and messages:
            last = messages[-1]
            if isinstance(last, dict):
                last_content = str(last.get("content", ""))
        text = f"{server.response_prefix} {last_content}"
        if streaming:
            self._send_sse_response(text, model)
        else:
            created = int(time.time())
            self._send_json(200, {
                "id": "chatcmpl-stub",
                "object": "chat.completion",
                "created": created,
                "model": model,
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
            })

    def _send_sse_response(self, text: str, model: str) -> None:
        chunk_id = "chatcmpl-stub"
        created = int(time.time())
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()

        def _emit(data: str) -> None:
            self.wfile.write(f"data: {data}\n\n".encode())
            self.wfile.flush()

        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))
        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
        }))
        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }))
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()

    def _handle_shard_download(self, parsed: urllib.parse.ParseResult):
        server: _StubHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model = params.get("model", [""])[0]
        try:
            shard_start = int(params.get("shard_start", [""])[0])
            shard_end = int(params.get("shard_end", [""])[0])
        except ValueError:
            self.send_response(400)
            self.end_headers()
            return

        if (
            server.shard_path is None
            or not server.shard_path.exists()
            or model != server.model
            or shard_start != server.shard_start
            or shard_end != server.shard_end
        ):
            self.send_response(404)
            self.end_headers()
            return

        self.send_response(200)
        self.send_header("Content-Type", "application/x-tar")
        self.send_header("X-Meshnet-Shard-Checksum", compute_shard_checksum(server.shard_path))
        self.end_headers()
        write_shard_archive(server.shard_path, self.wfile)

    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"{server.response_prefix} {prompt}"}).encode()
            else:
                payload = json.dumps({
                    "activations": {
                        "shape": [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM],
                        "dtype": "bfloat16",
                    }
                }).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"{server.response_prefix} {last_content}"}).encode()
            else:
                payload = json.dumps({
                    "activations": {
                        "shape": [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM],
                        "dtype": "bfloat16",
                    }
                }).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)

    def _handle_forward(self):
        server: _StubHTTPServer = self.server  # type: ignore[assignment]
        try:
            shape = _parse_shape(self.headers.get("X-Meshnet-Shape"))
            dtype = self.headers.get("X-Meshnet-Dtype", "")
            session = self.headers["X-Meshnet-Session"]
            chunk_index = self.headers["X-Meshnet-Chunk-Index"]
            chunk_total = self.headers["X-Meshnet-Chunk-Total"]
            encoding = self.headers.get("X-Meshnet-Encoding")
            length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(length)
            raw_body = _decompress_body(body, encoding)
            _validate_activation_body(raw_body, shape, dtype)
            chunk_index_value = int(chunk_index)
            chunk_total_value = int(chunk_total)
            if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value:
                raise ValueError("invalid chunk index/total")
        except (KeyError, ValueError, TypeError):
            self.send_response(400)
            self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
            self.end_headers()
            return

        server.forward_chunk_count += 1
        if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
            server.received_activations = True

        raw_payload = _make_stub_binary_activation(shape, dtype)
        payload = _compress_body(raw_payload, encoding)
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(payload)))
        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in shape))
        self.send_header("X-Meshnet-Dtype", dtype)
        self.send_header("X-Meshnet-Session", session)
        self.send_header("X-Meshnet-Chunk-Index", chunk_index)
        self.send_header("X-Meshnet-Chunk-Total", chunk_total)
        if encoding:
            self.send_header("X-Meshnet-Encoding", encoding)
        if server.is_last_shard:
            self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
        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.
    tracker_mode enables the /v1/chat/completions endpoint for head-shard nodes.
    """

    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,
        response_prefix: str = "stub response to:",
        model: str = "stub-model",
        shard_path: Path | None = None,
        tracker_mode: bool = False,
        tracker_url: str | None = None,
    ):
        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._response_prefix = response_prefix
        self._model = model
        self._shard_path = shard_path
        self._tracker_mode = tracker_mode
        self._tracker_url = tracker_url
        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

    @property
    def forward_chunk_count(self) -> int:
        """Number of binary /forward chunks handled since this node was started."""
        return self._server.forward_chunk_count if self._server is not None else 0

    @property
    def chat_completion_count(self) -> int:
        """Number of /v1/chat/completions requests handled since this node was started."""
        return self._server.chat_completion_count if self._server is not None else 0

    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._response_prefix,
            self._model,
            self._shard_path,
            self._tracker_mode,
            self._tracker_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