"""In-memory peer state registry for mesh networking."""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class PeerState:
    """State of a single mesh peer."""

    node_id: str
    host: str
    port: int
    capabilities: dict = field(default_factory=dict)  # {gpu_count, gpus, running_jobs, queued_jobs}
    last_seen: float = field(default_factory=time.monotonic)
    is_coordinator: bool = False


class PeerRegistry:
    """Thread-safe in-memory registry of mesh peers."""

    def __init__(self, stale_threshold: float = 15.0):
        self._peers: dict[str, PeerState] = {}
        self._lock = threading.Lock()
        self._stale_threshold = stale_threshold

    def update_peer(self, node_id: str, host: str, port: int, capabilities: dict) -> None:
        """Add or update a peer entry."""
        with self._lock:
            existing = self._peers.get(node_id)
            self._peers[node_id] = PeerState(
                node_id=node_id,
                host=host,
                port=port,
                capabilities=capabilities,
                last_seen=time.monotonic(),
                is_coordinator=existing.is_coordinator if existing else False,
            )

    def remove_peer(self, node_id: str) -> None:
        """Remove a peer from the registry."""
        with self._lock:
            self._peers.pop(node_id, None)

    def get_active_peers(self) -> list[PeerState]:
        """Return peers that have been seen within the stale threshold."""
        now = time.monotonic()
        with self._lock:
            return [
                p
                for p in self._peers.values()
                if (now - p.last_seen) <= self._stale_threshold
            ]

    def get_stale_peers(self) -> list[PeerState]:
        """Return peers that have not been seen within the stale threshold."""
        now = time.monotonic()
        with self._lock:
            return [
                p
                for p in self._peers.values()
                if (now - p.last_seen) > self._stale_threshold
            ]

    def get_aggregate_gpus(self) -> list[dict]:
        """Return all GPUs across active peers, tagged with node_id."""
        result = []
        for peer in self.get_active_peers():
            gpus = peer.capabilities.get("gpus", [])
            for gpu in gpus:
                tagged = dict(gpu)
                tagged["node_id"] = peer.node_id
                result.append(tagged)
        return result

    def get_peer(self, node_id: str) -> Optional[PeerState]:
        """Return a single peer by node_id, or None."""
        with self._lock:
            peer = self._peers.get(node_id)
            if peer is None:
                return None
            return peer

    def set_coordinator(self, node_id: str) -> None:
        """Mark a peer as the coordinator (clears previous coordinator flag)."""
        with self._lock:
            for p in self._peers.values():
                p.is_coordinator = (p.node_id == node_id)

    @property
    def peer_count(self) -> int:
        """Total number of registered peers (active + stale)."""
        with self._lock:
            return len(self._peers)
