"""Tests for SequentialChain."""

from __future__ import annotations

from unittest.mock import AsyncMock

import pytest

from koan.agents.chain import SequentialChain
from koan.agents.types import AgentResult, RunContext


def _mock_agent(name: str, output: str, iterations: int = 1, tool_calls: int = 0, error: bool = False) -> AsyncMock:
    """Create a mock agent with a preset result."""
    agent = AsyncMock()
    agent.name = name
    metadata = {"error": f"{name} failed"} if error else {}
    agent.run = AsyncMock(
        return_value=AgentResult(
            output=output,
            agent_name=name,
            iterations=iterations,
            tool_calls_made=tool_calls,
            metadata=metadata,
        )
    )
    return agent


class TestSequentialChain:
    async def test_three_agent_chain(self) -> None:
        """Output flows A -> B -> C."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "output-B")
        c = _mock_agent("C", "output-C")

        chain = SequentialChain([a, b, c])
        result = await chain.run("start")

        assert result.output == "output-C"
        assert "A" in result.agent_name
        assert "B" in result.agent_name
        assert "C" in result.agent_name
        assert result.iterations == 3
        assert result.tool_calls_made == 0

        # Verify input piping
        a.run.assert_called_once_with("start", context=None)
        b.run.assert_called_once_with("output-A", context=None)
        c.run.assert_called_once_with("output-B", context=None)

    async def test_failure_stops_chain(self) -> None:
        """Error in middle agent stops the chain."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "error-output", error=True)
        c = _mock_agent("C", "output-C")

        chain = SequentialChain([a, b, c])
        result = await chain.run("start")

        assert result.is_error
        assert "B" in result.metadata.get("failed_agent", "")
        # C should never have been called
        c.run.assert_not_called()

    async def test_failure_continues_when_stop_on_error_false(self) -> None:
        """Chain continues past errors when stop_on_error=False."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "error-output", error=True)
        c = _mock_agent("C", "output-C")

        chain = SequentialChain([a, b, c])
        result = await chain.run("start", stop_on_error=False)

        assert result.output == "output-C"
        c.run.assert_called_once()

    async def test_single_agent_chain(self) -> None:
        """Chain with one agent works correctly."""
        a = _mock_agent("solo", "solo-output", iterations=2, tool_calls=3)

        chain = SequentialChain([a])
        result = await chain.run("input")

        assert result.output == "solo-output"
        assert result.iterations == 2
        assert result.tool_calls_made == 3

    async def test_empty_chain_raises(self) -> None:
        """Empty chain raises ValueError."""
        with pytest.raises(ValueError, match="at least one agent"):
            SequentialChain([])

    async def test_context_passed_to_all_agents(self) -> None:
        """RunContext is forwarded to every agent."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "output-B")

        ctx = RunContext(user_id="u1", thread_id="t1")
        chain = SequentialChain([a, b])
        await chain.run("start", context=ctx)

        a.run.assert_called_once_with("start", context=ctx)
        b.run.assert_called_once_with("output-A", context=ctx)

    async def test_aggregates_iterations_and_tool_calls(self) -> None:
        """Chain sums iterations and tool calls across all agents."""
        a = _mock_agent("A", "out-A", iterations=2, tool_calls=1)
        b = _mock_agent("B", "out-B", iterations=3, tool_calls=5)

        chain = SequentialChain([a, b])
        result = await chain.run("start")

        assert result.iterations == 5
        assert result.tool_calls_made == 6
