"""Tests for mesh background sync loops."""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from radix_edge.mesh.coordinator import MeshCoordinator
from radix_edge.mesh.peer_registry import PeerRegistry
from radix_edge.mesh.sync import (
    _build_capabilities,
    mesh_heartbeat_loop,
    mesh_discovery_loop,
    election_monitor_loop,
    coordinator_loop,
)


class TestBuildCapabilities:
    def test_returns_dict_with_cpu_count(self, db):
        caps = _build_capabilities()
        assert "cpu_count" in caps
        assert isinstance(caps["cpu_count"], int)

    def test_returns_gpu_info(self, db):
        db.execute(
            "INSERT INTO gpus (gpu_index, name, memory_total_mb) VALUES (?, ?, ?)",
            (0, "A100", 81920),
        )
        db.commit()
        caps = _build_capabilities()
        assert caps["gpu_count"] == 1
        assert len(caps["gpus"]) == 1
        assert caps["gpus"][0]["name"] == "A100"

    def test_returns_job_counts(self, db):
        db.execute(
            "INSERT INTO jobs (job_id, user_id, image, status) VALUES (?, ?, ?, ?)",
            ("job-1", "user-1", "img:1", "running"),
        )
        db.execute(
            "INSERT INTO jobs (job_id, user_id, image, status) VALUES (?, ?, ?, ?)",
            ("job-2", "user-1", "img:1", "queued"),
        )
        db.commit()
        caps = _build_capabilities()
        assert caps["running_jobs"] == 1
        assert caps["queued_jobs"] == 1


class TestMeshHeartbeatLoop:
    @pytest.mark.asyncio
    async def test_sends_heartbeat_to_coordinator(self, db):
        config = MagicMock()
        config.mesh_heartbeat_interval = 0.01
        config.port = 8080

        reg = PeerRegistry()
        coord = MeshCoordinator("node-me", reg)
        reg.update_peer("node-coord", "10.0.0.1", 8080, {})
        coord.accept_coordinator("node-coord")

        mock_client = AsyncMock()
        peer_clients = {"node-coord": mock_client}

        task = asyncio.create_task(
            mesh_heartbeat_loop(config, reg, peer_clients, "node-me", coord)
        )
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

        mock_client.send_heartbeat.assert_called()

    @pytest.mark.asyncio
    async def test_sends_to_all_when_no_coordinator(self, db):
        config = MagicMock()
        config.mesh_heartbeat_interval = 0.01
        config.port = 8080

        reg = PeerRegistry()
        coord = MeshCoordinator("node-me", reg)
        # No coordinator elected yet

        client_a = AsyncMock()
        client_b = AsyncMock()
        peer_clients = {"node-a": client_a, "node-b": client_b}

        task = asyncio.create_task(
            mesh_heartbeat_loop(config, reg, peer_clients, "node-me", coord)
        )
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

        client_a.send_heartbeat.assert_called()
        client_b.send_heartbeat.assert_called()


class TestDiscoveryLoop:
    @pytest.mark.asyncio
    async def test_removes_stale_peers(self):
        config = MagicMock()
        config.mesh_discovery_interval = 0.01

        reg = PeerRegistry(stale_threshold=0.01)
        reg.update_peer("node-stale", "10.0.0.1", 8080, {})

        # Wait for staleness
        await asyncio.sleep(0.05)

        mock_discovery = MagicMock()
        mock_discovery.peers = []
        peer_clients: dict = {}

        task = asyncio.create_task(
            mesh_discovery_loop(config, mock_discovery, reg, peer_clients, "key")
        )
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

        assert reg.get_peer("node-stale") is None


class TestElectionMonitorLoop:
    @pytest.mark.asyncio
    async def test_triggers_election_when_coordinator_stale(self):
        config = MagicMock()
        config.mesh_stale_threshold = 0.01

        reg = PeerRegistry(stale_threshold=0.01)
        reg.update_peer("node-old-coord", "10.0.0.1", 8080, {})

        coord = MeshCoordinator("node-me", reg)
        coord.accept_coordinator("node-old-coord")

        # Wait for the old coordinator to go stale
        await asyncio.sleep(0.05)

        peer_clients: dict = {}
        task = asyncio.create_task(
            election_monitor_loop(config, coord, reg, peer_clients)
        )
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

        # After re-election, node-me should be coordinator (only active node)
        assert coord.is_coordinator is True


class TestCoordinatorLoop:
    @pytest.mark.asyncio
    async def test_runs_scheduling_when_coordinator(self):
        config = MagicMock()

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

        mock_client = AsyncMock()
        mock_client.get_queued_jobs.return_value = []
        peer_clients = {"node-a": mock_client}

        task = asyncio.create_task(coordinator_loop(config, coord, peer_clients))
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

        # schedule_across_mesh was called (indirectly via get_queued_jobs)
        mock_client.get_queued_jobs.assert_called()
