"""Tests for Orchestrator and routing strategies."""

from __future__ import annotations

from unittest.mock import AsyncMock

import pytest

from koan.agents.orchestrator import Orchestrator
from koan.agents.router import LLMRouter, RuleRouter
from koan.agents.types import AgentResult, RunContext
from koan.types import LLMResponse, Usage


def _mock_agent(name: str, output: str) -> AsyncMock:
    agent = AsyncMock()
    agent.name = name
    agent.run = AsyncMock(
        return_value=AgentResult(
            output=output,
            agent_name=name,
            iterations=1,
            tool_calls_made=0,
        )
    )
    return agent


class TestRuleRouter:
    async def test_keyword_routing(self) -> None:
        """Routes based on keyword match."""
        router = RuleRouter({"weather": "weather_agent", "code": "code_agent"})
        agents = {"weather_agent": _mock_agent("weather_agent", ""), "code_agent": _mock_agent("code_agent", "")}

        result = await router.route("What's the weather today?", agents)
        assert result == "weather_agent"

    async def test_case_insensitive(self) -> None:
        """Keyword matching is case-insensitive."""
        router = RuleRouter({"Weather": "weather_agent"})
        agents = {"weather_agent": _mock_agent("weather_agent", "")}

        result = await router.route("WEATHER forecast please", agents)
        assert result == "weather_agent"

    async def test_no_match_returns_none(self) -> None:
        """No keyword match returns None."""
        router = RuleRouter({"weather": "weather_agent"})
        agents = {"weather_agent": _mock_agent("weather_agent", "")}

        result = await router.route("Tell me about stocks", agents)
        assert result is None

    async def test_default_on_no_match(self) -> None:
        """Falls back to default when no keyword matches."""
        router = RuleRouter({"weather": "weather_agent"}, default="general_agent")
        agents = {
            "weather_agent": _mock_agent("weather_agent", ""),
            "general_agent": _mock_agent("general_agent", ""),
        }

        result = await router.route("Tell me about stocks", agents)
        assert result == "general_agent"

    async def test_first_match_wins(self) -> None:
        """First matching keyword is used."""
        router = RuleRouter({"weather": "weather_agent", "rain": "rain_agent"})
        agents = {
            "weather_agent": _mock_agent("weather_agent", ""),
            "rain_agent": _mock_agent("rain_agent", ""),
        }

        result = await router.route("Is it going to rain? weather says yes", agents)
        # "weather" appears in the rules first
        assert result == "weather_agent"


class TestLLMRouter:
    async def test_routes_to_llm_chosen_agent(self) -> None:
        """LLM classification picks the correct agent."""
        client = AsyncMock()
        client.complete = AsyncMock(
            return_value=LLMResponse(
                text="weather_agent",
                model="test",
                provider="test",
                usage=Usage(),
            )
        )
        config = AsyncMock()

        router = LLMRouter(client=client, config=config)
        agents = {
            "weather_agent": _mock_agent("weather_agent", ""),
            "code_agent": _mock_agent("code_agent", ""),
        }

        result = await router.route("What's the weather?", agents)
        assert result == "weather_agent"

    async def test_unknown_agent_returns_none(self) -> None:
        """LLM returns unknown agent name → None."""
        client = AsyncMock()
        client.complete = AsyncMock(
            return_value=LLMResponse(
                text="nonexistent_agent",
                model="test",
                provider="test",
                usage=Usage(),
            )
        )
        config = AsyncMock()

        router = LLMRouter(client=client, config=config)
        agents = {"weather_agent": _mock_agent("weather_agent", "")}

        result = await router.route("Hello", agents)
        assert result is None


class TestOrchestrator:
    async def test_routes_to_correct_agent(self) -> None:
        """Orchestrator routes to the correct sub-agent."""
        weather = _mock_agent("weather", "Sunny, 72F")
        code = _mock_agent("code", "def hello(): ...")

        router = RuleRouter({"weather": "weather", "code": "code"})
        orch = Orchestrator(agents={"weather": weather, "code": code}, router=router)

        result = await orch.run("What's the weather?")
        assert result.output == "Sunny, 72F"
        assert "weather" in result.agent_name
        assert result.metadata["routed_to"] == "weather"

    async def test_no_route_returns_error(self) -> None:
        """No route match returns error result."""
        agent = _mock_agent("agent", "output")
        router = RuleRouter({"weather": "agent"})
        orch = Orchestrator(agents={"agent": agent}, router=router)

        result = await orch.run("Tell me about stocks")
        assert result.is_error
        assert result.metadata["error"] == "no_route_found"

    async def test_fallback_agent(self) -> None:
        """Fallback agent used when no route matches."""
        specific = _mock_agent("specific", "specific output")
        fallback = _mock_agent("fallback", "fallback output")

        router = RuleRouter({"weather": "specific"})
        orch = Orchestrator(
            agents={"specific": specific, "fallback": fallback},
            router=router,
            fallback="fallback",
        )

        result = await orch.run("Something unrelated")
        assert result.output == "fallback output"
        assert "fallback" in result.agent_name

    async def test_context_forwarded(self) -> None:
        """RunContext is forwarded to the chosen agent."""
        agent = _mock_agent("agent", "output")
        router = RuleRouter({"test": "agent"})
        orch = Orchestrator(agents={"agent": agent}, router=router)

        ctx = RunContext(user_id="u1")
        await orch.run("test input", context=ctx)

        agent.run.assert_called_once_with("test input", context=ctx)

    async def test_empty_agents_raises(self) -> None:
        """Empty agents dict raises ValueError."""
        router = RuleRouter({})
        with pytest.raises(ValueError, match="at least one agent"):
            Orchestrator(agents={}, router=router)

    async def test_handoff_preserves_result(self) -> None:
        """Agent result metadata is preserved through orchestration."""
        agent = AsyncMock()
        agent.name = "detailed"
        agent.run = AsyncMock(
            return_value=AgentResult(
                output="detailed output",
                agent_name="detailed",
                iterations=3,
                tool_calls_made=5,
                metadata={"custom": "data"},
            )
        )

        router = RuleRouter({"go": "detailed"})
        orch = Orchestrator(agents={"detailed": agent}, router=router)

        result = await orch.run("go ahead")
        assert result.output == "detailed output"
        assert result.iterations == 3
        assert result.tool_calls_made == 5
        assert result.metadata["custom"] == "data"
        assert result.metadata["routed_to"] == "detailed"
