"""HTTP server for real PyTorch-backed shard nodes."""

from __future__ import annotations

import base64
import http.server
import json
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path
from typing import Any

from .model_backend import (
    InsufficientVRAMError,
    MissingModelDependencyError,
    Quantization,
    TorchModelShard,
    validate_quantization,
)
from .server import (
    _WIRE_VERSION,
    _compress_body,
    _decompress_body,
    _parse_shape,
    _validate_activation_body,
)


def _relay_hop(
    relay_addr: str,
    path: str,
    body: bytes,
    headers: dict[str, str],
    timeout: float = 120.0,
) -> tuple[int, dict[str, str], bytes]:
    """Send a single HTTP-shaped request through a relay RPC WebSocket.

    relay_addr is the wss://relay.../rpc/{peer_id} URL.
    Returns (status, response_headers_lower, response_body).
    Raises on connection failure so callers can fall back to direct.
    """
    import websockets.sync.client as wsc  # type: ignore[import]

    request_id = f"{time.time_ns():x}"
    payload = json.dumps({
        "request_id": request_id,
        "method": "POST",
        "path": path,
        "headers": headers,
        "body_base64": base64.b64encode(body).decode(),
    })
    with wsc.connect(relay_addr, open_timeout=timeout) as ws:
        ws.send(payload)
        raw = ws.recv(timeout=timeout)
    resp = json.loads(raw)
    status = int(resp.get("status", 503))
    resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
    body_b64 = resp.get("body_base64")
    resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
    return status, resp_headers, resp_body


class _TorchHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr,
        handler,
        backend: TorchModelShard,
        tracker_mode: bool = False,
        tracker_url: str | None = None,
        route_timeout: float = 30.0,
        debug: bool = False,
        max_loaded_shards: int = 1,
    ):
        super().__init__(addr, handler)
        self.backend = backend
        self.backends: dict[str, TorchModelShard] = {backend.model_id: backend}
        self.received_activations = False
        self.forward_chunk_count = 0
        self.tracker_mode = tracker_mode
        self.tracker_url = tracker_url
        self.route_timeout = route_timeout
        self.debug = debug
        self.max_loaded_shards = max(1, max_loaded_shards)
        self.total_requests: int = 0
        self.failed_requests: int = 0
        self.queue_depth: int = 0
        self._stats_lock = threading.Lock()

    def resolve_backend(self, model_name: str | None) -> TorchModelShard | None:
        if not model_name:
            return self.backend
        wanted = model_name.strip().lower()
        for key, shard_backend in self.backends.items():
            key_l = key.lower()
            if key_l == wanted or key_l.rsplit("/", 1)[-1] == wanted:
                return shard_backend
        return self.backend

    def chat_enabled(self) -> bool:
        return any(
            shard_backend.is_head
            for shard_backend in self.backends.values()
        )


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

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

    def _handle_infer(self) -> None:
        body = self._read_json_body()
        if body is None:
            return
        messages = body.get("messages", [])
        prompt = ""
        if isinstance(messages, list) and messages:
            last = messages[-1]
            if isinstance(last, dict):
                prompt = str(last.get("content", ""))
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        try:
            payload = server.backend.encode_prompt(prompt)
            if server.backend.is_tail:
                text = server.backend.decode_tail(
                    server.backend.torch.frombuffer(
                        bytearray(payload.body),
                        dtype=server.backend.torch.bfloat16,
                    )
                    .reshape(payload.shape)
                    .to(server.backend.device)
                )
                self._send_json(200, {"text": text})
                return
            self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}})
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})

    def _handle_forward(self) -> None:
        content_type = self.headers.get("Content-Type", "")
        if content_type.startswith("application/json"):
            self._handle_prompt_forward()
            return
        self._handle_binary_forward()

    def _handle_prompt_forward(self) -> None:
        body = self._read_json_body()
        if body is None:
            return
        prompt = str(body.get("prompt", ""))
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        try:
            payload = server.backend.encode_prompt(prompt)
        except Exception as exc:
            self._send_json(400, {"error": str(exc)})
            return
        self._send_activation(payload)

    def _handle_binary_forward(self) -> None:
        server: _TorchHTTPServer = 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)
            if dtype != "bfloat16":
                raise ValueError("real model backend requires bfloat16 activation input")
            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

        start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
        start_layer = int(start_layer_header) if start_layer_header else None

        try:
            result = server.backend.forward_bytes(
                raw_body,
                shape,
                self.headers.get("X-Meshnet-Attn-Mask"),
                self.headers.get("X-Meshnet-Position-Ids"),
                start_layer=start_layer,
            )
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})
            return

        if isinstance(result, str):
            self._send_json(200, {"text": result})
            return

        response_body = _compress_body(result.body, encoding)
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(response_body)))
        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape))
        self.send_header("X-Meshnet-Dtype", "bfloat16")
        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 result.attention_mask_header:
            self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
        if result.position_ids_header:
            self.send_header("X-Meshnet-Position-Ids", result.position_ids_header)
        self.end_headers()
        self.wfile.write(response_body)

    def _send_activation(self, payload) -> None:
        body = payload.body
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape))
        self.send_header("X-Meshnet-Dtype", "bfloat16")
        if payload.attention_mask_header:
            self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header)
        if payload.position_ids_header:
            self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header)
        self.end_headers()
        self.wfile.write(body)

    def _read_json_body(self) -> dict | None:
        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 None
        if not isinstance(body, dict):
            self._send_json(400, {"error": "JSON body must be an object"})
            return None
        return body

    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()
        try:
            self.wfile.write(payload)
        except BrokenPipeError:
            pass  # client disconnected before we could respond — not an error

    def _handle_chat_completions(self) -> None:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        with server._stats_lock:
            server.total_requests += 1
            server.queue_depth += 1
        try:
            self._do_chat_completions(server)
        finally:
            with server._stats_lock:
                server.queue_depth -= 1

    def _record_failed_request(self) -> None:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        with server._stats_lock:
            server.failed_requests += 1

    def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
        body = self._read_json_body()
        if body is None:
            return
        messages = body.get("messages", [])
        if not isinstance(messages, list):
            messages = []
        stream = bool(body.get("stream", False))
        model_name = str(body.get("model", ""))
        backend = server.resolve_backend(model_name)
        if backend is None or not backend.is_head:
            self._send_json(400, {"error": "model not loaded on this node"})
            return
        max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
        temperature = float(body.get("temperature") or 1.0)
        top_p = float(body.get("top_p") or 1.0)

        # Fast path: this node owns the complete model — use HF generate() with KV cache.
        # Avoids the single-token-per-forward-pass limitation of the distributed path.
        if backend.is_head and backend.is_tail:
            try:
                if stream:
                    self._stream_openai_response(
                        backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
                        model_name,
                    )
                else:
                    text = backend.generate_text(messages, max_tokens, temperature, top_p)
                    self._send_openai_response(text, model_name, False, messages, backend=backend)
            except Exception as exc:
                self._record_failed_request()
                self._send_json(500, {"error": f"generation failed: {exc}"})
            return

        # Distributed path: autoregressive generation across shards.
        # We do N single-step forward passes (no cross-node KV cache), which is slow
        # but correct. Each step: head encodes current sequence → forwards through route
        # → tail returns the next token string → append → repeat.
        remaining_route = self._get_remaining_route(model_name, backend=backend)
        print(
            f"  [node] chat route model={model_name!r} max_tokens={max_tokens} "
            f"downstream={remaining_route}",
            flush=True,
        )
        if not remaining_route:
            self._send_openai_response(
                "error: no downstream route — check tracker connectivity",
                model_name, False, messages, backend=backend,
            )
            return

        # Format with chat template so the model knows it's in assistant mode.
        try:
            if hasattr(backend.tokenizer, "apply_chat_template"):
                prompt_text: str = backend.tokenizer.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False,
                )
            else:
                raise AttributeError("no apply_chat_template")
        except Exception:
            prompt_text = " ".join(
                str(m.get("content", ""))
                for m in messages
                if isinstance(m, dict) and m.get("role") == "user"
            )

        eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
        generated: list[str] = []
        current_text = prompt_text

        stream_emit = None
        if stream:
            stream_emit = self._start_openai_stream(model_name)

        for _ in range(max_tokens):
            try:
                payload = backend.encode_prompt(current_text)
            except Exception as exc:
                print(f"  [node] distributed encode error: {exc}", flush=True)
                break
            token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
            if not token_str:
                break
            # Stop on error responses or EOS.
            if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
                break
            if eos_token and token_str == eos_token:
                break
            generated.append(token_str)
            if stream_emit is not None:
                stream_emit(token_str)
            current_text = current_text + token_str

        result_text = "".join(generated)
        if stream_emit is not None:
            stream_emit(None)
            return
        self._send_openai_response(result_text, model_name, stream, messages, backend=backend)

    def _get_remaining_route(self, model: str, *, backend: TorchModelShard | None = None) -> list[dict]:
        """Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr.

        Fast path reads X-Meshnet-Route header injected by the tracker.
        Slow path queries the tracker's /v1/route endpoint as a fallback.
        start_layer tells each downstream node which layer to begin from,
        enabling correct execution when shard ranges overlap.
        """
        # Fast path: tracker pre-resolved the downstream route and injected it as a header.
        injected = self.headers.get("X-Meshnet-Route")
        if injected:
            try:
                route = json.loads(injected)
                if isinstance(route, list):
                    hops: list[dict] = []
                    for item in route:
                        if isinstance(item, dict):
                            hop = {
                                "endpoint": str(item["endpoint"]),
                                "start_layer": int(item.get("start_layer", 0)),
                            }
                            if item.get("relay_addr"):
                                hop["relay_addr"] = str(item["relay_addr"])
                            hops.append(hop)
                        elif isinstance(item, str):
                            hops.append({"endpoint": item, "start_layer": 0})
                    print(f"  [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True)
                    return hops
            except (json.JSONDecodeError, TypeError, KeyError):
                pass

        # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        active_backend = backend or server.backend
        if server.tracker_url is None:
            return []
        route_model = getattr(active_backend, "model_id", None) or model
        try:
            url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
            with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
                route_resp = json.loads(r.read())
            own_port = server.server_address[1]
            nodes_info = route_resp.get("nodes", [])
            hops = []
            covered_up_to: int | None = None
            for node_info in nodes_info:
                ep = node_info.get("endpoint", "")
                if ep.rstrip("/").endswith(f":{own_port}"):
                    covered_up_to = node_info.get("shard_end")
                    continue
                if covered_up_to is None:
                    covered_up_to = (node_info.get("shard_start") or 1) - 1
                hop = {"endpoint": ep, "start_layer": covered_up_to + 1}
                if node_info.get("relay_addr"):
                    hop["relay_addr"] = str(node_info["relay_addr"])
                hops.append(hop)
                covered_up_to = node_info.get("shard_end", covered_up_to)
            print(f"  [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True)
            return hops
        except Exception as exc:
            print(f"  [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
            return []

    def _run_downstream_pipeline(self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None) -> str:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        active_backend = backend or server.backend
        if not route:
            # Partial shard at tail: decode the activation from the previous node.
            # Full single-node (head+tail) is handled before entering this method.
            if active_backend.is_tail:
                try:
                    tensor = active_backend.torch.frombuffer(
                        bytearray(payload.body),  # type: ignore[union-attr]
                        dtype=active_backend.torch.bfloat16,
                    ).reshape(payload.shape).to(active_backend.device)  # type: ignore[union-attr]
                    return active_backend.decode_tail(tensor)
                except Exception as exc:
                    return f"decode error: {exc}"
            return "no downstream route available for non-tail shard"

        session = str(uuid.uuid4())
        shape = payload.shape  # type: ignore[union-attr]
        attn_mask = payload.attention_mask_header  # type: ignore[union-attr]
        pos_ids = payload.position_ids_header  # type: ignore[union-attr]
        current_body = payload.body  # type: ignore[union-attr]
        current_shape = shape
        current_attn = attn_mask
        current_pos = pos_ids

        for hop_index, hop in enumerate(route):
            node_url = hop["endpoint"]
            start_layer = hop.get("start_layer", 0)
            relay_addr = hop.get("relay_addr")
            if server.debug:
                print(
                    f"  [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}"
                    + (f" relay={relay_addr}" if relay_addr else ""),
                    flush=True,
                )
            headers: dict[str, str] = {
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Wire": _WIRE_VERSION,
                "X-Meshnet-Shape": ",".join(str(d) for d in current_shape),
                "X-Meshnet-Dtype": "bfloat16",
                "X-Meshnet-Session": session,
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": str(hop_index),
                "X-Meshnet-Start-Layer": str(start_layer),
            }
            if current_attn:
                headers["X-Meshnet-Attn-Mask"] = current_attn
            if current_pos:
                headers["X-Meshnet-Position-Ids"] = current_pos
            if relay_addr:
                try:
                    status, resp_headers, resp_body = _relay_hop(
                        relay_addr, "/forward", current_body, headers, timeout=120.0,
                    )
                    if status >= 400:
                        print(
                            f"  [node] relay hop {hop_index} returned {status} from {relay_addr}",
                            flush=True,
                        )
                        return f"pipeline error at {node_url} via relay: status {status}"
                except Exception as exc:
                    print(
                        f"  [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
                        f"falling back to direct {node_url}",
                        flush=True,
                    )
                    relay_addr = None  # fall through to direct
            if not relay_addr:
                req = urllib.request.Request(
                    f"{node_url}/forward",
                    data=current_body,
                    headers=headers,
                    method="POST",
                )
                try:
                    with urllib.request.urlopen(req, timeout=120.0) as r:
                        resp_body = r.read()
                        resp_headers = {k.lower(): v for k, v in r.headers.items()}
                except Exception as exc:
                    print(f"  [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
                    return f"pipeline error at {node_url}: {exc}"
            content_type = resp_headers.get("content-type", "")
            if "application/json" in content_type:
                try:
                    data = json.loads(resp_body)
                    text = str(data.get("text", ""))
                    if server.debug:
                        print(f"  [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
                    return text
                except json.JSONDecodeError:
                    return resp_body.decode("utf-8", errors="replace")
            # Binary activation — update and forward to next node
            shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
            current_shape = _parse_shape(shape_header)
            current_body = resp_body
            current_attn = resp_headers.get("x-meshnet-attn-mask")
            current_pos = resp_headers.get("x-meshnet-position-ids")
        return ""

    def _stream_openai_response(self, token_iter, model: str) -> None:
        """Stream tokens from an iterator as SSE chunks."""
        emit = self._start_openai_stream(model)
        for token_text in token_iter:
            if not token_text:
                continue
            emit(token_text)
        emit(None)

    def _start_openai_stream(self, model: str):
        """Open an OpenAI-compatible SSE response and return a token emitter."""
        chunk_id = "chatcmpl-node"
        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:
            try:
                self.wfile.write(f"data: {data}\n\n".encode())
                self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError):
                pass

        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))

        def emit_token(token_text: str | None) -> None:
            if token_text is None:
                _emit(json.dumps({
                    "id": chunk_id, "object": "chat.completion.chunk", "created": created,
                    "model": model,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
                }))
                try:
                    self.wfile.write(b"data: [DONE]\n\n")
                    self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    pass
                return
            _emit(json.dumps({
                "id": chunk_id, "object": "chat.completion.chunk", "created": created,
                "model": model,
                "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
            }))

        return emit_token

    def _send_openai_response(
        self,
        text: str,
        model: str,
        stream: bool,
        messages: list[dict] | None = None,
        backend: TorchModelShard | None = None,
    ) -> None:
        chunk_id = "chatcmpl-node"
        created = int(time.time())
        active_backend = backend or self.server.backend  # type: ignore[attr-defined]
        if not stream:
            usage = _usage_for_response(active_backend, messages or [], text)
            self._send_json(200, {
                "id": chunk_id,
                "object": "chat.completion",
                "created": created,
                "model": model,
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }],
                "usage": usage,
            })
            return
        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:
            try:
                self.wfile.write(f"data: {data}\n\n".encode())
                self.wfile.flush()
            except BrokenPipeError:
                pass

        _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"}],
        }))
        try:
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
        except BrokenPipeError:
            pass


def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
    prompt_tokens = _backend_token_count(
        backend,
        "count_prompt_tokens",
        messages,
        fallback=_fallback_message_token_count(messages),
    )
    completion_tokens = _backend_token_count(
        backend,
        "count_text_tokens",
        completion_text,
        fallback=_fallback_text_token_count(completion_text),
    )
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
    }


def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
    method: Any = getattr(backend, method_name, None)
    if callable(method):
        try:
            return max(0, int(method(value)))
        except Exception:
            pass
    return max(0, int(fallback))


def _fallback_message_token_count(messages: list[dict]) -> int:
    text = " ".join(
        str(message.get("content", ""))
        for message in messages
        if isinstance(message, dict)
    )
    return _fallback_text_token_count(text)


def _fallback_text_token_count(text: str) -> int:
    parts = text.split()
    if parts:
        return len(parts)
    return 1 if text else 0


class TorchNodeServer:
    """HTTP server backed by a HuggingFace causal language model shard."""

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 0,
        model_id: str = "openai-community/gpt2",
        shard_start: int = 0,
        shard_end: int = 6,
        quantization: str = "bfloat16",
        backend: TorchModelShard | None = None,
        tracker_mode: bool | None = None,
        tracker_url: str | None = None,
        route_timeout: float = 30.0,
        cache_dir: Path | None = None,
        debug: bool = False,
        max_loaded_shards: int = 1,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._max_loaded_shards = max(1, max_loaded_shards)
        self._backend = backend or _load_backend(
            model_id,
            shard_start,
            shard_end,
            quantization,
            cache_dir,
        )
        self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
        # Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
        self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
        self._tracker_url = tracker_url
        self._route_timeout = route_timeout
        self._cache_dir = cache_dir
        self._debug = debug
        self._server: _TorchHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    @property
    def route_timeout(self) -> float:
        return self._route_timeout

    @property
    def backend(self) -> TorchModelShard:
        return self._backend

    @property
    def received_activations(self) -> bool:
        return self._server.received_activations if self._server is not None else False

    @property
    def forward_chunk_count(self) -> int:
        return self._server.forward_chunk_count if self._server is not None else 0

    @property
    def total_requests(self) -> int:
        return self._server.total_requests if self._server is not None else 0

    @property
    def failed_requests(self) -> int:
        return self._server.failed_requests if self._server is not None else 0

    @property
    def queue_depth(self) -> int:
        return self._server.queue_depth if self._server is not None else 0

    @property
    def loaded_model_ids(self) -> list[str]:
        return list(self._backends.keys())

    def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
        """Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more)."""
        add_directive = next(
            (directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"),
            None,
        )
        load_directive = next(
            (directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
            None,
        )
        directive = add_directive or load_directive
        if directive is None:
            return None
        shard_start = int(directive["shard_start"])
        shard_end = int(directive["shard_end"])
        quantization = str(directive.get("quantization") or self._backend.quantization)
        model_id = str(directive.get("model") or self._backend.model_id)
        replacing = directive.get("action") == "LOAD_SHARD"
        if not replacing and len(self._backends) >= self._max_loaded_shards:
            print(
                f"  [node] WARNING: ignoring ADD_SHARD for {model_id!r} — "
                f"loaded {len(self._backends)}/{self._max_loaded_shards} slots full",
                flush=True,
            )
            return None
        action_label = "reassigned" if replacing else "additional"
        print(
            f"  [node] loading {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
            flush=True,
        )
        try:
            new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
        except TypeError:
            new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
        self._backends[model_id] = new_backend
        if replacing or shard_start == 0:
            self._backend = new_backend
            self._tracker_mode = shard_start == 0
        print(
            f"  [node] loaded {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
            flush=True,
        )
        if self._server is not None:
            self._server.backends = dict(self._backends)
            if replacing or shard_start == 0:
                self._server.backend = new_backend
                self._server.tracker_mode = self._tracker_mode
        return {
            "action": directive.get("action"),
            "model": model_id,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "quantization": quantization,
            "tracker_mode": shard_start == 0,
        }

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("TorchNodeServer is already running")
        self._server = _TorchHTTPServer(
            (self._host, self._requested_port),
            _TorchHandler,
            self._backend,
            self._tracker_mode,
            self._tracker_url,
            self._route_timeout,
            self._debug,
            self._max_loaded_shards,
        )
        self._server.backends = dict(self._backends)
        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


def _load_backend(
    model_id: str,
    shard_start: int,
    shard_end: int,
    quantization: str,
    cache_dir: Path | None = None,
) -> TorchModelShard:
    from .model_backend import load_torch_shard

    quant = validate_quantization(quantization)
    try:
        return load_torch_shard(model_id, shard_start, shard_end, quant, cache_dir)
    except MissingModelDependencyError:
        raise
    except InsufficientVRAMError as exc:
        print(f"ERROR: {exc}", file=sys.stderr, flush=True)
        raise