"""Tracker HTTP server — node registry and route selection for the inference mesh.

HTTP API contract:
- POST /v1/nodes/register
  Request JSON: {
    "endpoint": "http://host:port", "shard_start": int, "shard_end": int,
    "model": str optional, "shard_checksum": str optional,
    "hardware_profile": object, "wallet_address": str optional,
    "score": number optional
  }
  Response 200: {"node_id": str}
  Response 400: {"error": str}
- POST /v1/nodes/{node_id}/heartbeat
  Response 200: {}
  Response 404: {"error": "node not found"}
- GET /v1/route?model=<preset>
  Response 200: {"route": ["http://node-a", "http://node-b"]}
  Response 400/404/503: {"error": str}
- GET /v1/routes?model=<preset>&redundancy=<n>
  Response 200: {"routes": [{"route": [...], "nodes": [...]}]}
  Response 400/404/503: {"error": str}
"""

import http.server
import hashlib
import json
import socketserver
import threading
import time
import urllib.parse
import urllib.request
import uuid
from typing import Any

from .gossip import NodeGossip
from .raft import RaftNode


DEFAULT_MODEL_PRESETS: dict[str, dict] = {
    "stub-model": {
        "layers_start": 0,
        "layers_end": 31,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
    },
    "openai-community/gpt2": {
        "layers_start": 0,
        "layers_end": 11,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
    },
}

DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024
DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024
DEFAULT_QUANTIZATIONS = ["bfloat16"]
DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0


class _NodeEntry:
    __slots__ = (
        "node_id", "endpoint", "shard_start", "shard_end",
        "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
        "score", "vram_bytes", "ram_bytes", "quantizations",
        "benchmark_tokens_per_sec", "quantization", "managed_assignment",
        "pending_directives", "last_heartbeat", "tracker_mode",
        "relay_addr", "cert_fingerprint", "peer_id",
    )

    def __init__(
        self,
        node_id: str,
        endpoint: str,
        shard_start: int | None,
        shard_end: int | None,
        model: str | None,
        shard_checksum: str | None,
        hardware_profile: dict,
        wallet_address: str | None,
        score: float,
        vram_bytes: int = DEFAULT_VRAM_BYTES,
        ram_bytes: int = DEFAULT_RAM_BYTES,
        quantizations: list[str] | None = None,
        benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC,
        quantization: str | None = None,
        managed_assignment: bool = False,
        tracker_mode: bool = False,
        hf_repo: str | None = None,
        num_layers: int | None = None,
        relay_addr: str | None = None,
        cert_fingerprint: str | None = None,
        peer_id: str | None = None,
    ) -> None:
        self.node_id = node_id
        self.endpoint = endpoint
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.model = model
        self.shard_checksum = shard_checksum
        self.hardware_profile = hardware_profile
        self.wallet_address = wallet_address
        self.score = score
        self.vram_bytes = vram_bytes
        self.ram_bytes = ram_bytes
        self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS)
        self.benchmark_tokens_per_sec = benchmark_tokens_per_sec
        self.quantization = quantization
        self.managed_assignment = managed_assignment
        self.tracker_mode = tracker_mode
        self.hf_repo = hf_repo
        self.num_layers = num_layers
        self.relay_addr = relay_addr
        self.cert_fingerprint = cert_fingerprint
        self.peer_id = peer_id
        self.pending_directives: list[dict] = []
        self.last_heartbeat: float = time.monotonic()


def _select_route(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> tuple[list[_NodeEntry], str]:
    """Greedy interval-cover. Returns (ordered route, error_message)."""
    candidates = sorted(
        [node for node in nodes if node.shard_start is not None and node.shard_end is not None],
        key=lambda n: (n.shard_start, -n.shard_end),  # type: ignore[operator]
    )
    route: list[_NodeEntry] = []
    covered_up_to = required_start - 1

    while covered_up_to < required_end:
        best: _NodeEntry | None = None
        for node in candidates:
            if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
                if best is None or node.shard_end > best.shard_end:
                    best = node
        if best is None:
            missing = covered_up_to + 1
            return [], f"no route available: no registered node covers layer {missing}"
        route.append(best)
        covered_up_to = best.shard_end
        candidates = [n for n in candidates if n is not best]

    return route, ""


def _coverage_percentage(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> float:
    required_layers = required_end - required_start + 1
    if required_layers <= 0:
        return 0.0

    intervals = sorted(
        (
            (max(required_start, node.shard_start), min(required_end, node.shard_end))
            for node in nodes
            if node.shard_start is not None
            and node.shard_end is not None
            if node.shard_end >= required_start and node.shard_start <= required_end
        ),
        key=lambda interval: interval[0],
    )
    covered = 0
    covered_until = required_start - 1
    for start, end in intervals:
        if end <= covered_until:
            continue
        next_start = max(start, covered_until + 1)
        if next_start > end:
            continue
        covered += end - next_start + 1
        covered_until = end
    return round((covered / required_layers) * 100, 2)


def _preset_layer_bounds(preset: dict) -> tuple[int, int]:
    start = int(preset.get("layers_start", 0))
    if "layers_end" in preset:
        return start, int(preset["layers_end"])
    return start, start + int(preset["total_layers"]) - 1


def _preset_bytes_per_layer(preset: dict) -> dict[str, int]:
    raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {}))
    if isinstance(raw, dict) and raw:
        return {str(quant): int(value) for quant, value in raw.items()}
    return {"bfloat16": 30 * 1024 * 1024}


def _node_quantization(node: _NodeEntry, preset: dict) -> str:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    if node.quantization in bytes_per_layer:
        return node.quantization
    for quantization in node.quantizations:
        if quantization in bytes_per_layer:
            return quantization
    return next(iter(bytes_per_layer))


def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    quantization = _node_quantization(node, preset)
    layer_bytes = bytes_per_layer[quantization]
    if layer_bytes <= 0:
        return 0
    return int((node.vram_bytes * 0.8) // layer_bytes)


def _coverage_map(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> list[dict]:
    layer_counts = []
    for layer in range(required_start, required_end + 1):
        count = 0
        for node in nodes:
            if node.shard_start is None or node.shard_end is None:
                continue
            if node.shard_start <= layer <= node.shard_end:
                count += 1
        layer_counts.append((layer, count))

    coverage: list[dict] = []
    for layer, count in layer_counts:
        if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1:
            coverage[-1]["end_layer"] = layer
        else:
            coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count})
    return coverage


def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
    return [
        (segment["start_layer"], segment["end_layer"])
        for segment in coverage
        if segment["node_count"] == 0
    ]


def _relay_http_request(
    relay_addr: str,
    path: str,
    body: bytes,
    headers: dict[str, str],
    timeout: float = 310.0,
) -> dict | None:
    """Send an HTTP-shaped request through a relay RPC WebSocket."""
    try:
        import websockets.sync.client as wsc  # type: ignore[import]

        request_id = str(uuid.uuid4())
        with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws:
            ws.send(json.dumps({
                "request_id": request_id,
                "method": "POST",
                "path": path,
                "headers": headers,
                "body": body.decode(errors="replace"),
            }))
            raw = ws.recv(timeout=timeout)
        response = json.loads(raw)
        if response.get("request_id") not in {None, request_id}:
            return None
        return response
    except Exception:
        return None


def _nodes_and_bounds_for_model(
    server: "_TrackerHTTPServer",
    model: str,
) -> tuple[list[_NodeEntry], int, int] | None:
    preset = server.model_presets.get(model)
    if preset is not None:
        required_start, required_end = _preset_layer_bounds(preset)
        return [node for node in server.registry.values() if node.model == model], required_start, required_end

    nodes = [
        node for node in server.registry.values()
        if (node.hf_repo == model or node.model == model)
        and node.shard_start is not None
        and node.shard_end is not None
        and node.num_layers is not None
    ]
    if not nodes:
        return None
    return nodes, 0, max(node.num_layers for node in nodes) - 1


def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict:
    return {
        "action": "LOAD_SHARD",
        "model": model,
        "start_layer": start,
        "end_layer": end,
        "shard_start": start,
        "shard_end": end,
        "quantization": quantization,
    }


def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict:
    return {
        "action": "DROP_SHARD",
        "model": model,
        "start_layer": start,
        "end_layer": end,
        "shard_start": start,
        "shard_end": end,
        "quantization": quantization,
    }


def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
    now = time.monotonic()
    expired_ids = [
        node_id for node_id, entry in server.registry.items()
        if (now - entry.last_heartbeat) > server.heartbeat_timeout
    ]
    for node_id in expired_ids:
        entry = server.registry.pop(node_id)
        print(
            f"[tracker] node expired: {node_id}  {entry.endpoint}  "
            f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
            flush=True,
        )
    if expired_ids:
        _rebalance_all_locked(server)
    return expired_ids


def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
    preset = server.model_presets.get(model)
    if preset is None:
        return
    required_start, required_end = _preset_layer_bounds(preset)
    total_layers = required_end - required_start + 1
    model_nodes = [node for node in server.registry.values() if node.model == model]
    managed_nodes = [node for node in model_nodes if node.managed_assignment]
    if not managed_nodes:
        return

    previous_ranges = {
        node.node_id: (node.shard_start, node.shard_end, node.quantization)
        for node in managed_nodes
    }
    for node in managed_nodes:
        node.shard_start = None
        node.shard_end = None

    managed_nodes.sort(
        key=lambda node: (
            -node.benchmark_tokens_per_sec,
            -_node_layer_capacity(node, preset),
            node.node_id,
        )
    )
    base_nodes = [node for node in model_nodes if not node.managed_assignment]
    coverage = _coverage_map(base_nodes, required_start, required_end)
    gaps = _coverage_gaps(coverage)
    if not gaps:
        gaps = [(required_start, required_end)]

    eligible_nodes = [
        node for node in managed_nodes
        if _node_layer_capacity(node, preset) > 0
    ]
    node_index = 0
    for gap_start, gap_end in gaps:
        cursor = gap_start
        while cursor <= gap_end and node_index < len(eligible_nodes):
            node = eligible_nodes[node_index]
            remaining_layers = gap_end - cursor + 1
            remaining_nodes_after = len(eligible_nodes) - node_index - 1
            capacity = min(
                _node_layer_capacity(node, preset),
                total_layers,
                max(1, remaining_layers - remaining_nodes_after),
            )
            if capacity <= 0:
                node_index += 1
                continue
            quantization = _node_quantization(node, preset)
            node.quantization = quantization
            node.shard_start = cursor
            node.shard_end = min(gap_end, cursor + capacity - 1)
            cursor = node.shard_end + 1
            node_index += 1

    for node in managed_nodes:
        previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
        current_range = (node.shard_start, node.shard_end, node.quantization)
        if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
            continue
        if previous_start is not None and previous_end is not None:
            node.pending_directives.append(
                _drop_directive(
                    node,
                    model,
                    previous_start,
                    previous_end,
                    previous_quantization or _node_quantization(node, preset),
                )
            )
        node.pending_directives.append(
            _load_directive(
                node,
                model,
                node.shard_start,
                node.shard_end,
                node.quantization or _node_quantization(node, preset),
            )
        )


def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None:
    for model in list(server.model_presets):
        _rebalance_model_locked(server, model)


def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
    if contracts is None or not wallet_address:
        return None
    if contracts.registry.get_wallet(wallet_address).banned:
        return "wallet is banned"
    return None


def _node_id_for_registration(
    endpoint: str,
    model: str,
    wallet_address: str | None,
    shard_start: int | None,
    shard_end: int | None,
    hf_repo: str | None,
) -> str:
    wallet_prefix = wallet_address[:8] if wallet_address else "anon"
    stable_key = "|".join([
        wallet_address or "",
        endpoint.rstrip("/"),
        model,
        hf_repo or "",
        "" if shard_start is None else str(shard_start),
        "" if shard_end is None else str(shard_end),
    ])
    digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12]
    return f"{wallet_prefix}-{digest}"


class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True

    def __init__(
        self,
        addr: tuple,
        handler,
        registry: dict,
        lock: threading.Lock,
        heartbeat_timeout: float,
        model_presets: dict,
        contracts: Any | None,
        minimum_stake: int,
        relay_url: str | None = None,
        raft: "RaftNode | None" = None,
        gossip: "NodeGossip | None" = None,
    ) -> None:
        super().__init__(addr, handler)
        self.registry = registry
        self.lock = lock
        self.heartbeat_timeout = heartbeat_timeout
        self.model_presets = model_presets
        self.contracts = contracts
        self.minimum_stake = minimum_stake
        self.relay_url = relay_url.rstrip("/") if relay_url else None
        self.raft = raft
        self.gossip = gossip


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

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

    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:
            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 _purge_expired_nodes(self) -> None:
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        _purge_expired_nodes_locked(server)

    def do_POST(self):
        if self.path == "/v1/chat/completions":
            self._handle_proxy_chat()
            return
        if self.path == "/v1/nodes/register":
            self._handle_register()
            return
        if self.path == "/v1/raft/vote":
            self._handle_raft_vote()
            return
        if self.path == "/v1/raft/append":
            self._handle_raft_append()
            return
        if self.path == "/v1/gossip":
            self._handle_gossip()
            return
        parts = self.path.split("/")
        # /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
        if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
            self._handle_heartbeat(parts[3])
            return
        self.send_response(404)
        self.end_headers()

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/v1/route":
            self._handle_route(parsed)
        elif parsed.path == "/v1/routes":
            self._handle_routes(parsed)
        elif parsed.path == "/v1/nodes/assign":
            self._handle_assign(parsed)
        elif parsed.path == "/v1/network/assign":
            self._handle_network_assign(parsed)
        elif parsed.path == "/v1/network/map":
            self._handle_network_map()
        elif parsed.path == "/v1/models":
            self._handle_models()
        elif parsed.path.startswith("/v1/coverage/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/"))
            self._handle_coverage(model)
        elif parsed.path.startswith("/v1/tracker-nodes/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/"))
            self._handle_tracker_nodes(model)
        elif parsed.path == "/v1/raft/status":
            self._handle_raft_status()
        elif parsed.path == "/v1/health":
            self._send_json(200, {"status": "ok"})
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_models(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        with server.lock:
            self._purge_expired_nodes()
            alive = list(server.registry.values())
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        data = []
        seen_ids: set[str] = set()
        for name, preset in server.model_presets.items():
            model_nodes = [node for node in alive if node.model == name]
            if not model_nodes:
                continue
            required_start, required_end = _preset_layer_bounds(preset)
            coverage = _coverage_percentage(
                model_nodes,
                required_start,
                required_end,
            )
            aliases = [name]
            hf_repo = preset.get("hf_repo")
            if hf_repo and hf_repo not in aliases:
                aliases.append(hf_repo)
            data.append({
                "id": name,
                "object": "model",
                "created": created,
                "owned_by": "meshnet",
                "name": name,
                "hf_repo": hf_repo,
                "aliases": aliases,
                "shard_coverage_percentage": coverage,
            })
            seen_ids.add(name)

        hf_model_ids = sorted({
            node.hf_repo or node.model
            for node in alive
            if node.model is not None
            and node.model not in server.model_presets
            and node.shard_start is not None
            and node.shard_end is not None
            and node.num_layers is not None
        })
        for model_id in hf_model_ids:
            if model_id is None or model_id in seen_ids:
                continue
            model_nodes = [
                node for node in alive
                if node.shard_start is not None
                and node.shard_end is not None
                and node.num_layers is not None
                and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id))
            ]
            if not model_nodes:
                continue
            short_names = sorted({node.model for node in model_nodes if node.model})
            aliases = [model_id, *[name for name in short_names if name != model_id]]
            required_start = 0
            required_end = max(node.num_layers for node in model_nodes) - 1
            data.append({
                "id": model_id,
                "object": "model",
                "created": created,
                "owned_by": "meshnet",
                "name": short_names[0] if short_names else model_id,
                "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None,
                "aliases": aliases,
                "shard_coverage_percentage": _coverage_percentage(
                    model_nodes,
                    required_start,
                    required_end,
                ),
            })
            seen_ids.add(model_id)
        self._send_json(200, {"object": "list", "data": data})

    def _handle_coverage(self, model: str):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        with server.lock:
            self._purge_expired_nodes()
            resolved = _nodes_and_bounds_for_model(server, model)
        if resolved is None:
            self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
            return
        alive, required_start, required_end = resolved
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        coverage = _coverage_map(alive, required_start, required_end)
        self._send_json(200, {"model": model, "coverage": coverage})

    def _handle_tracker_nodes(self, model: str):
        """Return nodes registered with tracker_mode=True whose shard starts at layer 0."""
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return
        required_start, _ = _preset_layer_bounds(preset)
        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        tracker_nodes = [
            node for node in alive
            if node.shard_start is not None
            and node.shard_start == required_start
            and node.tracker_mode
        ]
        self._send_json(200, {
            "model": model,
            "tracker_nodes": [
                {
                    "node_id": node.node_id,
                    "endpoint": node.endpoint,
                    "relay_addr": node.relay_addr,
                    "peer_id": node.peer_id,
                    "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
                }
                for node in tracker_nodes
            ],
        })

    def _handle_network_map(self) -> None:
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        with server.lock:
            self._purge_expired_nodes()
            nodes = list(server.registry.values())
        self._send_json(200, {
            "relay_url": server.relay_url,
            "nodes": [
                {
                    "node_id": node.node_id,
                    "endpoint": node.endpoint,
                    "relay_addr": node.relay_addr,
                    "peer_id": node.peer_id,
                    "model": node.model,
                    "hf_repo": node.hf_repo,
                    "shard_start": node.shard_start,
                    "shard_end": node.shard_end,
                    "tracker_mode": node.tracker_mode,
                    "last_heartbeat": node.last_heartbeat,
                }
                for node in nodes
            ],
        })

    # ---------------------------------------------------------------- OpenAI proxy

    def _handle_proxy_chat(self) -> None:
        """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node.

        Picks a live tracker-mode node for the requested model using round-robin,
        then forwards the request verbatim and relays the response (including
        streaming SSE chunks) back to the caller.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]

        length = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(length) if length else b"{}"

        try:
            body = json.loads(raw_body)
        except json.JSONDecodeError:
            self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}})
            return

        model: str = body.get("model", "")
        is_stream: bool = bool(body.get("stream", False))

        # Find a live tracker-mode node for this model
        with server.lock:
            self._purge_expired_nodes()
            candidates = [
                n for n in server.registry.values()
                if n.tracker_mode and (n.model == model or n.hf_repo == model)
            ]

        if not candidates:
            # Fall back: any node serving shard_start=0 for this model
            with server.lock:
                candidates = [
                    n for n in server.registry.values()
                    if n.shard_start == 0 and (n.model == model or n.hf_repo == model)
                ]

        if not candidates:
            self._send_json(503, {"error": {
                "message": f"no nodes available for model {model!r}",
                "type": "service_unavailable",
                "code": "model_not_available",
            }})
            return

        # Simple round-robin via list length modulo (stateless, good enough)
        node = candidates[int(time.time() * 1000) % len(candidates)]
        target_url = f"{node.endpoint}/v1/chat/completions"
        request_id = str(body.get("id") or f"req-{time.time_ns():x}")

        # Pre-resolve the downstream route so the first-shard node skips its own
        # tracker query.  We already hold the full registry picture — no need for
        # a second round-trip.
        route_model = node.hf_repo or node.model or model
        with server.lock:
            if route_model in server.model_presets:
                preset = server.model_presets[route_model]
                rs, re = _preset_layer_bounds(preset)
                all_nodes: list = [n for n in server.registry.values() if n.model == route_model]
            else:
                all_nodes = [
                    n for n in server.registry.values()
                    if (n.hf_repo == route_model or n.model == route_model)
                    and n.shard_start is not None and n.num_layers is not None
                ]
                rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
        route_nodes, _ = _select_route(all_nodes, rs, re)
        # Strip the first-shard node we're about to proxy to — it's already handling the request.
        downstream_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint])
        route_debug = " -> ".join(
            f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
            for n in route_nodes
        )
        print(
            f"[tracker] proxy route {request_id}: model={model!r} "
            f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} "
            f"route={route_debug or '<empty>'}",
            flush=True,
        )

        req = urllib.request.Request(
            target_url,
            data=raw_body,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": downstream_urls,
            },
            method="POST",
        )
        # Copy Authorization header from client if present
        auth = self.headers.get("Authorization")
        if auth:
            req.add_header("Authorization", auth)

        try:
            upstream = urllib.request.urlopen(req, timeout=300.0)
            print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True)
        except urllib.error.HTTPError as exc:
            # Relay error status + body from node
            err_body = exc.read()
            self.send_response(exc.code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(err_body)))
            self.end_headers()
            try:
                self.wfile.write(err_body)
            except BrokenPipeError:
                pass
            return
        except Exception as exc:
            if node.relay_addr:
                print(
                    f"[tracker] direct proxy failed {request_id}: {target_url}: {exc}; "
                    f"trying relay {node.relay_addr}",
                    flush=True,
                )
                relayed = _relay_http_request(
                    node.relay_addr,
                    path="/v1/chat/completions",
                    body=raw_body,
                    headers={
                        "Content-Type": "application/json",
                        "X-Meshnet-Route": downstream_urls,
                        **({"Authorization": auth} if auth else {}),
                    },
                )
                if relayed is not None:
                    self._send_relayed_response(relayed)
                    return
            print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True)
            self._send_json(503, {"error": {
                "message": f"upstream node unreachable: {exc}",
                "type": "service_unavailable",
                "code": "upstream_error",
            }})
            return

        with upstream:
            content_type = upstream.headers.get("Content-Type", "application/json")
            if is_stream or "text/event-stream" in content_type:
                # Relay SSE stream chunk-by-chunk
                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()
                try:
                    while True:
                        line = upstream.readline()
                        if not line:
                            break
                        self.wfile.write(line)
                        self.wfile.flush()
                except BrokenPipeError:
                    pass
            else:
                # Non-streaming: buffer and relay
                resp_body = upstream.read()
                print(
                    f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}",
                    flush=True,
                )
                self.send_response(200)
                self.send_header("Content-Type", content_type)
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                try:
                    self.wfile.write(resp_body)
                except BrokenPipeError:
                    pass

    def _send_relayed_response(self, response: dict) -> None:
        status = int(response.get("status", 503))
        headers = response.get("headers") if isinstance(response.get("headers"), dict) else {}
        body_text = response.get("body") or ""
        body = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
        self.send_response(status)
        self.send_header("Content-Type", headers.get("Content-Type", "application/json"))
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        try:
            self.wfile.write(body)
        except BrokenPipeError:
            pass

    def _handle_register(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return

        # --- Raft cluster mode: forward to leader or propose via Raft ---
        if server.raft is not None:
            if not server.raft.is_leader:
                leader = server.raft.leader()
                if leader is None:
                    self._send_json(503, {"error": "no leader elected — retry in a moment"})
                    return
                # Proxy to leader
                try:
                    data = json.dumps(body).encode()
                    req = urllib.request.Request(
                        f"{leader}/v1/nodes/register",
                        data=data,
                        headers={"Content-Type": "application/json"},
                        method="POST",
                    )
                    with urllib.request.urlopen(req, timeout=5.0) as r:
                        resp_body = r.read()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(resp_body)))
                    self.end_headers()
                    self.wfile.write(resp_body)
                except Exception as exc:
                    self._send_json(503, {"error": f"leader proxy failed: {exc}"})
                return
            # Leader path: fall through to normal registration, then replicate via Raft.
            # We let the registration run first (to generate node_id), then asynchronously
            # propose to the Raft log so followers can replicate the entry.
            # Raft proposal happens after the response is sent (fire-and-forget replication).

        endpoint = body.get("endpoint")
        if not isinstance(endpoint, str) or not endpoint:
            self._send_json(400, {"error": "endpoint is required"})
            return
        parsed_endpoint = urllib.parse.urlparse(endpoint)
        if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc:
            self._send_json(400, {"error": "endpoint must be an http(s) URL"})
            return

        shard_start: int | None
        shard_end: int | None
        explicit_shard = "shard_start" in body or "shard_end" in body
        if explicit_shard:
            try:
                shard_start = int(body["shard_start"])
                shard_end = int(body["shard_end"])
            except (KeyError, TypeError, ValueError):
                self._send_json(400, {"error": "shard_start and shard_end must be numeric"})
                return
            if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
                self._send_json(400, {"error": "shard range must be non-negative and ordered"})
                return
        else:
            shard_start = None
            shard_end = None
        try:
            score = float(body.get("score", 1.0))
        except (TypeError, ValueError):
            self._send_json(400, {"error": "score must be numeric"})
            return

        hardware_profile = body.get("hardware_profile", {})
        if not isinstance(hardware_profile, dict):
            self._send_json(400, {"error": "hardware_profile must be an object"})
            return
        model = body.get("model")
        if model is None:
            model = "stub-model"
        if not isinstance(model, str):
            self._send_json(400, {"error": "model must be a string"})
            return
        shard_checksum = body.get("shard_checksum")
        if shard_checksum is not None and not isinstance(shard_checksum, str):
            self._send_json(400, {"error": "shard_checksum must be a string"})
            return
        try:
            vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES))
            ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES))
            benchmark_tokens_per_sec = float(
                body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC)
            )
        except (TypeError, ValueError):
            self._send_json(400, {"error": "vram_bytes, ram_bytes, and benchmark_tokens_per_sec must be numeric"})
            return
        if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0:
            self._send_json(400, {"error": "capability values must be positive"})
            return
        quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS)
        if not (
            isinstance(quantizations_body, list)
            and quantizations_body
            and all(isinstance(item, str) and item for item in quantizations_body)
        ):
            self._send_json(400, {"error": "quantizations must be a non-empty string array"})
            return
        quantizations = list(quantizations_body)
        quantization = body.get("quantization")
        if quantization is not None and not isinstance(quantization, str):
            self._send_json(400, {"error": "quantization must be a string"})
            return
        wallet_address = body.get("wallet_address")
        if wallet_address is not None and not isinstance(wallet_address, str):
            self._send_json(400, {"error": "wallet_address must be a string"})
            return
        ban_error = _registration_ban_error(server.contracts, wallet_address)
        if ban_error:
            self._send_json(403, {"error": ban_error})
            return

        tracker_mode = bool(body.get("tracker_mode", False))
        hf_repo = body.get("hf_repo")
        if hf_repo is not None and not isinstance(hf_repo, str):
            self._send_json(400, {"error": "hf_repo must be a string"})
            return
        num_layers_body = body.get("num_layers")
        num_layers: int | None = None
        if num_layers_body is not None:
            try:
                num_layers = int(num_layers_body)
            except (TypeError, ValueError):
                self._send_json(400, {"error": "num_layers must be an integer"})
                return
        relay_addr = body.get("relay_addr") or None
        cert_fingerprint = body.get("cert_fingerprint") or None
        peer_id = body.get("peer_id") or None

        node_id = _node_id_for_registration(
            endpoint,
            model,
            wallet_address,
            shard_start,
            shard_end,
            hf_repo,
        )
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            model=model,
            shard_checksum=shard_checksum,
            hardware_profile=hardware_profile,
            wallet_address=wallet_address,
            score=score,
            vram_bytes=vram_bytes,
            ram_bytes=ram_bytes,
            quantizations=quantizations,
            benchmark_tokens_per_sec=benchmark_tokens_per_sec,
            quantization=quantization,
            managed_assignment=not explicit_shard,
            tracker_mode=tracker_mode,
            hf_repo=hf_repo,
            num_layers=num_layers,
            relay_addr=relay_addr,
            cert_fingerprint=cert_fingerprint,
            peer_id=peer_id,
        )
        with server.lock:
            self._purge_expired_nodes()
            # Dedup: if this endpoint is already registered, remove the old entry first.
            stale_ids = [
                eid for eid, e in server.registry.items()
                if e.endpoint == entry.endpoint.rstrip("/")
            ]
            for eid in stale_ids:
                old = server.registry.pop(eid)
                print(
                    f"[tracker] node re-registered: replaced {eid} with {node_id}"
                    f"  {old.endpoint}",
                    flush=True,
                )
            server.registry[node_id] = entry
            if entry.managed_assignment:
                _rebalance_model_locked(server, model)
            assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None
            if assignment_directive is not None:
                entry.pending_directives.clear()

        shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
        repo_info = f" [{hf_repo}]" if hf_repo else ""
        print(
            f"[tracker] node registered: {node_id}  {endpoint}  {model}{repo_info}  {shard_info}",
            flush=True,
        )

        payload = {"node_id": node_id}
        if assignment_directive is not None:
            payload["directive"] = assignment_directive
        self._send_json(200, payload)

        # Raft replication: leader proposes this registration to followers so their
        # registries stay in sync.  Fire-and-forget (async) — the client already
        # got its node_id; replication happens in the background.
        if server.raft is not None and server.raft.is_leader:
            raft_payload = dict(body)
            raft_payload["node_id"] = node_id  # include the generated ID
            threading.Thread(
                target=server.raft.propose,
                args=("register", raft_payload),
                daemon=True,
            ).start()

    def _handle_heartbeat(self, node_id: str):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body: dict = {}
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length > 0:
            try:
                body = json.loads(self.rfile.read(content_length))
            except Exception:
                pass
        with server.lock:
            self._purge_expired_nodes()
            entry = server.registry.get(node_id)
            if entry is None:
                self._send_json(404, {"error": "node not found"})
                return
            entry.last_heartbeat = time.monotonic()
            if body.get("relay_addr"):
                entry.relay_addr = body["relay_addr"]
            if body.get("cert_fingerprint"):
                entry.cert_fingerprint = body["cert_fingerprint"]
            if body.get("peer_id"):
                entry.peer_id = body["peer_id"]
            _rebalance_model_locked(server, entry.model or "stub-model")
            directives = list(entry.pending_directives)
            entry.pending_directives.clear()
        if server.gossip is not None:
            server.gossip.record(node_id)
        if directives:
            self._send_json(200, {"directives": directives})
        else:
            self._send_json(200, {})

    # ---------------------------------------------------------------- Raft handlers

    def _handle_raft_vote(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_request_vote(body)
        self._send_json(200, result)

    def _handle_raft_append(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_append_entries(body)
        self._send_json(200, result)

    def _handle_raft_status(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        if server.raft is None:
            self._send_json(200, {"role": "standalone", "term": 0, "leader": None})
            return
        self._send_json(200, server.raft.status())

    def _handle_gossip(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.gossip is not None and isinstance(body, dict):
            server.gossip.merge({k: float(v) for k, v in body.items()})
        self._send_json(200, {})

    def _handle_assign(self, parsed: urllib.parse.ParseResult):
        """Return an optimal shard assignment for a node given its hardware profile.

        Query params:
          model   — model preset name (default: first preset)
          device  — "cuda" | "cpu"
          vram_mb — integer VRAM in MB (0 for CPU)

        The greedy strategy: find the first gap in current layer coverage
        and assign it.  If no gap exists, assign the full model range so the
        node provides redundant coverage.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)

        model_list = params.get("model")
        if not model_list:
            model = next(iter(server.model_presets), None)
            if model is None:
                self._send_json(503, {"error": "no model presets configured"})
                return
        else:
            model = model_list[0]

        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start, required_end = _preset_layer_bounds(preset)

        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        device = params.get("device", ["cpu"])[0]
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        max_layers = required_end - required_start + 1
        if device != "cuda" or vram_mb < 8192:
            max_layers = min(max_layers, 16)

        # Collect covered intervals sorted by start layer.
        covered = sorted(
            [
                (n.shard_start, n.shard_end)
                for n in alive
                if n.shard_start is not None and n.shard_end is not None
            ],
            key=lambda t: t[0],
        )

        # Walk from required_start to find the first uncovered layer.
        gap_start = required_start
        for s, e in covered:
            if s <= gap_start:
                gap_start = max(gap_start, e + 1)
            else:
                break  # gap found before this node

        if gap_start > required_end:
            # Full coverage exists — assign the full range for redundancy.
            shard_start = required_start
            shard_end = min(required_end, shard_start + max_layers - 1)
        else:
            shard_start = gap_start
            shard_end = min(required_end, shard_start + max_layers - 1)

        peers = [
            {"endpoint": node.endpoint, "checksum": node.shard_checksum}
            for node in alive
            if node.model == model
            and node.shard_start == shard_start
            and node.shard_end == shard_end
            and node.shard_checksum
        ]

        self._send_json(200, {
            "shard_start": shard_start,
            "shard_end": shard_end,
            "model": model,
            "model_layers_end": required_end,
            "peers": peers,
            **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
        })

    def _handle_network_assign(self, parsed: urllib.parse.ParseResult):
        """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes.

        Query params:
          vram_mb  — integer VRAM in MB (0 = CPU-only node)
          device   — "cuda" | "cpu"
          hf_repo  — optional; if set, restrict search to this repo only

        Returns:
          {hf_repo, shard_start, shard_end, num_layers, gap_found}
          gap_found=true means a real uncovered gap was assigned; false means redundancy.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        device = params.get("device", ["cpu"])[0]
        filter_repo = params.get("hf_repo", [None])[0]  # optional repo filter

        with server.lock:
            self._purge_expired_nodes()
            all_nodes = list(server.registry.values())

        # Collect only nodes that registered a real HF model (have hf_repo + shard bounds).
        hf_nodes = [
            n for n in all_nodes
            if n.hf_repo
            and n.shard_start is not None
            and n.shard_end is not None
            and n.num_layers is not None
            and (filter_repo is None or n.hf_repo == filter_repo)
        ]

        if not hf_nodes:
            msg = (
                f"no HF-model nodes registered for {filter_repo!r}"
                if filter_repo
                else "no HF-model nodes registered; cannot assign shards"
            )
            self._send_json(503, {"error": msg})
            return

        # Group by hf_repo; pick the one with the largest total_layers and biggest gap.
        from collections import defaultdict
        repo_groups: dict = defaultdict(list)
        repo_layers: dict = {}
        for n in hf_nodes:
            repo_groups[n.hf_repo].append(n)
            # Use the largest num_layers seen for this repo.
            if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]:
                repo_layers[n.hf_repo] = n.num_layers

        # Pick the repo where the gap is largest (most work to do).
        best_repo = None
        best_gap_size = -1
        best_gap_start = 0
        best_num_layers = 0

        for repo, nodes in repo_groups.items():
            total = repo_layers[repo]
            covered = sorted(
                [(n.shard_start, n.shard_end) for n in nodes],
                key=lambda t: t[0],
            )
            # Walk from 0 to find first uncovered layer.
            gap_start = 0
            for s, e in covered:
                if s <= gap_start:
                    gap_start = max(gap_start, e + 1)
                else:
                    break
            gap_size = max(0, (total - 1) - gap_start + 1)  # layers remaining uncovered
            if gap_size > best_gap_size:
                best_gap_size = gap_size
                best_gap_start = gap_start
                best_repo = repo
                best_num_layers = total

        gap_found = best_gap_size > 0
        if not gap_found:
            # All shards are covered — still assign to the model with most nodes for redundancy.
            best_repo = max(repo_groups, key=lambda r: len(repo_groups[r]))
            best_gap_start = 0
            best_num_layers = repo_layers[best_repo]

        # Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM.
        total_l = best_num_layers
        if device == "cuda" and vram_mb >= 8192:
            max_layers = total_l
        else:
            max_layers = max(1, total_l // 2)

        shard_start = best_gap_start
        shard_end = min(total_l - 1, shard_start + max_layers - 1)

        self._send_json(200, {
            "hf_repo": best_repo,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "num_layers": total_l,
            "gap_found": gap_found,
        })

    def _handle_route(self, parsed: urllib.parse.ParseResult):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model_list = params.get("model")
        if not model_list:
            self._send_json(400, {"error": "missing 'model' query parameter"})
            return

        model = model_list[0]
        preset = server.model_presets.get(model)

        with server.lock:
            self._purge_expired_nodes()
            if preset is not None:
                # Preset-based routing (stub-model system).
                alive = [node for node in server.registry.values() if node.model == model]
                required_start, required_end = _preset_layer_bounds(preset)
            else:
                # HF model routing: match by hf_repo (full) or model short name.
                alive = [
                    node for node in server.registry.values()
                    if (node.hf_repo == model or node.model == model)
                    and node.shard_start is not None
                    and node.shard_end is not None
                    and node.num_layers is not None
                ]
                if not alive:
                    self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
                    return
                required_start = 0
                required_end = max(n.num_layers for n in alive) - 1  # type: ignore[type-var]

        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        route, error = _select_route(alive, required_start, required_end)
        if error:
            self._send_json(503, {"error": error})
            return

        self._send_json(200, {
            "route": [e.endpoint for e in route],
            "nodes": [
                {
                    "node_id": e.node_id,
                    "endpoint": e.endpoint,
                    "relay_addr": e.relay_addr,
                    "peer_id": e.peer_id,
                    "wallet_address": e.wallet_address,
                    "shard_start": e.shard_start,
                    "shard_end": e.shard_end,
                    "model": e.model,
                    "shard_checksum": e.shard_checksum,
                    "score": e.score,
                }
                for e in route
            ],
        })

    def _handle_routes(self, parsed: urllib.parse.ParseResult):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model_list = params.get("model")
        if not model_list:
            self._send_json(400, {"error": "missing 'model' query parameter"})
            return

        try:
            redundancy = int(params.get("redundancy", ["1"])[0])
        except ValueError:
            self._send_json(400, {"error": "redundancy must be an integer"})
            return
        if redundancy < 1:
            self._send_json(400, {"error": "redundancy must be at least 1"})
            return

        model = model_list[0]
        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start, required_end = _preset_layer_bounds(preset)

        with server.lock:
            self._purge_expired_nodes()
            candidates = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            candidates = [
                node for node in candidates
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        routes = []
        remaining = list(candidates)
        for _ in range(redundancy):
            route, error = _select_route(remaining, required_start, required_end)
            if error:
                self._send_json(503, {"error": error})
                return
            route_endpoints = {node.endpoint for node in route}
            remaining = [node for node in remaining if node.endpoint not in route_endpoints]
            routes.append({
                "route": [e.endpoint for e in route],
                "nodes": [
                    {
                        "node_id": e.node_id,
                        "endpoint": e.endpoint,
                        "relay_addr": e.relay_addr,
                        "peer_id": e.peer_id,
                        "wallet_address": e.wallet_address,
                        "shard_start": e.shard_start,
                        "shard_end": e.shard_end,
                        "model": e.model,
                        "shard_checksum": e.shard_checksum,
                        "score": e.score,
                    }
                    for e in route
                ],
            })

        self._send_json(200, {"routes": routes})


class TrackerServer:
    """HTTP tracker that manages node registration and resolves inference routes.

    Nodes register via POST /v1/nodes/register and keep themselves alive with
    POST /v1/nodes/<id>/heartbeat.  The gateway queries GET /v1/route?model=<name>
    to obtain an ordered list of node endpoints whose shards cover all layers.
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 0,
        heartbeat_timeout: float = 90.0,
        rebalance_interval: float = 30.0,
        model_presets: dict | None = None,
        contracts: Any | None = None,
        minimum_stake: int = 0,
        cluster_peers: list[str] | None = None,
        cluster_self_url: str | None = None,
        relay_url: str | None = None,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._heartbeat_timeout = heartbeat_timeout
        self._rebalance_interval = rebalance_interval
        self._model_presets: dict = (
            model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS)
        )
        self._contracts = contracts
        self._minimum_stake = minimum_stake
        self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else []
        self._cluster_self_url: str | None = cluster_self_url
        self._relay_url = relay_url
        self._registry: dict[str, _NodeEntry] = {}
        self._lock = threading.Lock()
        self._server: _TrackerHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self._rebalance_stop = threading.Event()
        self._rebalance_thread: threading.Thread | None = None
        self._raft: RaftNode | None = None
        self._gossip: NodeGossip | None = None
        self.port: int | None = None

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

        # Start HTTP server first so we know our port
        self._server = _TrackerHTTPServer(
            (self._host, self._requested_port),
            _TrackerHandler,
            self._registry,
            self._lock,
            self._heartbeat_timeout,
            self._model_presets,
            self._contracts,
            self._minimum_stake,
            self._relay_url,
        )
        self.port = self._server.server_address[1]

        # Start Raft + gossip if cluster peers are configured
        if self._cluster_peers:
            self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
            self._raft = RaftNode(
                self_url=self_url,
                peers=self._cluster_peers,
                apply_fn=self._raft_apply,
            )
            self._gossip = NodeGossip(peers=self._cluster_peers)
            self._server.raft = self._raft
            self._server.gossip = self._gossip
            self._raft.start()
            self._gossip.start()

        self._rebalance_stop.clear()
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True)
        self._rebalance_thread.start()
        return self.port

    def _raft_apply(self, command: str, payload: dict) -> None:
        """Called by RaftNode when a log entry is committed — replicate to local registry."""
        if command != "register":
            return
        # Re-apply the registration to our local registry (follower path).
        # On the leader this is a no-op since it already registered locally.
        node_id = payload.get("node_id")
        if not node_id or node_id in self._registry:
            return  # already present (leader case) or malformed
        endpoint = payload.get("endpoint", "")
        try:
            shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None
            shard_end   = int(payload["shard_end"])   if payload.get("shard_end")   is not None else None
        except (TypeError, ValueError):
            return
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            model=payload.get("model", "stub-model"),
            shard_checksum=payload.get("shard_checksum"),
            hardware_profile=payload.get("hardware_profile", {}),
            wallet_address=payload.get("wallet_address"),
            score=float(payload.get("score", 1.0)),
            tracker_mode=bool(payload.get("tracker_mode", False)),
            hf_repo=payload.get("hf_repo"),
            num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None,
        )
        with self._lock:
            self._registry[node_id] = entry

    def _rebalance_loop(self) -> None:
        while not self._rebalance_stop.wait(self._rebalance_interval):
            server = self._server
            if server is None:
                return
            with self._lock:
                _purge_expired_nodes_locked(server)
                _rebalance_all_locked(server)

    def stop(self) -> None:
        if self._raft is not None:
            self._raft.stop()
        if self._gossip is not None:
            self._gossip.stop()
        if self._server is None:
            return
        self._rebalance_stop.set()
        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        if self._rebalance_thread is not None:
            self._rebalance_thread.join(timeout=1)
        self._server = None
        self._thread = None
        self._rebalance_thread = None
        self._raft = None
        self._gossip = None
        self.port = None