"""
Unit tests for the composition module.

Tests for AgentPipeline, ConditionalAgent, and ParallelAgents.
"""

import asyncio
from typing import Any

import pytest

from persuade.agents.base import AgentPhaseResult, TwoPhaseAgent
from persuade.agents.composition import (
    AgentPipeline,
    ComposedResult,
    ConditionalAgent,
    ParallelAgents,
)


class MockAgent(TwoPhaseAgent):
    """Mock agent for testing composition patterns."""

    def __init__(
        self,
        name: str = "MockAgent",
        plan_success: bool = True,
        execute_success: bool = True,
        plan_data: dict[str, Any] | None = None,
        execute_data: dict[str, Any] | None = None,
    ):
        self.name = name
        self.plan_success = plan_success
        self.execute_success = execute_success
        self.plan_data = plan_data or {}
        self.execute_data = execute_data or {}
        self.plan_called = False
        self.execute_called = False
        self.received_previous_results: list[AgentPhaseResult] = []

    async def plan(self, **kwargs: Any) -> AgentPhaseResult:
        self.plan_called = True
        self.received_previous_results = kwargs.get("previous_results", [])
        if self.plan_success:
            return AgentPhaseResult(
                success=True,
                data=self.plan_data,
                messages=[f"{self.name} plan completed"],
            )
        return AgentPhaseResult(
            success=False,
            error=f"{self.name} plan failed",
        )

    async def execute(self, plan_result: AgentPhaseResult, **kwargs: Any) -> AgentPhaseResult:
        self.execute_called = True
        if self.execute_success:
            return AgentPhaseResult(
                success=True,
                data=self.execute_data,
                messages=[f"{self.name} execute completed"],
            )
        return AgentPhaseResult(
            success=False,
            error=f"{self.name} execute failed",
        )


class TestComposedResult:
    """Tests for ComposedResult dataclass."""

    def test_success_result(self) -> None:
        """Test creating a successful composed result."""
        result = ComposedResult(
            overall_success=True,
            results=[
                AgentPhaseResult(success=True),
                AgentPhaseResult(success=True),
            ],
        )
        assert result.overall_success is True
        assert len(result.results) == 2
        assert result.failed_at_index is None
        assert result.error is None

    def test_failure_result(self) -> None:
        """Test creating a failed composed result."""
        result = ComposedResult(
            overall_success=False,
            results=[
                AgentPhaseResult(success=True),
                AgentPhaseResult(success=False, error="Failed"),
            ],
            failed_at_index=1,
            error="Agent 1 failed",
        )
        assert result.overall_success is False
        assert result.failed_at_index == 1
        assert result.error is not None
        assert "failed" in result.error.lower()


class TestAgentPipeline:
    """Tests for AgentPipeline."""

    def test_empty_pipeline_raises(self) -> None:
        """Test that empty pipeline raises ValueError."""
        with pytest.raises(ValueError, match="at least one agent"):
            AgentPipeline([])

    def test_pipeline_length(self) -> None:
        """Test __len__ returns number of agents."""
        pipeline = AgentPipeline([MockAgent(), MockAgent(), MockAgent()])
        assert len(pipeline) == 3

    @pytest.mark.asyncio
    async def test_single_agent_success(self) -> None:
        """Test pipeline with single successful agent."""
        agent = MockAgent(name="Single")
        pipeline = AgentPipeline([agent])

        result = await pipeline.run(auto_approve=True)

        assert result.overall_success is True
        assert len(result.results) == 1
        assert agent.plan_called is True
        assert agent.execute_called is True

    @pytest.mark.asyncio
    async def test_multiple_agents_success(self) -> None:
        """Test pipeline with multiple successful agents."""
        agents = [
            MockAgent(name="First"),
            MockAgent(name="Second"),
            MockAgent(name="Third"),
        ]
        pipeline = AgentPipeline(agents)

        result = await pipeline.run(auto_approve=True)

        assert result.overall_success is True
        assert len(result.results) == 3
        for agent in agents:
            assert agent.plan_called is True
            assert agent.execute_called is True

    @pytest.mark.asyncio
    async def test_failure_stops_pipeline(self) -> None:
        """Test that failure stops execution of remaining agents."""
        agents = [
            MockAgent(name="First"),
            MockAgent(name="Second", plan_success=False),
            MockAgent(name="Third"),
        ]
        pipeline = AgentPipeline(agents)

        result = await pipeline.run(auto_approve=True)

        assert result.overall_success is False
        assert result.failed_at_index == 1
        assert result.error is not None
        assert "Second" in result.error
        # First two agents should have been called
        assert agents[0].plan_called is True
        assert agents[1].plan_called is True
        # Third should not have been called
        assert agents[2].plan_called is False

    @pytest.mark.asyncio
    async def test_previous_results_passed(self) -> None:
        """Test that previous results are passed to subsequent agents."""
        agents = [
            MockAgent(name="First", execute_data={"from": "first"}),
            MockAgent(name="Second"),
        ]
        pipeline = AgentPipeline(agents)

        await pipeline.run(auto_approve=True)

        # Second agent should have received first agent's result (from run() = execute result)
        assert len(agents[1].received_previous_results) == 1
        assert agents[1].received_previous_results[0].data.get("from") == "first"


class TestConditionalAgent:
    """Tests for ConditionalAgent."""

    @pytest.mark.asyncio
    async def test_condition_true_executes_if_true(self) -> None:
        """Test that if_true agent executes when condition is true."""
        if_true_agent = MockAgent(name="TrueAgent")

        agent = ConditionalAgent(
            condition=lambda r: r.data.get("run_next", False),
            if_true=if_true_agent,
        )

        # Simulate previous result with condition met
        previous = [AgentPhaseResult(success=True, data={"run_next": True})]
        result = await agent.plan(previous_results=previous)

        assert result.success is True
        assert if_true_agent.plan_called is True

    @pytest.mark.asyncio
    async def test_condition_false_executes_if_false(self) -> None:
        """Test that if_false agent executes when condition is false."""
        if_true_agent = MockAgent(name="TrueAgent")
        if_false_agent = MockAgent(name="FalseAgent")

        agent = ConditionalAgent(
            condition=lambda r: r.data.get("run_next", False),
            if_true=if_true_agent,
            if_false=if_false_agent,
        )

        # Simulate previous result with condition not met
        previous = [AgentPhaseResult(success=True, data={"run_next": False})]
        result = await agent.plan(previous_results=previous)

        assert result.success is True
        assert if_true_agent.plan_called is False
        assert if_false_agent.plan_called is True

    @pytest.mark.asyncio
    async def test_condition_false_no_if_false_skips(self) -> None:
        """Test that skipping occurs when condition false and no if_false."""
        if_true_agent = MockAgent(name="TrueAgent")

        agent = ConditionalAgent(
            condition=lambda r: r.data.get("run_next", False),
            if_true=if_true_agent,
            if_false=None,
        )

        previous = [AgentPhaseResult(success=True, data={"run_next": False})]
        result = await agent.plan(previous_results=previous)

        assert result.success is True
        assert result.data.get("skipped") is True
        assert if_true_agent.plan_called is False

    @pytest.mark.asyncio
    async def test_execute_skipped_returns_success(self) -> None:
        """Test that execute on skipped agent returns success."""
        if_true_agent = MockAgent(name="TrueAgent")

        agent = ConditionalAgent(
            condition=lambda r: False,  # Always false
            if_true=if_true_agent,
        )

        plan_result = await agent.plan(previous_results=[])
        exec_result = await agent.execute(plan_result)

        assert exec_result.success is True
        assert exec_result.data.get("skipped") is True

    @pytest.mark.asyncio
    async def test_no_previous_results_evaluates_to_false(self) -> None:
        """Test that no previous results means condition evaluates to skip."""
        if_true_agent = MockAgent(name="TrueAgent")

        agent = ConditionalAgent(
            condition=lambda r: True,  # Would be true if evaluated
            if_true=if_true_agent,
        )

        # No previous results
        result = await agent.plan(previous_results=[])

        # Should skip because no previous result to evaluate
        assert result.success is True
        assert result.data.get("skipped") is True


class TestParallelAgents:
    """Tests for ParallelAgents."""

    def test_empty_parallel_raises(self) -> None:
        """Test that empty parallel raises ValueError."""
        with pytest.raises(ValueError, match="at least one agent"):
            ParallelAgents([])

    def test_parallel_length(self) -> None:
        """Test __len__ returns number of agents."""
        parallel = ParallelAgents([MockAgent(), MockAgent()])
        assert len(parallel) == 2

    @pytest.mark.asyncio
    async def test_all_success(self) -> None:
        """Test parallel execution with all successful."""
        agents = [
            MockAgent(name="A"),
            MockAgent(name="B"),
            MockAgent(name="C"),
        ]
        parallel = ParallelAgents(agents)

        result = await parallel.run(auto_approve=True)

        assert result.overall_success is True
        assert len(result.results) == 3
        for agent in agents:
            assert agent.plan_called is True

    @pytest.mark.asyncio
    async def test_one_failure(self) -> None:
        """Test parallel execution with one failure."""
        agents = [
            MockAgent(name="A"),
            MockAgent(name="B", plan_success=False),
            MockAgent(name="C"),
        ]
        parallel = ParallelAgents(agents)

        result = await parallel.run(auto_approve=True)

        assert result.overall_success is False
        assert result.failed_at_index == 1
        assert result.error is not None
        assert "B" in result.error
        # All agents should have run (parallel)
        assert len(result.results) == 3

    @pytest.mark.asyncio
    async def test_exception_handled(self) -> None:
        """Test that exceptions are converted to failed results."""

        class ExceptionAgent(TwoPhaseAgent):
            async def plan(self, **kwargs: Any) -> AgentPhaseResult:
                raise RuntimeError("Unexpected error")

            async def execute(
                self, plan_result: AgentPhaseResult, **kwargs: Any
            ) -> AgentPhaseResult:
                return AgentPhaseResult(success=True)

        agents = [MockAgent(name="A"), ExceptionAgent()]
        parallel = ParallelAgents(agents)

        result = await parallel.run(auto_approve=True)

        assert result.overall_success is False
        assert result.failed_at_index == 1
        assert result.error is not None
        assert "exception" in result.error.lower()

    @pytest.mark.asyncio
    async def test_runs_concurrently(self) -> None:
        """Test that agents actually run concurrently."""
        start_times: list[float] = []
        end_times: list[float] = []

        class TimedAgent(TwoPhaseAgent):
            async def plan(self, **kwargs: Any) -> AgentPhaseResult:
                import time

                start_times.append(time.time())
                await asyncio.sleep(0.1)  # Small delay
                end_times.append(time.time())
                return AgentPhaseResult(success=True)

            async def execute(
                self, plan_result: AgentPhaseResult, **kwargs: Any
            ) -> AgentPhaseResult:
                return AgentPhaseResult(success=True)

        parallel = ParallelAgents([TimedAgent(), TimedAgent(), TimedAgent()])

        result = await parallel.run(auto_approve=True)

        assert result.overall_success is True
        # All should have started at roughly the same time
        assert max(start_times) - min(start_times) < 0.05
