"""Tests for the smart router."""

from __future__ import annotations

import pytest

from a2a_mesh.exceptions import NoCapableAgentError, QueueFullError
from a2a_mesh.models import (
    AgentCard,
    RoutingPolicy,
    RoutingStrategy,
    Task,
)
from a2a_mesh.registry import AgentRegistry
from a2a_mesh.router import Router


class TestRouter:
    """Tests for Router."""

    def test_route_by_agent_name(self, router: Router) -> None:
        task = Task(name="t", agent="research-agent")
        agent = router.route(task)
        assert agent.card.name == "research-agent"

    def test_route_by_capability(self, router: Router) -> None:
        task = Task(name="t", required_capabilities=["web_search"])
        agent = router.route(task)
        assert agent.card.name == "research-agent"

    def test_route_no_match_raises(self, router: Router) -> None:
        task = Task(name="t", required_capabilities=["quantum_physics"])
        with pytest.raises(NoCapableAgentError):
            router.route(task)

    def test_route_unknown_agent_raises(self, router: Router) -> None:
        task = Task(name="t", agent="nonexistent-agent")
        with pytest.raises(NoCapableAgentError):
            router.route(task)

    def test_least_cost_strategy(self, populated_registry: AgentRegistry) -> None:
        router = Router(
            populated_registry,
            policy=RoutingPolicy(strategy=RoutingStrategy.LEAST_COST),
        )
        task = Task(name="t", required_capabilities=["summarization"])
        agent = router.route(task)
        # writing-agent has cost=0.01, research=0.02, analysis=0.05
        assert agent.card.name == "writing-agent"

    def test_least_load_strategy(self, populated_registry: AgentRegistry) -> None:
        router = Router(
            populated_registry,
            policy=RoutingPolicy(strategy=RoutingStrategy.LEAST_LOAD),
        )
        populated_registry.agents["writing-agent"].current_load = 5
        populated_registry.agents["research-agent"].current_load = 1
        populated_registry.agents["analysis-agent"].current_load = 3
        task = Task(name="t", required_capabilities=["summarization"])
        agent = router.route(task)
        assert agent.card.name == "research-agent"

    def test_least_latency_strategy(self, populated_registry: AgentRegistry) -> None:
        router = Router(
            populated_registry,
            policy=RoutingPolicy(strategy=RoutingStrategy.LEAST_LATENCY),
        )
        populated_registry.agents["writing-agent"].avg_latency_ms = 500
        populated_registry.agents["research-agent"].avg_latency_ms = 100
        populated_registry.agents["analysis-agent"].avg_latency_ms = 200
        task = Task(name="t", required_capabilities=["summarization"])
        agent = router.route(task)
        assert agent.card.name == "research-agent"

    def test_round_robin_strategy(self, populated_registry: AgentRegistry) -> None:
        router = Router(
            populated_registry,
            policy=RoutingPolicy(strategy=RoutingStrategy.ROUND_ROBIN),
        )
        task = Task(name="t", required_capabilities=["summarization"])

        first = router.route(task)
        second = router.route(task)
        third = router.route(task)

        names = {first.card.name, second.card.name, third.card.name}
        # After 3 rounds, we should have cycled through the agents
        assert len(names) >= 2  # at least 2 different agents

    def test_queue_full_raises(self, populated_registry: AgentRegistry) -> None:
        router = Router(
            populated_registry,
            policy=RoutingPolicy(
                strategy=RoutingStrategy.LEAST_LOAD,
                max_queue_depth=2,
            ),
        )
        # Fill up all agents
        for agent in populated_registry.agents.values():
            agent.current_load = 2
        task = Task(name="t", required_capabilities=["summarization"])
        with pytest.raises(QueueFullError):
            router.route(task)

    def test_route_multi(self, router: Router) -> None:
        task = Task(name="t", required_capabilities=["summarization"])
        agents = router.route_multi(task, count=2)
        assert len(agents) == 2

    def test_route_multi_limited_by_available(self, router: Router) -> None:
        task = Task(name="t", required_capabilities=["web_search"])
        agents = router.route_multi(task, count=5)
        # Only research-agent has web_search
        assert len(agents) == 1
