"""Tests for ParallelFanOut."""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock

import pytest

from koan.agents.parallel import ParallelFanOut
from koan.agents.types import AgentResult, RunContext


def _mock_agent(
    name: str,
    output: str,
    iterations: int = 1,
    tool_calls: int = 0,
    delay: float = 0,
    error: Exception | None = None,
) -> AsyncMock:
    agent = AsyncMock()
    agent.name = name

    async def _run(input_text, *, context=None):
        if delay:
            await asyncio.sleep(delay)
        if error:
            raise error
        return AgentResult(
            output=output,
            agent_name=name,
            iterations=iterations,
            tool_calls_made=tool_calls,
        )

    agent.run = AsyncMock(side_effect=_run)
    return agent


class TestParallelFanOut:
    async def test_three_agents_run_concurrently(self) -> None:
        """All agents receive the same input and results are collected."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "output-B")
        c = _mock_agent("C", "output-C")

        fan_out = ParallelFanOut([a, b, c])
        result = await fan_out.run("start")

        assert "output-A" in result.output
        assert "output-B" in result.output
        assert "output-C" in result.output
        assert result.iterations == 3
        assert not result.is_error

        # All agents called with same input
        a.run.assert_called_once_with("start", context=None)
        b.run.assert_called_once_with("start", context=None)
        c.run.assert_called_once_with("start", context=None)

    async def test_one_agent_fails_partial_results(self) -> None:
        """One failing agent doesn't prevent others from completing."""
        a = _mock_agent("A", "output-A")
        b = _mock_agent("B", "", error=RuntimeError("boom"))
        c = _mock_agent("C", "output-C")

        fan_out = ParallelFanOut([a, b, c])
        result = await fan_out.run("start")

        assert "output-A" in result.output
        assert "output-C" in result.output
        assert "errors" in result.metadata
        assert len(result.metadata["errors"]) == 1
        assert "boom" in result.metadata["errors"][0]

    async def test_timeout_on_slow_agent(self) -> None:
        """Slow agent is cancelled after timeout."""
        a = _mock_agent("A", "fast", delay=0)
        b = _mock_agent("B", "slow", delay=10)

        fan_out = ParallelFanOut([a, b], timeout=0.1)
        result = await fan_out.run("start")

        assert "fast" in result.output
        assert "errors" in result.metadata

    async def test_single_agent(self) -> None:
        """Works with a single agent."""
        a = _mock_agent("solo", "solo-output", iterations=2, tool_calls=3)

        fan_out = ParallelFanOut([a])
        result = await fan_out.run("input")

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

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

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

        ctx = RunContext(user_id="u1")
        fan_out = ParallelFanOut([a, b])
        await fan_out.run("start", context=ctx)

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

    async def test_aggregates_iterations_and_tool_calls(self) -> None:
        """Iterations and tool calls summed 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)

        fan_out = ParallelFanOut([a, b])
        result = await fan_out.run("start")

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

    async def test_metadata_contains_individual_results(self) -> None:
        """Metadata includes individual AgentResult objects."""
        a = _mock_agent("A", "out-A")
        b = _mock_agent("B", "out-B")

        fan_out = ParallelFanOut([a, b])
        result = await fan_out.run("start")

        assert "results" in result.metadata
        assert len(result.metadata["results"]) == 2
