"""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,
    "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}
"""

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


DEFAULT_MODEL_PRESETS: dict[str, dict] = {
    "stub-model": {"layers_start": 0, "layers_end": 31},
}


class _NodeEntry:
    __slots__ = (
        "node_id", "endpoint", "shard_start", "shard_end",
        "hardware_profile", "wallet_address", "score", "last_heartbeat",
    )

    def __init__(
        self,
        node_id: str,
        endpoint: str,
        shard_start: int,
        shard_end: int,
        hardware_profile: dict,
        wallet_address: str | None,
        score: float,
    ) -> None:
        self.node_id = node_id
        self.endpoint = endpoint
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.hardware_profile = hardware_profile
        self.wallet_address = wallet_address
        self.score = score
        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(nodes, key=lambda n: (n.shard_start, -n.shard_end))
    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, ""


class _TrackerHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr: tuple,
        handler,
        registry: dict,
        lock: threading.Lock,
        heartbeat_timeout: float,
        model_presets: dict,
        contracts: Any | None,
        minimum_stake: int,
    ) -> 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


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()
        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:
            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]
        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:
            del server.registry[node_id]

    def do_POST(self):
        if self.path == "/v1/nodes/register":
            self._handle_register()
            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/nodes/assign":
            self._handle_assign(parsed)
        elif parsed.path == "/v1/models":
            self._handle_models()
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_models(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        data = [
            {"id": name, "object": "model", "created": created, "owned_by": "meshnet"}
            for name in server.model_presets
        ]
        self._send_json(200, {"object": "list", "data": data})

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

        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

        try:
            shard_start = int(body["shard_start"])
            shard_end = int(body["shard_end"])
            score = float(body.get("score", 1.0))
        except (KeyError, TypeError, ValueError):
            self._send_json(400, {"error": "shard_start, shard_end, and score 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

        hardware_profile = body.get("hardware_profile", {})
        if not isinstance(hardware_profile, dict):
            self._send_json(400, {"error": "hardware_profile must be an object"})
            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

        node_id = str(uuid.uuid4())
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            hardware_profile=hardware_profile,
            wallet_address=wallet_address,
            score=score,
        )
        with server.lock:
            self._purge_expired_nodes()
            server.registry[node_id] = entry

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

    def _handle_heartbeat(self, node_id: str):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        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()
        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: int = preset["layers_start"]
        required_end: int = preset["layers_end"]

        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
            ]

        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],
            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)

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

    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)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start: int = preset["layers_start"]
        required_end: int = preset["layers_end"]

        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
            ]

        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": [
                {
                    "endpoint": e.endpoint,
                    "wallet_address": e.wallet_address,
                    "shard_start": e.shard_start,
                    "shard_end": e.shard_end,
                    "score": e.score,
                }
                for e in route
            ],
        })


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 = 30.0,
        model_presets: dict | None = None,
        contracts: Any | None = None,
        minimum_stake: int = 0,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._heartbeat_timeout = heartbeat_timeout
        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._registry: dict[str, _NodeEntry] = {}
        self._lock = threading.Lock()
        self._server: _TrackerHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("TrackerServer is already running")
        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.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self.port

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