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

from __future__ import annotations

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

from .model_backend import (
    InsufficientVRAMError,
    KVCacheMiss,
    MissingModelDependencyError,
    Quantization,
    TailTokenResult,
    TorchModelShard,
    _tensor_from_bfloat16_bytes,
    validate_quantization,
)
from .seam_telemetry import GenerationTelemetry
from .activation_compression import (
    CompressionPolicies,
    CompressionPolicy,
    compress_activation,
    decompress_activation,
)


class _PipelineCacheMiss(Exception):
    """A downstream hop reported 409 cache_miss — head must re-prefill."""


class _RelayRequestUncertainError(ConnectionError):
    """A relay request may have reached the peer but produced no response."""


class _DirectRequestUncertainError(ConnectionError):
    """A direct request may have reached the downstream node but did not finish."""


from .server import (
    _WIRE_VERSION,
    _parse_shape,
    _validate_activation_body,
)


def _endpoint_key(url: str) -> str:
    """Normalize http(s) endpoints for host:port comparison."""
    parsed = urllib.parse.urlparse(url.rstrip("/"))
    host = (parsed.hostname or "").lower()
    if not host:
        return url.rstrip("/").lower()
    port = parsed.port
    if port is None:
        port = 443 if parsed.scheme == "https" else 80
    return f"{host}:{port}"


def _own_endpoint_key(server: _TorchHTTPServer) -> str:
    advertised = getattr(server, "advertised_endpoint", None)
    if advertised:
        return _endpoint_key(advertised)
    host, port = server.server_address
    return _endpoint_key(f"http://{host}:{port}")


def _clamp_downstream_hops(
    hops: list[dict],
    backend: TorchModelShard | None,
) -> list[dict]:
    """Ensure downstream start_layer continues after this shard's layers."""
    if not hops or backend is None:
        return hops
    shard_end = getattr(backend, "shard_end", None)
    if shard_end is None:
        return hops
    min_start = int(shard_end) + 1
    clamped: list[dict] = []
    for hop in hops:
        adjusted = dict(hop)
        if int(adjusted.get("start_layer", 0)) < min_start:
            adjusted["start_layer"] = min_start
        clamped.append(adjusted)
    return clamped


def _format_downstream_route(hops: list[dict]) -> str:
    return ", ".join(
        f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops
    )


def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None:
    """Rewrite one in-place progress line (\\r) or finish with a newline."""
    if final:
        if state[0]:
            sys.stdout.write("\r" + message + "\n")
            state[0] = False
        else:
            print(message, flush=True)
        return
    if state[0]:
        sys.stdout.write("\r" + message)
    else:
        sys.stdout.write(message)
        state[0] = True
    sys.stdout.flush()


class _RelayHopClient:
    """Persistent relay connection scoped to one generation handler."""

    def __init__(self, relay_addr: str, timeout: float = 120.0) -> None:
        self.relay_addr = relay_addr
        self.timeout = timeout
        self._ws = None
        self._lock = threading.RLock()

    def request(
        self,
        path: str,
        body: bytes,
        headers: dict[str, str],
    ) -> tuple[int, dict[str, str], bytes]:
        import websockets.sync.client as wsc  # type: ignore[import]

        from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size

        with self._lock:
            if self._ws is None:
                self._ws = wsc.connect(
                    self.relay_addr,
                    open_timeout=self.timeout,
                    max_size=ws_max_size(),
                    compression=None,
                )
            request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex
            frame = encode_binary_frame({
                "request_id": request_id,
                "method": "POST",
                "path": path,
                "headers": headers,
            }, body)
            try:
                # A send failure is uncertain too: bytes may already have been
                # accepted by the kernel or peer before the exception surfaced.
                self._ws.send(frame)
                raw = self._ws.recv(timeout=self.timeout)
                if isinstance(raw, (bytes, bytearray)):
                    resp_header, resp_body = decode_binary_frame(bytes(raw))
                    response_id = str(resp_header.get("request_id") or "")
                    status = int(resp_header.get("status", 503))
                    resp_headers = {
                        k.lower(): v for k, v in (resp_header.get("headers") or {}).items()
                    }
                else:
                    resp = json.loads(raw)
                    response_id = str(resp.get("request_id") or "")
                    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()
                    )
                if response_id and response_id != request_id:
                    raise ValueError("relay response request_id did not match request")
                return status, resp_headers, resp_body
            except Exception as exc:
                self.close()
                raise _RelayRequestUncertainError(
                    "relay connection failed after forwarding request; refusing replay"
                ) from exc

    def close(self) -> None:
        with self._lock:
            if self._ws is not None:
                try:
                    self._ws.close()
                except Exception:
                    pass
                finally:
                    self._ws = None


class _DirectHopClient:
    """One serialized HTTP/1.1 connection to one downstream hop.

    Generation handlers own these clients, so a cached Route Session reuses a
    TCP connection without sharing it between concurrent Route Sessions.
    """

    def __init__(self, endpoint: str, timeout: float = 120.0) -> None:
        parsed = urllib.parse.urlsplit(endpoint.rstrip("/"))
        if parsed.scheme not in {"http", "https"} or not parsed.hostname:
            raise ValueError(f"invalid downstream endpoint: {endpoint!r}")
        self.timeout = timeout
        self._scheme = parsed.scheme
        self._host = parsed.hostname
        self._port = parsed.port
        self._base_path = parsed.path.rstrip("/")
        self._connection: http.client.HTTPConnection | None = None
        self._lock = threading.RLock()

    def _connect(self) -> http.client.HTTPConnection:
        connection_type = (
            http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
        )
        return connection_type(self._host, self._port, timeout=self.timeout)

    def request(
        self, path: str, body: bytes, headers: Mapping[str, str],
    ) -> tuple[int, dict[str, str], bytes]:
        request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
        with self._lock:
            if self._connection is None:
                self._connection = self._connect()
            try:
                self._connection.request("POST", request_path, body=body, headers=dict(headers))
                response = self._connection.getresponse()
                response_body = response.read()
                response_headers = {key.lower(): value for key, value in response.headers.items()}
                return response.status, response_headers, response_body
            except Exception as exc:
                self.close()
                raise _DirectRequestUncertainError(
                    "direct connection failed after forwarding request; refusing replay"
                ) from exc

    def close(self) -> None:
        with self._lock:
            if self._connection is not None:
                try:
                    self._connection.close()
                finally:
                    self._connection = None


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 one request through a short-lived relay client (compatibility API).

    relay_addr is the wss://relay.../rpc/{peer_id} URL. The request and any
    binary response travel as binary frames (JSON header + raw body); relay
    error responses and legacy peers still answer with JSON text frames.
    Returns (status, response_headers_lower, response_body).
    Raises on connection failure so callers can fall back to direct.
    """
    client = _RelayHopClient(relay_addr, timeout)
    try:
        return client.request(path, body, headers)
    finally:
        client.close()


_COMPRESSION_POLICIES = CompressionPolicies()


def _maybe_compress_activation(
    body: bytes, policy: CompressionPolicy | None = None,
) -> tuple[bytes, str | None]:
    """Compatibility wrapper for callers that only need wire body and encoding."""
    result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill"))
    return result.body, result.encoding


def _is_cache_miss_body(body: bytes) -> bool:
    try:
        return json.loads(body).get("error") == "cache_miss"
    except (json.JSONDecodeError, AttributeError, UnicodeDecodeError):
        return False


def _response_error_snippet(body: bytes, limit: int = 500) -> str:
    """Return a compact error string from a downstream JSON/text response body."""
    try:
        payload = json.loads(body)
        if isinstance(payload, dict):
            message = payload.get("error") or payload.get("detail") or payload
            return str(message)[:limit]
    except (json.JSONDecodeError, TypeError, UnicodeDecodeError):
        pass
    return body.decode("utf-8", errors="replace")[:limit]


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.advertised_endpoint: str | None = None
        self.total_requests: int = 0
        self.failed_requests: int = 0
        self.queue_depth: int = 0
        self._stats_lock = threading.Lock()
        self._active_requests: dict[str, dict[str, Any]] = {}
        self._decode_log: dict[str, dict[str, float]] = {}

    def note_decode_step(
        self, session: str, now: float | None = None,
    ) -> int | None:
        """Count one decode forward; return the cumulative step count when a
        log line is due (first step of a session, then every 5s), else None."""
        if now is None:
            now = time.monotonic()
        with self._stats_lock:
            rec = self._decode_log.get(session)
            if rec is None:
                if len(self._decode_log) >= 64:
                    stale = [
                        sid for sid, r in self._decode_log.items()
                        if now - r["seen"] > 600.0
                    ]
                    for sid in stale:
                        del self._decode_log[sid]
                    while len(self._decode_log) >= 64:
                        self._decode_log.pop(next(iter(self._decode_log)))
                self._decode_log[session] = {"steps": 1.0, "logged": now, "seen": now}
                return 1
            rec["steps"] += 1
            rec["seen"] = now
            if now - rec["logged"] >= 5.0:
                rec["logged"] = now
                return int(rec["steps"])
            return None

    def snapshot_current_requests(self) -> list[dict[str, Any]]:
        """In-flight request snapshots for tracker heartbeats."""
        now = time.monotonic()
        with self._stats_lock:
            out: list[dict[str, Any]] = []
            for rec in self._active_requests.values():
                elapsed = max(now - float(rec["started"]), 1e-6)
                tokens = int(rec.get("tokens") or 0)
                out.append({
                    "request_id": str(rec["request_id"]),
                    "model": str(rec.get("model") or ""),
                    "kind": str(rec.get("kind") or "chat"),
                    "tokens": tokens,
                    "elapsed_seconds": round(elapsed, 1),
                    "tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
                    "routing_complete": bool(rec.get("routing_complete")),
                    "telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None,
                })
            return out

    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):
    # HTTP/1.1 is required for Route Session-owned downstream connections.
    # Finite responses below provide Content-Length; streams are chunked.
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def _request_id(self) -> str:
        return (
            self.headers.get("X-Meshnet-Request-Id")
            or self.headers.get("X-Request-Id")
            or f"local-{time.time_ns():x}"
        )

    def _request_log_suffix(self) -> str:
        activation_id = self.headers.get("X-Meshnet-Activation-Id")
        if activation_id:
            return f" activation_id={activation_id}"
        req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id")
        return f" request_id={req_id}" if req_id else ""

    def _track_request_begin(
        self,
        server: "_TorchHTTPServer",
        request_id: str,
        model: str,
    ) -> None:
        with server._stats_lock:
            server._active_requests[request_id] = {
                "request_id": request_id,
                "model": model,
                "kind": "chat",
                "started": time.monotonic(),
                "tokens": 0,
                "routing_complete": False,
                "telemetry": None,
            }

    def _track_request_progress(
        self,
        server: "_TorchHTTPServer",
        request_id: str,
        *,
        tokens: int,
        routing_complete: bool = False,
    ) -> None:
        with server._stats_lock:
            rec = server._active_requests.get(request_id)
            if rec is None:
                return
            rec["tokens"] = tokens
            if routing_complete:
                rec["routing_complete"] = True

    def _track_request_end(self, server: "_TorchHTTPServer", request_id: str) -> None:
        with server._stats_lock:
            server._active_requests.pop(request_id, None)

    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.send_header("Content-Length", "0")
            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_activation(body, encoding).body
            _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.send_header("Content-Length", "0")
            self.end_headers()
            return

        server.forward_chunk_count += 1
        hop_index = int(self.headers.get("X-Meshnet-Hop-Index", "0"))
        if hop_index > 0:
            server.received_activations = True
        # Session KV-cache protocol: prefill establishes per-session state on
        # this node's layer range; decode reuses it. Absent header = legacy
        # stateless call (also the signature fake backends implement).
        cache_mode = self.headers.get("X-Meshnet-Cache")
        if chunk_index_value == 0:
            shard_start = getattr(server.backend, "shard_start", "?")
            shard_end = getattr(server.backend, "shard_end", "?")
            if cache_mode == "decode":
                # One decode forward arrives per generated token — log a
                # periodic per-session summary instead of one line per token.
                steps = server.note_decode_step(session)
                if steps is not None:
                    print(
                        f"  [node] decoding layers={shard_start}-{shard_end} "
                        f"session={session[:8]} steps={steps}"
                        f"{self._request_log_suffix()}",
                        flush=True,
                    )
            else:
                print(
                    f"  [node] forward hop={hop_index} "
                    f"layers={shard_start}-{shard_end} "
                    f"session={session[:8]}{self._request_log_suffix()}",
                    flush=True,
                )

        start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
        start_layer = int(start_layer_header) if start_layer_header else None
        forward_kwargs: dict[str, object] = {}
        if cache_mode in ("prefill", "decode"):
            past_len_header = self.headers.get("X-Meshnet-Past-Len")
            forward_kwargs = {
                "session_id": session,
                "cache_mode": cache_mode,
                "past_len": int(past_len_header) if past_len_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,
                **forward_kwargs,
            )
        except KVCacheMiss as exc:
            self._send_json(409, {"error": "cache_miss", "detail": str(exc)})
            return
        except Exception as exc:
            print(
                f"  [node] forward failed layers={getattr(server.backend, 'shard_start', '?')}-"
                f"{getattr(server.backend, 'shard_end', '?')} session={session[:8]}: {exc}"
                f"{self._request_log_suffix()}",
                flush=True,
            )
            self._send_json(500, {"error": str(exc)})
            return

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

        route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
        phase_condition = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
        response_compression = compress_activation(
            result.body, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
        )
        response_body = response_compression.body
        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 response_compression.encoding:
            self.send_header("X-Meshnet-Encoding", response_compression.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]
        request_id = self._request_id()
        with server._stats_lock:
            server.total_requests += 1
            server.queue_depth += 1
        try:
            self._do_chat_completions(server, request_id)
        finally:
            self._track_request_end(server, request_id)
            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", request_id: str) -> 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 5120)
        temperature = float(body.get("temperature") or 1.0)
        top_p = float(body.get("top_p") or 1.0)

        self._track_request_begin(server, request_id, model_name)
        print(
            f"  [node] processing chat model={model_name!r} stream={stream} "
            f"max_tokens={max_tokens}{self._request_log_suffix()}",
            flush=True,
        )

        # 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:
            gen_started = time.monotonic()
            progress_line = [False]
            try:
                if stream:
                    token_count = 0

                    def _counting_stream():
                        nonlocal token_count
                        for token_text in backend.generate_text_streaming(
                            messages, max_tokens, temperature, top_p,
                        ):
                            if token_text:
                                token_count += 1
                                self._track_request_progress(
                                    server, request_id, tokens=token_count, routing_complete=True,
                                )
                            yield token_text

                    self._stream_openai_response(_counting_stream(), model_name)
                    elapsed = time.monotonic() - gen_started
                    tps = token_count / max(elapsed, 1e-6)
                    _write_progress_line(
                        progress_line,
                        f"  [node] chat complete (stream) tokens={token_count} "
                        f"elapsed_s={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}",
                        final=True,
                    )
                else:
                    text = backend.generate_text(messages, max_tokens, temperature, top_p)
                    completion_tokens = _backend_token_count(
                        backend, "count_text_tokens", text, fallback=len(text.split()) or 1,
                    )
                    print(
                        f"  [node] chat complete tokens={completion_tokens} "
                        f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}",
                        flush=True,
                    )
                    self._send_openai_response(text, model_name, False, messages, backend=backend)
            except Exception as exc:
                self._record_failed_request()
                print(
                    f"  [node] chat failed after {time.monotonic() - gen_started:.1f}s: {exc}"
                    f"{self._request_log_suffix()}",
                    flush=True,
                )
                self._send_json(500, {"error": f"generation failed: {exc}"})
            return

        # Distributed path: autoregressive generation across shards with a
        # sharded per-node KV cache. Step 0 prefills the full prompt through the
        # route (each node caches state for its own layer range, keyed by a
        # per-generation session id); steps 1+ send only the newest token's
        # hidden state. A 409 cache_miss from any hop (eviction/restart/route
        # change) falls back to a full re-prefill — the old stateless behavior.
        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

        session_id = str(uuid.uuid4())
        telemetry = GenerationTelemetry(session_id)
        with server._stats_lock:
            current = server._active_requests.get(request_id)
            if current is not None:
                current["telemetry"] = telemetry
        use_kv = bool(getattr(backend, "supports_kv_cache", False))
        # EOS detection by id must work on the stateless path too: the tail
        # returns token_id regardless of caching, and EOS usually decodes to
        # "" (skip_special_tokens), so the text comparison never fires.
        eos_ids: set[int] = set()
        try:
            eos_ids = set(backend.eos_token_ids())
        except Exception:
            eos_ids = set()

        stream_emit = None
        if stream:
            stream_emit = self._start_openai_stream(model_name)
            self._track_request_progress(server, request_id, tokens=0, routing_complete=True)

        _GENERATION_LOG_INTERVAL = 5.0
        gen_started = time.monotonic()
        last_gen_log = gen_started
        progress_line = [False]
        last_token_id: int | None = None
        failure_reason: str | None = None
        relay_clients: dict[str, _RelayHopClient] = {}
        direct_clients: dict[str, _DirectHopClient] = {}

        def _prefill_step() -> tuple[str, int | None]:
            """Full-sequence prefill: initial step and cache-miss recovery."""
            payload = (
                backend.encode_prompt(current_text, session_id=session_id)
                if use_kv
                else backend.encode_prompt(current_text)
            )
            return self._run_downstream_pipeline(
                payload, remaining_route, backend=backend,
                session=session_id, cache_mode="prefill" if use_kv else None,
                relay_clients=relay_clients,
                direct_clients=direct_clients if use_kv else None,
                telemetry=telemetry,
            )

        for step in range(max_tokens):
            try:
                if use_kv and step > 0 and last_token_id is not None:
                    try:
                        payload = backend.encode_next_token(last_token_id, session_id)
                        token_str, token_id = self._run_downstream_pipeline(
                            payload, remaining_route, backend=backend,
                            session=session_id, cache_mode="decode",
                            relay_clients=relay_clients,
                            direct_clients=direct_clients,
                            telemetry=telemetry,
                        )
                    except (KVCacheMiss, _PipelineCacheMiss) as miss:
                        # Evicted/restarted node or head lost its own session:
                        # re-prefill the whole sequence once and continue cached.
                        print(
                            f"  [node] kv cache miss at step {step} ({miss}); "
                            f"re-prefilling {len(current_text)} chars",
                            flush=True,
                        )
                        token_str, token_id = _prefill_step()
                else:
                    token_str, token_id = _prefill_step()
            except _PipelineCacheMiss as exc:
                print(f"  [node] unexpected cache miss on prefill session={session_id[:8]}: {exc}", flush=True)
                failure_reason = f"cache miss on prefill: {exc}"
                break
            except Exception as exc:
                print(f"  [node] distributed encode error session={session_id[:8]}: {exc}", flush=True)
                failure_reason = f"distributed encode error: {exc}"
                break
            # Stop on error responses or EOS.
            if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
                failure_reason = token_str
                break
            if token_id is not None and token_id in eos_ids:
                break
            if eos_token and token_str == eos_token:
                break
            if not token_str and token_id is None:
                break
            last_token_id = token_id
            # token_str can be empty for a skipped special token that is not
            # EOS — keep generating from its token_id without emitting text.
            if token_str:
                generated.append(token_str)
                if stream_emit is not None:
                    stream_emit(token_str)
                current_text = current_text + token_str
            self._track_request_progress(
                server,
                request_id,
                tokens=len(generated),
                routing_complete=True,
            )
            telemetry.note_tokens(len(generated))
            now = time.monotonic()
            if telemetry.report_due:
                summary = telemetry.snapshot(now=now)
                seams = summary["seams"]
                if seams:
                    latest = seams[-1]
                    print(
                        f"  [node] seam telemetry session={session_id[:8]} "
                        f"phase={latest['phase']} hop={latest['hop']} "
                        f"activations={latest['activations']} "
                        f"avg_ms={latest['avg_latency_ms']:.2f} "
                        f"wire_bytes={latest['wire_bytes']} "
                        f"response_bytes={latest['response_bytes']} "
                        f"tps={summary['tokens_per_sec']:.2f} "
                        f"activation_id={latest['last_activation_id']}",
                        flush=True,
                    )
                telemetry.mark_reported(now=now)
            if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
                elapsed = now - gen_started
                token_count = len(generated)
                tps = token_count / max(elapsed, 1e-6)
                _write_progress_line(
                    progress_line,
                    f"  [node] generating step={step + 1}/{max_tokens} session={session_id[:8]} "
                    f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
                )
                last_gen_log = now

        if use_kv:
            try:
                backend.release_session(session_id)
            except Exception:
                pass
        for relay_client in relay_clients.values():
            relay_client.close()
        for direct_client in direct_clients.values():
            direct_client.close()

        if generated:
            elapsed = time.monotonic() - gen_started
            token_count = len(generated)
            tps = token_count / max(elapsed, 1e-6)
            _write_progress_line(
                progress_line,
                f"  [node] generation complete session={session_id[:8]} tokens={token_count} "
                f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
                final=True,
            )
        telemetry.close()

        result_text = "".join(generated)
        # A failure before the first token is an upstream error, not an empty
        # completion — tell the client instead of returning a blank 200.
        if failure_reason and not generated:
            if stream_emit is not None:
                stream_emit(None, error=failure_reason)
                return
            self._send_json(502, {
                "error": {"message": failure_reason, "type": "upstream_error"},
            })
            return
        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.
        """
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        active_backend = backend or server.backend

        # 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})
                    hops = _clamp_downstream_hops(hops, active_backend)
                    print(
                        f"  [node] using injected downstream route: {_format_downstream_route(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).
        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_key = _own_endpoint_key(server)
            nodes_info = route_resp.get("nodes", [])
            hops: list[dict] = []
            passed_self = False
            for node_info in nodes_info:
                ep = node_info.get("endpoint", "")
                if not ep:
                    continue
                if _endpoint_key(ep) == own_key:
                    passed_self = True
                    continue
                if not passed_self:
                    continue
                hop = {
                    "endpoint": ep,
                    "start_layer": int(node_info.get("start_layer", 0)),
                }
                if node_info.get("relay_addr"):
                    hop["relay_addr"] = str(node_info["relay_addr"])
                hops.append(hop)
            hops = _clamp_downstream_hops(hops, active_backend)
            print(
                f"  [node] tracker downstream route: {_format_downstream_route(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,
        session: str | None = None,
        cache_mode: str | None = None,
        relay_clients: dict[str, _RelayHopClient] | None = None,
        direct_clients: dict[str, _DirectHopClient] | None = None,
        telemetry: GenerationTelemetry | None = None,
    ) -> tuple[str, int | None]:
        """Forward an activation through the downstream route.

        Returns (token_text, token_id) — token_id is None when a hop predates
        the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds
        409 cache_miss (evicted/restarted node) so the caller can re-prefill.
        """
        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 = _tensor_from_bfloat16_bytes(
                        payload.body, payload.shape, active_backend.torch,  # type: ignore[union-attr]
                    ).to(active_backend.device)
                    if hasattr(active_backend, "decode_tail_token"):
                        tail = active_backend.decode_tail_token(tensor)
                        return tail.text, tail.token_id
                    return active_backend.decode_tail(tensor), None
                except Exception as exc:
                    return f"decode error: {exc}", None
            return "no downstream route available for non-tail shard", None

        # Session is stable across all steps of one generation when the caller
        # provides it (KV-cache protocol); fresh per call otherwise (legacy).
        session = session or 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,
                )
            phase = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
            route_condition = "relay" if relay_addr else "lan"
            compression = compress_activation(
                current_body, _COMPRESSION_POLICIES.for_condition(route_condition, phase),
            )
            wire_body, wire_encoding = compression.body, compression.encoding
            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),
                "X-Meshnet-Activation-Id": uuid.uuid4().hex,
                "X-Meshnet-Compression-Route": route_condition,
            }
            if wire_encoding:
                headers["X-Meshnet-Encoding"] = wire_encoding
            if telemetry is not None:
                telemetry.record_compression(
                    phase=phase, hop=hop_index, node=node_url,
                    input_bytes=compression.input_bytes, output_bytes=compression.output_bytes,
                    elapsed_seconds=compression.elapsed_seconds,
                )
            if cache_mode:
                headers["X-Meshnet-Cache"] = cache_mode
                past_len = getattr(payload, "past_len", None)
                if cache_mode == "decode" and past_len is not None:
                    headers["X-Meshnet-Past-Len"] = str(past_len)
            if current_attn:
                headers["X-Meshnet-Attn-Mask"] = current_attn
            if current_pos:
                headers["X-Meshnet-Position-Ids"] = current_pos
            if relay_addr:
                connection_reused = bool(
                    relay_clients is not None and relay_addr in relay_clients
                )
                seam_started = time.monotonic()
                try:
                    if relay_clients is None:
                        status, resp_headers, resp_body = _relay_hop(
                            relay_addr, "/forward", wire_body, headers, timeout=120.0,
                        )
                    else:
                        relay_client = relay_clients.setdefault(
                            relay_addr, _RelayHopClient(relay_addr, timeout=120.0),
                        )
                        status, resp_headers, resp_body = relay_client.request(
                            "/forward", wire_body, headers,
                        )
                    if telemetry is not None:
                        telemetry.record_seam(
                            activation_id=headers["X-Meshnet-Activation-Id"],
                            phase=cache_mode or "prefill",
                            hop=hop_index,
                            node=node_url,
                            latency_seconds=time.monotonic() - seam_started,
                            wire_bytes=len(wire_body),
                            response_bytes=len(resp_body),
                            connection_reused=connection_reused,
                        )
                    if status == 409 and _is_cache_miss_body(resp_body):
                        raise _PipelineCacheMiss(node_url)
                    if status >= 400:
                        detail = _response_error_snippet(resp_body)
                        print(
                            f"  [node] relay hop {hop_index} session={session[:8]} returned "
                            f"{status} from {relay_addr}: {detail}",
                            flush=True,
                        )
                        return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
                except _PipelineCacheMiss:
                    raise
                except _RelayRequestUncertainError as exc:
                    # The activation may already have mutated the downstream
                    # KV cache.  Do not replay it on a direct connection.
                    print(
                        f"  [node] relay hop {hop_index} session={session[:8]} outcome is uncertain at "
                        f"{relay_addr}: {exc}",
                        flush=True,
                    )
                    return f"pipeline relay outcome uncertain at {node_url}: {exc}", None
                except Exception as exc:
                    print(
                        f"  [node] relay hop {hop_index} session={session[:8]} 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:
                connection_reused = bool(
                    direct_clients is not None and node_url in direct_clients
                )
                seam_started = time.monotonic()
                try:
                    if direct_clients is None:
                        direct_client = _DirectHopClient(node_url, timeout=120.0)
                        try:
                            status, resp_headers, resp_body = direct_client.request(
                                "/forward", wire_body, headers,
                            )
                        finally:
                            direct_client.close()
                    else:
                        direct_client = direct_clients.setdefault(
                            node_url, _DirectHopClient(node_url, timeout=120.0),
                        )
                        status, resp_headers, resp_body = direct_client.request(
                            "/forward", wire_body, headers,
                        )
                    if telemetry is not None:
                        telemetry.record_seam(
                            activation_id=headers["X-Meshnet-Activation-Id"],
                            phase=cache_mode or "prefill",
                            hop=hop_index,
                            node=node_url,
                            latency_seconds=time.monotonic() - seam_started,
                            wire_bytes=len(wire_body),
                            response_bytes=len(resp_body),
                            connection_reused=connection_reused,
                        )
                    if status == 409 and _is_cache_miss_body(resp_body):
                        raise _PipelineCacheMiss(node_url)
                    if status >= 400:
                        detail = _response_error_snippet(resp_body)
                        print(
                            f"  [node] pipeline hop {hop_index} session={session[:8]} failed at {node_url}: "
                            f"status {status}: {detail}", flush=True,
                        )
                        return f"pipeline error at {node_url}: status {status}: {detail}", None
                except _PipelineCacheMiss:
                    raise
                except _DirectRequestUncertainError as exc:
                    print(
                        f"  [node] pipeline hop {hop_index} session={session[:8]} outcome is uncertain "
                        f"at {node_url}: {exc}",
                        flush=True,
                    )
                    return f"pipeline direct outcome uncertain at {node_url}: {exc}", None
                except Exception as exc:
                    print(
                        f"  [node] pipeline hop {hop_index} session={session[:8]} "
                        f"failed at {node_url}: {exc}", flush=True,
                    )
                    return f"pipeline error at {node_url}: {exc}", None
            content_type = resp_headers.get("content-type", "")
            if "application/json" in content_type:
                try:
                    data = json.loads(resp_body)
                    text = str(data.get("text", ""))
                    token_id = data.get("token_id")
                    if server.debug:
                        print(f"  [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
                    return text, int(token_id) if token_id is not None else None
                except json.JSONDecodeError:
                    return resp_body.decode("utf-8", errors="replace"), None
            # 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)
            try:
                decompression = decompress_activation(
                    resp_body, resp_headers.get("x-meshnet-encoding")
                )
                current_body = decompression.body
                if telemetry is not None:
                    telemetry.record_compression(
                        phase=phase, hop=hop_index, node=node_url,
                        input_bytes=decompression.input_bytes, output_bytes=decompression.output_bytes,
                        elapsed_seconds=decompression.elapsed_seconds, decompression=True,
                    )
            except ValueError as exc:
                print(
                    f"  [node] pipeline hop {hop_index} session={session[:8]} "
                    f"bad response encoding: {exc}", flush=True,
                )
                return f"pipeline error at {node_url}: {exc}", None
            current_attn = resp_headers.get("x-meshnet-attn-mask")
            current_pos = resp_headers.get("x-meshnet-position-ids")
        return "", None

    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.send_header("Transfer-Encoding", "chunked")
        self.end_headers()

        def _emit(data: str) -> bool:
            try:
                body = f"data: {data}\n\n".encode()
                self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
                self.wfile.flush()
                return True
            except (BrokenPipeError, ConnectionResetError):
                return False

        _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, *, error: str | None = None) -> None:
            if error is not None:
                # OpenAI-style mid-stream error frame; clients surface it
                # instead of showing an empty completion.
                _emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
                try:
                    self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
                    self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    pass
                return
            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"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\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.send_header("Transfer-Encoding", "chunked")
        self.end_headers()

        def _emit(data: str) -> bool:
            try:
                body = f"data: {data}\n\n".encode()
                self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
                self.wfile.flush()
                return True
            except (BrokenPipeError, ConnectionResetError):
                return False

        _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"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
            self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            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,
        force_cpu: bool = False,
        recipe_params: Mapping[str, Any] | None = None,
    ) -> 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,
            force_cpu=force_cpu,
            recipe_params=recipe_params,
        )
        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 current_requests(self) -> list[dict[str, Any]]:
        if self._server is None:
            return []
        return self._server.snapshot_current_requests()

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

    def backend_for(self, model_id: str) -> TorchModelShard | None:
        """The loaded backend serving `model_id` — full repo id or short name."""
        backend = self._backends.get(model_id)
        if backend is not None:
            return backend
        short = model_id.split("/")[-1].lower()
        for key, candidate in self._backends.items():
            if key.split("/")[-1].lower() == short:
                return candidate
        return None

    def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
        """Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more)."""
        drop_all_directive = next(
            (directive for directive in reversed(directives) if directive.get("action") == "DROP_ALL_SHARDS"),
            None,
        )
        if drop_all_directive is not None:
            self._backends.clear()
            self._backend = None
            self._tracker_mode = False
            if self._server is not None:
                self._server.backends = {}
                self._server.backend = None
                self._server.tracker_mode = False
            return {"action": "DROP_ALL_SHARDS"}
        drop_directive = next(
            (directive for directive in reversed(directives) if directive.get("action") == "DROP_SHARD"),
            None,
        )
        if drop_directive is not None:
            model_id = str(drop_directive.get("model") or "")
            removed = self._backends.pop(model_id, None)
            if removed is not None:
                if self._backends:
                    self._backend = next(iter(self._backends.values()))
                    self._tracker_mode = self._backend.shard_start == 0
                else:
                    self._backend = None
                    self._tracker_mode = False
                if self._server is not None:
                    self._server.backends = dict(self._backends)
                    self._server.backend = self._backend
                    self._server.tracker_mode = self._tracker_mode
            # A DROP_SHARD may be paired with a subsequent LOAD_SHARD in the
            # same heartbeat; continue so the replacement is applied atomically.
            if not any(
                directive.get("action") in {"ADD_SHARD", "LOAD_SHARD"}
                for directive in directives
            ):
                return {"action": "DROP_SHARD", "model": model_id} if removed is not None else None
        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 if self._backend is not None else "auto")
        )
        model_id = str(directive.get("model") or (self._backend.model_id if self._backend is not None else ""))
        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:
            if replacing:
                self._backends.clear()
            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 set_advertised_endpoint(self, endpoint: str) -> None:
        """Set the LAN-facing endpoint used for route self-detection."""
        if self._server is not None:
            self._server.advertised_endpoint = endpoint

    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,
    force_cpu: bool = False,
    recipe_params: Mapping[str, Any] | 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,
            force_cpu=force_cpu,
            recipe_params=recipe_params,
        )
    except MissingModelDependencyError:
        raise
    except InsufficientVRAMError as exc:
        print(f"ERROR: {exc}", file=sys.stderr, flush=True)
        raise
