"""Device pool: aggregating device state across hosts with staleness tracking.

Capability: ``device-pool``.

The pool never reaches out over the network. Each host's own agent calls
``DevicePool.sync_host_devices(host_id, snapshot)`` to push its current
``DeviceManager.list_devices()`` result. A host whose heartbeat goes stale
has its devices lazily reported as ``unreachable`` on the next read.
"""

from __future__ import annotations

from dataclasses import dataclass, field, replace
from datetime import datetime
from typing import TYPE_CHECKING, Literal

from core.models import Device, utc_now

if TYPE_CHECKING:
    from cloud.config import CloudConfig
    from cloud.store import CloudStore


PooledDeviceStatus = Literal["idle", "busy", "offline", "error", "unreachable"]

# Statuses that map 1:1 from a host's snapshot. ``unreachable`` is pool-only.
_HOST_REPORTED_STATUSES = ("idle", "busy", "offline", "error")


@dataclass(frozen=True)
class HostRegistration:
    """A host process that has registered itself with the pool."""

    host_id: str
    address: str | None
    last_seen_at: datetime
    planner_transport: Literal["direct", "cloud"] = "direct"


@dataclass(frozen=True)
class PooledDevice:
    """A device owned by a registered host, as seen by the pool."""

    device_id: str
    host_id: str
    driver_type: str
    status: PooledDeviceStatus
    capability_tags: list[str] = field(default_factory=list)
    synced_at: datetime | None = None
    mcp_busy: bool = False


class DevicePool:
    """Aggregates ``Device`` snapshots pushed by many host processes."""

    def __init__(self, store: "CloudStore", config: "CloudConfig") -> None:
        self.store = store
        self.config = config

    def sync_host_devices(
        self,
        host_id: str,
        snapshot: list[Device],
        *,
        address: str | None = None,
        planner_transport: Literal["direct", "cloud"] = "direct",
        allow_device_takeover: bool = False,
        mcp_busy_device_ids: list[str] | None = None,
    ) -> None:
        """Push a host's current device snapshot into the pool.

        Updates the host's ``last_seen_at`` and atomically replaces its
        previously-stored device rows with the new snapshot. Devices from
        other hosts are untouched.
        """
        now = utc_now()
        self.store.upsert_host(
            host_id,
            address=address,
            last_seen_at=now,
            planner_transport=planner_transport,
        )
        busy_set = set(mcp_busy_device_ids or [])
        devices = [
            self._to_pooled(device, host_id, now, mcp_busy=device.id in busy_set)
            for device in snapshot
        ]
        if allow_device_takeover:
            self.store.replace_host_devices(
                host_id,
                devices,
                allow_device_takeover=True,
            )
        else:
            self.store.replace_host_devices(host_id, devices)

    def list_devices(self) -> list[PooledDevice]:
        devices = self.store.list_devices()
        if not devices:
            return []
        hosts = {h.host_id: h for h in self.store.list_hosts()}
        now = utc_now()
        result: list[PooledDevice] = []
        for device in devices:
            host = hosts.get(device.host_id)
            if host is not None and self._is_stale(host, now):
                device = self._as_unreachable(device)
            result.append(device)
        return result

    def get_device(self, device_id: str) -> PooledDevice | None:
        device = self.store.get_device(device_id)
        if device is None:
            return None
        host = self.store.get_host(device.host_id)
        if host is not None and self._is_stale(host, utc_now()):
            return self._as_unreachable(device)
        return device

    def list_hosts(self) -> list[HostRegistration]:
        return self.store.list_hosts()

    def _to_pooled(
        self,
        device: Device,
        host_id: str,
        synced_at: datetime,
        *,
        mcp_busy: bool = False,
    ) -> PooledDevice:
        raw_status = (
            device.status if device.status in _HOST_REPORTED_STATUSES else "idle"
        )
        tags = list(device.capability_tags or [])
        return PooledDevice(
            device_id=device.id,
            host_id=host_id,
            driver_type=device.driver_type,
            status=raw_status,  # type: ignore[arg-type]
            capability_tags=tags,
            synced_at=synced_at,
            mcp_busy=mcp_busy,
        )

    def _is_stale(self, host: HostRegistration, now: datetime) -> bool:
        age = (now - host.last_seen_at).total_seconds()
        return age > self.config.stale_after_seconds

    def _as_unreachable(self, device: PooledDevice) -> PooledDevice:
        if device.status == "unreachable":
            return device
        return replace(device, status="unreachable")