"""Leader election and distributed scheduling coordinator."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Optional

from radix_edge.mesh.peer_registry import PeerRegistry

if TYPE_CHECKING:
    from radix_edge.mesh.peer_client import PeerClient

logger = logging.getLogger("radix_edge.mesh.coordinator")


class MeshCoordinator:
    """Bully-algorithm leader election and coordinator-side distributed scheduling.

    The node with the highest ``node_id`` (lexicographic) among active peers wins
    the election and becomes the coordinator.
    """

    def __init__(self, node_id: str, registry: PeerRegistry):
        self.node_id = node_id
        self.registry = registry
        self._is_coordinator: bool = False
        self._coordinator_id: Optional[str] = None

    def elect(self) -> str:
        """Run a bully election. Returns the coordinator node_id."""
        all_ids = [p.node_id for p in self.registry.get_active_peers()] + [self.node_id]
        # Deduplicate (self may also be in the registry in some edge cases)
        unique_ids = sorted(set(all_ids))
        winner = max(unique_ids)
        self._coordinator_id = winner
        self._is_coordinator = (winner == self.node_id)
        self.registry.set_coordinator(winner)

        if self._is_coordinator:
            logger.info("Election result: this node (%s) is coordinator", self.node_id)
        else:
            logger.info("Election result: coordinator is %s", winner)

        return winner

    @property
    def is_coordinator(self) -> bool:
        """Whether this node is the current coordinator."""
        return self._is_coordinator

    @property
    def coordinator_id(self) -> Optional[str]:
        """The node_id of the current coordinator, or None if no election has run."""
        return self._coordinator_id

    def accept_coordinator(self, coordinator_id: str) -> None:
        """Accept an externally announced coordinator (from election announce message)."""
        self._coordinator_id = coordinator_id
        self._is_coordinator = (coordinator_id == self.node_id)
        self.registry.set_coordinator(coordinator_id)
        logger.info("Accepted coordinator: %s (self=%s)", coordinator_id, self._is_coordinator)

    async def schedule_across_mesh(self, peer_clients: dict[str, "PeerClient"]) -> int:
        """Coordinator: collect queued jobs from all peers, assign to best peer.

        1. Gather queued jobs from all peers via peer_clients.
        2. Gather aggregate GPU state from registry.
        3. For each job, pick best peer based on GPU availability.
        4. Assign via peer_client.assign_job().

        Returns the number of jobs assigned.
        """
        if not self._is_coordinator:
            return 0

        assigned_count = 0

        # Collect queued jobs from all peers
        all_jobs: list[tuple[str, dict]] = []  # (source_node_id, job_data)
        for nid, client in peer_clients.items():
            try:
                jobs = await client.get_queued_jobs()
                for job in jobs:
                    all_jobs.append((nid, job))
            except Exception:
                logger.warning("Failed to get queued jobs from peer %s", nid)

        if not all_jobs:
            return 0

        # Build peer capacity map: node_id -> available GPU count
        capacity: dict[str, int] = {}
        for peer in self.registry.get_active_peers():
            caps = peer.capabilities
            gpu_count = caps.get("gpu_count", 0)
            running = caps.get("running_jobs", 0)
            available = max(0, gpu_count - running)
            capacity[peer.node_id] = available

        # Include self capacity (not in peer registry)
        # Self capacity is handled by the local scheduler

        # Assign jobs to peers with available capacity
        for source_nid, job_data in all_jobs:
            num_gpus = job_data.get("num_gpus", 1)

            # Find best peer: most available GPUs
            best_nid = None
            best_available = -1
            for nid, avail in capacity.items():
                if avail >= num_gpus and avail > best_available:
                    best_available = avail
                    best_nid = nid

            if best_nid is None:
                logger.debug("No peer with capacity for job %s (needs %d GPUs)", job_data.get("job_id"), num_gpus)
                continue

            # If the best peer is the source, skip (it can schedule locally)
            if best_nid == source_nid:
                continue

            # Assign to the best peer
            if best_nid in peer_clients:
                try:
                    job_id = job_data.get("job_id", "unknown")
                    await peer_clients[best_nid].assign_job(job_id, list(range(num_gpus)))
                    capacity[best_nid] -= num_gpus
                    assigned_count += 1
                    logger.info("Assigned job %s to peer %s", job_id, best_nid)
                except Exception:
                    logger.warning("Failed to assign job %s to peer %s", job_data.get("job_id"), best_nid)

        return assigned_count
