"""Tests for the agent registry."""

from __future__ import annotations

import pytest

from a2a_mesh.exceptions import AgentAlreadyRegisteredError, AgentNotFoundError
from a2a_mesh.models import AgentCard, AgentStatus
from a2a_mesh.registry import AgentRegistry


class TestAgentRegistry:
    """Tests for AgentRegistry."""

    def test_register_agent(self, registry: AgentRegistry, research_card: AgentCard) -> None:
        agent = registry.register(research_card)
        assert agent.card.name == "research-agent"
        assert agent.status == AgentStatus.UNKNOWN
        assert "research-agent" in registry.agents

    def test_register_duplicate_raises(
        self, registry: AgentRegistry, research_card: AgentCard
    ) -> None:
        registry.register(research_card)
        with pytest.raises(AgentAlreadyRegisteredError):
            registry.register(research_card)

    def test_register_duplicate_with_force(
        self, registry: AgentRegistry, research_card: AgentCard
    ) -> None:
        agent1 = registry.register(research_card)
        agent2 = registry.register(research_card, force=True)
        assert agent1.agent_id != agent2.agent_id

    def test_deregister_agent(
        self, registry: AgentRegistry, research_card: AgentCard
    ) -> None:
        registry.register(research_card)
        registry.deregister("research-agent")
        assert "research-agent" not in registry.agents

    def test_deregister_unknown_raises(self, registry: AgentRegistry) -> None:
        with pytest.raises(AgentNotFoundError):
            registry.deregister("nonexistent")

    def test_get_agent(
        self, registry: AgentRegistry, research_card: AgentCard
    ) -> None:
        registry.register(research_card)
        agent = registry.get("research-agent")
        assert agent.card.name == "research-agent"

    def test_get_unknown_raises(self, registry: AgentRegistry) -> None:
        with pytest.raises(AgentNotFoundError):
            registry.get("nonexistent")

    def test_list_agents(self, populated_registry: AgentRegistry) -> None:
        agents = populated_registry.list_agents()
        assert len(agents) == 3
        names = {a.card.name for a in agents}
        assert "research-agent" in names
        assert "analysis-agent" in names
        assert "writing-agent" in names

    def test_find_by_capability(self, populated_registry: AgentRegistry) -> None:
        matches = populated_registry.find_by_capability(
            ["summarization"], healthy_only=False
        )
        # research, analysis, and writing agents all have summarization
        assert len(matches) == 3

    def test_find_by_specific_capability(self, populated_registry: AgentRegistry) -> None:
        matches = populated_registry.find_by_capability(
            ["web_search"], healthy_only=False
        )
        assert len(matches) == 1
        assert matches[0].card.name == "research-agent"

    def test_find_by_multiple_capabilities(self, populated_registry: AgentRegistry) -> None:
        matches = populated_registry.find_by_capability(
            ["data_analysis", "financial_analysis"], healthy_only=False
        )
        assert len(matches) == 1
        assert matches[0].card.name == "analysis-agent"

    def test_find_no_match(self, populated_registry: AgentRegistry) -> None:
        matches = populated_registry.find_by_capability(
            ["quantum_computing"], healthy_only=False
        )
        assert len(matches) == 0

    def test_find_excludes_unhealthy(self, populated_registry: AgentRegistry) -> None:
        populated_registry.agents["research-agent"].status = AgentStatus.UNHEALTHY
        matches = populated_registry.find_by_capability(
            ["web_search"], healthy_only=True
        )
        assert len(matches) == 0

    def test_find_sorted_by_load(self, populated_registry: AgentRegistry) -> None:
        populated_registry.agents["writing-agent"].current_load = 5
        populated_registry.agents["research-agent"].current_load = 1
        populated_registry.agents["analysis-agent"].current_load = 3
        matches = populated_registry.find_by_capability(
            ["summarization"], healthy_only=False
        )
        loads = [m.current_load for m in matches]
        assert loads == sorted(loads)

    @pytest.mark.asyncio
    async def test_health_check_no_url(self, registry: AgentRegistry) -> None:
        card = AgentCard(name="local-agent")
        registry.register(card)
        status = await registry.check_health("local-agent")
        assert status == AgentStatus.UNKNOWN
