"""Tests for mesh peer registry."""

from __future__ import annotations

import threading
import time

from radix_edge.mesh.peer_registry import PeerRegistry, PeerState


class TestPeerRegistry:
    def test_update_and_get_peer(self):
        reg = PeerRegistry(stale_threshold=15.0)
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 2})

        peer = reg.get_peer("node-a")
        assert peer is not None
        assert peer.node_id == "node-a"
        assert peer.host == "10.0.0.1"
        assert peer.port == 8080
        assert peer.capabilities["gpu_count"] == 2

    def test_update_existing_peer(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 1})
        reg.update_peer("node-a", "10.0.0.1", 8080, {"gpu_count": 4})

        peer = reg.get_peer("node-a")
        assert peer.capabilities["gpu_count"] == 4

    def test_remove_peer(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {})
        assert reg.get_peer("node-a") is not None

        reg.remove_peer("node-a")
        assert reg.get_peer("node-a") is None

    def test_remove_nonexistent_peer(self):
        reg = PeerRegistry()
        reg.remove_peer("ghost")  # Should not raise

    def test_get_active_peers(self):
        reg = PeerRegistry(stale_threshold=15.0)
        reg.update_peer("node-a", "10.0.0.1", 8080, {})
        reg.update_peer("node-b", "10.0.0.2", 8080, {})

        active = reg.get_active_peers()
        ids = {p.node_id for p in active}
        assert ids == {"node-a", "node-b"}

    def test_stale_peer_detection(self):
        reg = PeerRegistry(stale_threshold=0.01)  # Very short threshold
        reg.update_peer("node-a", "10.0.0.1", 8080, {})

        # Immediately should be active
        assert len(reg.get_active_peers()) == 1
        assert len(reg.get_stale_peers()) == 0

        # Wait for staleness
        time.sleep(0.05)
        assert len(reg.get_active_peers()) == 0
        assert len(reg.get_stale_peers()) == 1

    def test_aggregate_gpus(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {
            "gpus": [{"name": "A100", "memory": "80 GiB"}],
            "gpu_count": 1,
        })
        reg.update_peer("node-b", "10.0.0.2", 8080, {
            "gpus": [{"name": "H100", "memory": "80 GiB"}, {"name": "H100", "memory": "80 GiB"}],
            "gpu_count": 2,
        })

        agg = reg.get_aggregate_gpus()
        assert len(agg) == 3
        node_ids = [g["node_id"] for g in agg]
        assert node_ids.count("node-a") == 1
        assert node_ids.count("node-b") == 2

    def test_set_coordinator(self):
        reg = PeerRegistry()
        reg.update_peer("node-a", "10.0.0.1", 8080, {})
        reg.update_peer("node-b", "10.0.0.2", 8080, {})

        reg.set_coordinator("node-b")

        peer_a = reg.get_peer("node-a")
        peer_b = reg.get_peer("node-b")
        assert peer_a.is_coordinator is False
        assert peer_b.is_coordinator is True

    def test_peer_count(self):
        reg = PeerRegistry()
        assert reg.peer_count == 0
        reg.update_peer("node-a", "10.0.0.1", 8080, {})
        assert reg.peer_count == 1
        reg.update_peer("node-b", "10.0.0.2", 8080, {})
        assert reg.peer_count == 2
        reg.remove_peer("node-a")
        assert reg.peer_count == 1

    def test_thread_safety(self):
        """Multiple threads updating the registry concurrently should not raise."""
        reg = PeerRegistry()
        errors = []

        def updater(node_id):
            try:
                for _ in range(100):
                    reg.update_peer(node_id, "10.0.0.1", 8080, {"gpu_count": 1})
                    reg.get_active_peers()
                    reg.get_peer(node_id)
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=updater, args=(f"node-{i}",)) for i in range(5)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert errors == []
        assert reg.peer_count == 5
