"""Tests for mesh coordinator (bully election + distributed scheduling)."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

import pytest

from radix_edge.mesh.coordinator import MeshCoordinator
from radix_edge.mesh.peer_registry import PeerRegistry


class TestBullyElection:
    def test_highest_id_wins(self):
        reg = PeerRegistry()
        reg.update_peer("node-aaa", "10.0.0.1", 8080, {})
        reg.update_peer("node-bbb", "10.0.0.2", 8080, {})
        reg.update_peer("node-ccc", "10.0.0.3", 8080, {})

        coord = MeshCoordinator("node-ddd", reg)
        winner = coord.elect()
        assert winner == "node-ddd"
        assert coord.is_coordinator is True

    def test_peer_wins_when_higher(self):
        reg = PeerRegistry()
        reg.update_peer("node-zzz", "10.0.0.1", 8080, {})
        reg.update_peer("node-aaa", "10.0.0.2", 8080, {})

        coord = MeshCoordinator("node-bbb", reg)
        winner = coord.elect()
        assert winner == "node-zzz"
        assert coord.is_coordinator is False

    def test_single_node_becomes_coordinator(self):
        reg = PeerRegistry()
        coord = MeshCoordinator("node-solo", reg)
        winner = coord.elect()
        assert winner == "node-solo"
        assert coord.is_coordinator is True

    def test_coordinator_id_property(self):
        reg = PeerRegistry()
        coord = MeshCoordinator("node-a", reg)
        assert coord.coordinator_id is None

        coord.elect()
        assert coord.coordinator_id == "node-a"

    def test_accept_coordinator(self):
        reg = PeerRegistry()
        reg.update_peer("node-other", "10.0.0.1", 8080, {})
        coord = MeshCoordinator("node-me", reg)

        coord.accept_coordinator("node-other")
        assert coord.coordinator_id == "node-other"
        assert coord.is_coordinator is False

    def test_accept_self_as_coordinator(self):
        reg = PeerRegistry()
        coord = MeshCoordinator("node-me", reg)

        coord.accept_coordinator("node-me")
        assert coord.is_coordinator is True

    def test_re_election_after_peer_removed(self):
        reg = PeerRegistry()
        reg.update_peer("node-zzz", "10.0.0.1", 8080, {})

        coord = MeshCoordinator("node-bbb", reg)
        winner = coord.elect()
        assert winner == "node-zzz"
        assert coord.is_coordinator is False

        # Remove the winner peer
        reg.remove_peer("node-zzz")
        winner = coord.elect()
        assert winner == "node-bbb"
        assert coord.is_coordinator is True

    def test_duplicate_self_in_registry(self):
        """If self node_id is also in the registry, should not break."""
        reg = PeerRegistry()
        reg.update_peer("node-me", "10.0.0.1", 8080, {})

        coord = MeshCoordinator("node-me", reg)
        winner = coord.elect()
        assert winner == "node-me"
        assert coord.is_coordinator is True


class TestScheduleAcrossMesh:
    @pytest.mark.asyncio
    async def test_schedule_assigns_jobs(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 4, "running_jobs": 0})
        reg.update_peer("node-b", "10.0.0.2", 8080, {"gpu_count": 2, "running_jobs": 1})

        coord = MeshCoordinator("node-coord", reg)
        coord.elect()
        assert coord.is_coordinator is True

        # Mock peer clients
        client_a = AsyncMock()
        client_a.get_queued_jobs.return_value = []
        client_a.assign_job.return_value = {"status": "accepted"}

        client_b = AsyncMock()
        client_b.get_queued_jobs.return_value = [
            {"job_id": "job-1", "num_gpus": 1},
        ]
        client_b.assign_job.return_value = {"status": "accepted"}

        peer_clients = {"node-a": client_a, "node-b": client_b}
        assigned = await coord.schedule_across_mesh(peer_clients)

        # job-1 from node-b should be assigned to node-a (more capacity)
        assert assigned == 1
        client_a.assign_job.assert_called_once()

    @pytest.mark.asyncio
    async def test_non_coordinator_does_nothing(self):
        reg = PeerRegistry()
        reg.update_peer("node-zzz", "10.0.0.1", 8080, {"gpu_count": 4, "running_jobs": 0})

        coord = MeshCoordinator("node-aaa", reg)
        coord.elect()
        assert coord.is_coordinator is False

        assigned = await coord.schedule_across_mesh({})
        assert assigned == 0

    @pytest.mark.asyncio
    async def test_no_capacity_no_assignment(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 1, "running_jobs": 1})

        coord = MeshCoordinator("node-coord", reg)
        coord.elect()

        client_a = AsyncMock()
        client_a.get_queued_jobs.return_value = [{"job_id": "job-1", "num_gpus": 2}]

        assigned = await coord.schedule_across_mesh({"node-a": client_a})
        assert assigned == 0

    @pytest.mark.asyncio
    async def test_peer_failure_handled_gracefully(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 4, "running_jobs": 0})
        reg.update_peer("node-b", "10.0.0.2", 8080, {"gpu_count": 2, "running_jobs": 0})

        coord = MeshCoordinator("node-coord", reg)
        coord.elect()

        client_a = AsyncMock()
        client_a.get_queued_jobs.side_effect = Exception("Connection refused")

        client_b = AsyncMock()
        client_b.get_queued_jobs.return_value = []

        # Should not raise
        assigned = await coord.schedule_across_mesh({"node-a": client_a, "node-b": client_b})
        assert assigned == 0

    @pytest.mark.asyncio
    async def test_no_jobs_returns_zero(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 4, "running_jobs": 0})

        coord = MeshCoordinator("node-coord", reg)
        coord.elect()

        client_a = AsyncMock()
        client_a.get_queued_jobs.return_value = []

        assigned = await coord.schedule_across_mesh({"node-a": client_a})
        assert assigned == 0
