"""Tests for streaming on SequentialChain, ParallelFanOut, and Orchestrator."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

from koan.agents.chain import SequentialChain
from koan.agents.events import AgentComplete, AgentError, AgentStart
from koan.agents.orchestrator import Orchestrator
from koan.agents.parallel import ParallelFanOut
from koan.agents.router import RuleRouter
from koan.agents.types import AgentResult, RunContext


def _mock_agent(name: str, output: str, streaming: bool = False) -> MagicMock:
    """Create a mock agent. If streaming=True, adds run_stream()."""
    agent = MagicMock()
    agent.name = name
    agent.run = AsyncMock(
        return_value=AgentResult(
            output=output,
            agent_name=name,
            iterations=1,
            tool_calls_made=0,
        )
    )

    if streaming:

        async def _run_stream(input_text, *, context=None):
            yield AgentComplete(output=output, iterations=1, tool_calls_made=0)

        agent.run_stream = _run_stream
    else:
        agent.run_stream = None

    return agent


def _mock_error_agent(name: str) -> MagicMock:
    agent = MagicMock()
    agent.name = name
    agent.run = AsyncMock(
        return_value=AgentResult(
            output="",
            agent_name=name,
            iterations=1,
            tool_calls_made=0,
            metadata={"error": f"{name} failed"},
        )
    )
    agent.run_stream = None
    return agent


class TestSequentialChainStream:
    async def test_streams_events_from_agents(self) -> None:
        """Chain streams AgentStart + AgentComplete for each agent."""
        a = _mock_agent("A", "out-A")
        b = _mock_agent("B", "out-B")

        chain = SequentialChain([a, b])
        events = [e async for e in chain.run_stream("start")]

        starts = [e for e in events if isinstance(e, AgentStart)]
        completes = [e for e in events if isinstance(e, AgentComplete)]

        assert len(starts) == 2
        assert starts[0].agent_name == "A"
        assert starts[1].agent_name == "B"
        assert len(completes) == 2

    async def test_pipes_output_between_agents(self) -> None:
        """Second agent receives first agent's output as input."""
        a = _mock_agent("A", "from-A")
        b = _mock_agent("B", "from-B")

        chain = SequentialChain([a, b])
        _ = [e async for e in chain.run_stream("start")]

        b.run.assert_called_once_with("from-A", context=None)

    async def test_stops_on_error(self) -> None:
        """Chain stops streaming on error when stop_on_error=True."""
        a = _mock_agent("A", "out-A")
        b = _mock_error_agent("B")
        c = _mock_agent("C", "out-C")

        chain = SequentialChain([a, b, c])
        events = [e async for e in chain.run_stream("start")]

        errors = [e for e in events if isinstance(e, AgentError)]
        assert len(errors) == 1
        # C should not have started
        c_starts = [e for e in events if isinstance(e, AgentStart) and e.agent_name == "C"]
        assert len(c_starts) == 0

    async def test_delegates_to_streaming_agent(self) -> None:
        """If sub-agent has run_stream(), chain uses it."""
        a = _mock_agent("A", "out-A", streaming=True)

        chain = SequentialChain([a])
        events = [e async for e in chain.run_stream("start")]

        completes = [e for e in events if isinstance(e, AgentComplete)]
        assert len(completes) == 1
        assert completes[0].output == "out-A"

    async def test_context_forwarded(self) -> None:
        """RunContext is forwarded to agents."""
        a = _mock_agent("A", "out-A")

        ctx = RunContext(user_id="u1")
        chain = SequentialChain([a])
        _ = [e async for e in chain.run_stream("start", context=ctx)]

        a.run.assert_called_once_with("start", context=ctx)


class TestParallelFanOutStream:
    async def test_streams_events_from_all_agents(self) -> None:
        """All agents produce events."""
        a = _mock_agent("A", "out-A")
        b = _mock_agent("B", "out-B")

        fan_out = ParallelFanOut([a, b])
        events = [e async for e in fan_out.run_stream("start")]

        starts = [e for e in events if isinstance(e, AgentStart)]
        completes = [e for e in events if isinstance(e, AgentComplete)]

        assert len(starts) == 2
        assert len(completes) == 2
        start_names = {e.agent_name for e in starts}
        assert start_names == {"A", "B"}

    async def test_handles_agent_error(self) -> None:
        """Error from an agent yields AgentError event."""
        a = _mock_agent("A", "out-A")

        b = MagicMock()
        b.name = "B"
        b.run = AsyncMock(side_effect=RuntimeError("boom"))
        b.run_stream = None

        fan_out = ParallelFanOut([a, b])
        events = [e async for e in fan_out.run_stream("start")]

        errors = [e for e in events if isinstance(e, AgentError)]
        assert len(errors) == 1
        assert "boom" in errors[0].error

    async def test_delegates_to_streaming_agent(self) -> None:
        """Uses run_stream() on agents that support it."""
        a = _mock_agent("A", "out-A", streaming=True)

        fan_out = ParallelFanOut([a])
        events = [e async for e in fan_out.run_stream("start")]

        completes = [e for e in events if isinstance(e, AgentComplete)]
        assert len(completes) == 1


class TestOrchestratorStream:
    async def test_streams_routed_agent(self) -> None:
        """Orchestrator streams events from the routed agent."""
        weather = _mock_agent("weather", "Sunny")
        router = RuleRouter({"weather": "weather"})
        orch = Orchestrator(agents={"weather": weather}, router=router)

        events = [e async for e in orch.run_stream("weather forecast")]

        starts = [e for e in events if isinstance(e, AgentStart)]
        completes = [e for e in events if isinstance(e, AgentComplete)]

        assert len(starts) == 1
        assert starts[0].agent_name == "weather"
        assert len(completes) == 1
        assert completes[0].output == "Sunny"

    async def test_no_route_yields_error(self) -> None:
        """No route match yields AgentError."""
        agent = _mock_agent("agent", "output")
        router = RuleRouter({"weather": "agent"})
        orch = Orchestrator(agents={"agent": agent}, router=router)

        events = [e async for e in orch.run_stream("stocks")]

        errors = [e for e in events if isinstance(e, AgentError)]
        assert len(errors) == 1
        assert errors[0].error == "no_route_found"

    async def test_delegates_to_streaming_agent(self) -> None:
        """Uses run_stream() on the routed agent if available."""
        agent = _mock_agent("weather", "Sunny", streaming=True)
        router = RuleRouter({"weather": "weather"})
        orch = Orchestrator(agents={"weather": agent}, router=router)

        events = [e async for e in orch.run_stream("weather please")]

        completes = [e for e in events if isinstance(e, AgentComplete)]
        assert len(completes) == 1

    async def test_fallback_streams(self) -> None:
        """Fallback agent is used and streams when no route matches."""
        fallback = _mock_agent("fallback", "fallback output")
        router = RuleRouter({"weather": "specific"})
        orch = Orchestrator(
            agents={"specific": _mock_agent("specific", ""), "fallback": fallback},
            router=router,
            fallback="fallback",
        )

        events = [e async for e in orch.run_stream("something else")]

        starts = [e for e in events if isinstance(e, AgentStart)]
        assert len(starts) == 1
        assert starts[0].agent_name == "fallback"
