"""Tests for distributed mesh scheduling across multiple simulated peers."""

from __future__ import annotations

from unittest.mock import AsyncMock

import pytest

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


# Use "node-zzz" as coordinator so it always wins the bully election (highest ID).


class TestDistributedScheduling:
    @pytest.mark.asyncio
    async def test_two_peers_job_moves_to_capacity(self):
        """Job on a busy peer should be assigned to an idle peer."""
        reg = PeerRegistry()
        reg.update_peer("node-busy", "10.0.0.1", 8080, {"gpu_count": 2, "running_jobs": 2})
        reg.update_peer("node-idle", "10.0.0.2", 8080, {"gpu_count": 4, "running_jobs": 0})

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

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

        client_idle = AsyncMock()
        client_idle.get_queued_jobs.return_value = []
        client_idle.assign_job.return_value = {"status": "accepted"}

        peer_clients = {"node-busy": client_busy, "node-idle": client_idle}
        assigned = await coord.schedule_across_mesh(peer_clients)

        assert assigned == 1
        client_idle.assign_job.assert_called_once_with("job-1", [0])

    @pytest.mark.asyncio
    async def test_three_peers_best_capacity_wins(self):
        """Among multiple peers, the one with most available GPUs gets the job."""
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 2, "running_jobs": 1})  # 1 free
        reg.update_peer("node-b", "10.0.0.2", 8080, {"gpu_count": 4, "running_jobs": 0})  # 4 free
        reg.update_peer("node-c", "10.0.0.3", 8080, {"gpu_count": 8, "running_jobs": 2})  # 6 free

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

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

        client_b = AsyncMock()
        client_b.get_queued_jobs.return_value = []
        client_b.assign_job.return_value = {"status": "accepted"}

        client_c = AsyncMock()
        client_c.get_queued_jobs.return_value = []
        client_c.assign_job.return_value = {"status": "accepted"}

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

        assert assigned == 1
        # node-c has 6 free GPUs, most capacity
        client_c.assign_job.assert_called_once_with("job-1", [0, 1])

    @pytest.mark.asyncio
    async def test_job_stays_on_source_if_best(self):
        """If the source peer has the most capacity, job should NOT be reassigned."""
        reg = PeerRegistry()
        reg.update_peer("node-source", "10.0.0.1", 8080, {"gpu_count": 8, "running_jobs": 0})
        reg.update_peer("node-other", "10.0.0.2", 8080, {"gpu_count": 2, "running_jobs": 1})

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

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

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

        peer_clients = {"node-source": client_source, "node-other": client_other}
        assigned = await coord.schedule_across_mesh(peer_clients)

        # Job stays on source (source has best capacity, so skip re-assignment)
        assert assigned == 0

    @pytest.mark.asyncio
    async def test_multi_gpu_job_needs_sufficient_capacity(self):
        """A job requiring 4 GPUs should only go to a peer with 4+ free GPUs."""
        reg = PeerRegistry()
        reg.update_peer("node-small", "10.0.0.1", 8080, {"gpu_count": 2, "running_jobs": 0})
        reg.update_peer("node-large", "10.0.0.2", 8080, {"gpu_count": 8, "running_jobs": 0})

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

        client_small = AsyncMock()
        client_small.get_queued_jobs.return_value = [
            {"job_id": "job-big", "num_gpus": 4},
        ]

        client_large = AsyncMock()
        client_large.get_queued_jobs.return_value = []
        client_large.assign_job.return_value = {"status": "accepted"}

        peer_clients = {"node-small": client_small, "node-large": client_large}
        assigned = await coord.schedule_across_mesh(peer_clients)

        assert assigned == 1
        client_large.assign_job.assert_called_once()

    @pytest.mark.asyncio
    async def test_fallback_when_peer_unreachable(self):
        """If assignment to the best peer fails, the job count should reflect that."""
        reg = PeerRegistry()
        reg.update_peer("node-source", "10.0.0.2", 8080, {"gpu_count": 1, "running_jobs": 1})
        reg.update_peer("node-target", "10.0.0.1", 8080, {"gpu_count": 4, "running_jobs": 0})

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

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

        client_target = AsyncMock()
        client_target.get_queued_jobs.return_value = []
        client_target.assign_job.side_effect = Exception("Connection refused")

        peer_clients = {"node-source": client_source, "node-target": client_target}
        assigned = await coord.schedule_across_mesh(peer_clients)

        # Assignment failed, count should be 0
        assert assigned == 0

    @pytest.mark.asyncio
    async def test_multiple_jobs_from_different_peers(self):
        """Jobs from multiple peers can be distributed."""
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 1, "running_jobs": 1})  # 0 free
        reg.update_peer("node-b", "10.0.0.2", 8080, {"gpu_count": 1, "running_jobs": 1})  # 0 free
        reg.update_peer("node-c", "10.0.0.3", 8080, {"gpu_count": 8, "running_jobs": 0})  # 8 free

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

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

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

        client_c = AsyncMock()
        client_c.get_queued_jobs.return_value = []
        client_c.assign_job.return_value = {"status": "accepted"}

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

        # Both jobs should go to node-c (only one with capacity)
        assert assigned == 2
        assert client_c.assign_job.call_count == 2
