"""Mesh router -- routes requests across devices in the multi-device mesh.

Decision factors:
- Model availability (which device has this model?)
- Latency (prefer lower latency devices)
- Load (prefer less busy devices)
- Cost (local GPU = $0, remote GPU = $0, cloud = $$)
- Privacy (PII-sensitive requests stay local)
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING, Any, cast

import httpx
from pydantic import BaseModel

if TYPE_CHECKING:
    from llmhosts.mesh.manager import DeviceManager
    from llmhosts.mesh.models import Device

logger = logging.getLogger(__name__)


class MeshRoute(BaseModel):
    """The result of a routing decision -- where to send a request."""

    device_id: str
    device_name: str
    tunnel_url: str | None
    is_local: bool
    model: str
    estimated_latency_ms: float | None
    reason: str  # "Model available locally" / "Forwarded to home-desktop (24GB GPU)"


class MeshRouter:
    """Routes requests across devices in the mesh.

    Priority order:
    1. Local device (if model available)
    2. Lowest-latency remote device with model
    3. Any remote device (may need to pull model)
    4. Cloud fallback (returns None -- caller handles cloud routing)
    """

    def __init__(self, device_manager: DeviceManager) -> None:
        self._manager = device_manager

    async def route(self, model: str, *, prefer_local: bool = True) -> MeshRoute | None:
        """Find the best device to handle a request for *model*.

        Returns
        -------
        MeshRoute | None
            The routing decision, or ``None`` if no mesh device can serve
            the model and cloud fallback is needed.
        """
        devices = await self._manager.list_devices()
        if not devices:
            return None

        # Separate into candidates that have the model vs. those that don't
        local_with_model: Device | None = None
        remote_with_model: list[Device] = []
        remote_without_model: list[Device] = []

        for device in devices:
            if not device.healthy:
                continue

            has_model = model in device.models

            if device.is_local and has_model:
                local_with_model = device
            elif not device.is_local and has_model:
                remote_with_model.append(device)
            elif not device.is_local:
                remote_without_model.append(device)

        # Priority 1: Local device with model
        if prefer_local and local_with_model is not None:
            return MeshRoute(
                device_id=local_with_model.id,
                device_name=local_with_model.name,
                tunnel_url=None,
                is_local=True,
                model=model,
                estimated_latency_ms=0.0,
                reason="Model available locally",
            )

        # Priority 2: Lowest-latency remote device with model
        if remote_with_model:
            # Sort by latency (None latency = high value to push to end)
            remote_with_model.sort(key=lambda d: d.latency_ms if d.latency_ms is not None else 999999.0)
            best = remote_with_model[0]
            hw_info = f" ({best.hardware_summary})" if best.hardware_summary else ""
            return MeshRoute(
                device_id=best.id,
                device_name=best.name,
                tunnel_url=best.tunnel_url,
                is_local=False,
                model=model,
                estimated_latency_ms=best.latency_ms,
                reason=f"Forwarded to {best.name}{hw_info}",
            )

        # Priority 3: Local device without model (caller may need to pull)
        if prefer_local and local_with_model is None:
            local = await self._manager.get_local_device()
            if local and local.healthy:
                return MeshRoute(
                    device_id=local.id,
                    device_name=local.name,
                    tunnel_url=None,
                    is_local=True,
                    model=model,
                    estimated_latency_ms=0.0,
                    reason=f"Model '{model}' not cached locally -- may need pull",
                )

        # Priority 4: Any remote device (may need to pull model)
        if remote_without_model:
            remote_without_model.sort(key=lambda d: d.latency_ms if d.latency_ms is not None else 999999.0)
            best = remote_without_model[0]
            return MeshRoute(
                device_id=best.id,
                device_name=best.name,
                tunnel_url=best.tunnel_url,
                is_local=False,
                model=model,
                estimated_latency_ms=best.latency_ms,
                reason=f"Forwarded to {best.name} (model may need pull)",
            )

        # No mesh device available -- cloud fallback
        return None

    async def forward_request(self, device: Device, path: str, body: dict[str, Any]) -> dict[str, Any]:
        """Forward a request to a remote device via its tunnel URL.

        Sends ``POST {device.tunnel_url}{path}`` with *body* as JSON.

        Parameters
        ----------
        device:
            The target device (must have a ``tunnel_url``).
        path:
            The API path, e.g. ``/v1/chat/completions``.
        body:
            The JSON request body.

        Returns
        -------
        dict[str, Any]
            The parsed JSON response from the remote device.

        Raises
        ------
        ValueError
            If the device has no tunnel URL.
        httpx.HTTPError
            On network failures.
        """
        if not device.tunnel_url:
            raise ValueError(f"Device '{device.name}' has no tunnel URL configured")

        url = device.tunnel_url.rstrip("/") + path
        logger.info("Forwarding request to %s: %s", device.name, url)

        start = time.monotonic()
        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(url, json=body)
        elapsed_ms = (time.monotonic() - start) * 1000

        logger.info(
            "Response from %s: status=%d, latency=%.0fms",
            device.name,
            resp.status_code,
            elapsed_ms,
        )

        resp.raise_for_status()
        return cast("dict[str, Any]", resp.json())
