"""
Agent composition patterns.

Provides patterns for composing multiple agents together:
- AgentPipeline: Execute agents in sequence
- ConditionalAgent: Execute based on conditions

These patterns enable building complex workflows from simple agent components.
"""

from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from typing import Any

from persuade.agents.base import AgentPhaseResult, TwoPhaseAgent

# Type alias for agent sequences - use Sequence for covariance
# (allows list[SubclassAgent] where Sequence[TwoPhaseAgent] is expected)
AgentSequence = Sequence[TwoPhaseAgent]


@dataclass
class ComposedResult:
    """Result from a composed agent execution.

    Aggregates results from multiple agents in a pipeline.

    Attributes:
        overall_success: True if all agents succeeded
        results: List of results from each agent in order
        failed_at_index: Index of first failed agent, or None if all succeeded
        error: Error message from the first failure
    """

    overall_success: bool
    results: list[AgentPhaseResult] = field(default_factory=list)
    failed_at_index: int | None = None
    error: str | None = None


class AgentPipeline:
    """Execute multiple agents in sequence, passing results forward.

    Each agent receives `previous_results` kwarg containing all prior results.
    Execution stops on first failure.

    Example:
        pipeline = AgentPipeline([
            FetchFeedbackAgent(pr_number=67),
            ClassifyFeedbackAgent(),
            ApplyFixesAgent(),
        ])
        result = await pipeline.run(auto_approve=True)

        if result.overall_success:
            print("All agents completed successfully")
        else:
            print(f"Failed at agent {result.failed_at_index}: {result.error}")
    """

    def __init__(self, agents: AgentSequence):
        """Initialize pipeline with agents to execute in order.

        Args:
            agents: Sequence of agents to execute sequentially
        """
        if not agents:
            raise ValueError("AgentPipeline requires at least one agent")
        self.agents = list(agents)  # Convert to list for internal use

    async def run(self, auto_approve: bool = False, **kwargs: Any) -> ComposedResult:
        """Run all agents in sequence.

        Each agent's result is available to subsequent agents via the
        `previous_results` kwarg passed to plan() and execute().

        Args:
            auto_approve: If True, auto-approve all agents
            **kwargs: Additional arguments passed to all agents

        Returns:
            ComposedResult with all individual results
        """
        results: list[AgentPhaseResult] = []

        for i, agent in enumerate(self.agents):
            # Pass a copy of previous results to prevent mutation issues
            result = await agent.run(
                auto_approve=auto_approve,
                previous_results=list(results),
                **kwargs,
            )
            results.append(result)

            if not result.success:
                return ComposedResult(
                    overall_success=False,
                    results=results,
                    failed_at_index=i,
                    error=f"Agent {i} ({agent.__class__.__name__}) failed: {result.error}",
                )

        return ComposedResult(overall_success=True, results=results)

    def __len__(self) -> int:
        """Return number of agents in pipeline."""
        return len(self.agents)


class ConditionalAgent(TwoPhaseAgent):
    """Agent that conditionally executes based on a predicate.

    Evaluates a condition against the previous agent's result to decide
    which agent to execute. Useful for building branching workflows.

    Example:
        # Only apply fixes if there are items to fix
        agent = ConditionalAgent(
            condition=lambda r: len(r.data.get("items", [])) > 0,
            if_true=ApplyFixesAgent(),
            if_false=None,  # Skip if no items
        )

        # In a pipeline
        pipeline = AgentPipeline([
            FetchFeedbackAgent(pr_number=67),
            ConditionalAgent(
                condition=lambda r: r.data.get("needs_iteration", False),
                if_true=IteratePRAgent(pr_number=67),
            ),
        ])
    """

    def __init__(
        self,
        condition: Callable[[AgentPhaseResult], bool],
        if_true: TwoPhaseAgent,
        if_false: TwoPhaseAgent | None = None,
    ):
        """Initialize conditional agent.

        Args:
            condition: Function that takes the previous result and returns bool
            if_true: Agent to execute when condition is True
            if_false: Agent to execute when condition is False, or None to skip
        """
        self.condition = condition
        self.if_true = if_true
        self.if_false = if_false
        self._selected_agent: TwoPhaseAgent | None = None

    async def plan(self, **kwargs: Any) -> AgentPhaseResult:
        """Plan phase: evaluate condition and delegate to selected agent.

        The condition is evaluated against the last result in `previous_results`.
        If no previous results exist and if_false is None, returns success with skip message.
        """
        previous: list[AgentPhaseResult] = kwargs.get("previous_results", [])
        last_result = previous[-1] if previous else None

        # Evaluate condition
        if last_result is not None and self.condition(last_result):
            self._selected_agent = self.if_true
        else:
            self._selected_agent = self.if_false

        # If no agent selected (condition false and no if_false), skip
        if self._selected_agent is None:
            return AgentPhaseResult(
                success=True,
                messages=["Condition not met, skipping agent"],
                data={"skipped": True},
            )

        return await self._selected_agent.plan(**kwargs)

    async def execute(self, plan_result: AgentPhaseResult, **kwargs: Any) -> AgentPhaseResult:
        """Execute phase: delegate to selected agent or skip.

        If agent was skipped in plan phase, returns success.
        """
        if self._selected_agent is None:
            return AgentPhaseResult(
                success=True,
                messages=["Skipped (condition not met)"],
                data={"skipped": True},
            )
        return await self._selected_agent.execute(plan_result, **kwargs)


class ParallelAgents:
    """Execute multiple agents in parallel.

    All agents run concurrently and results are collected.
    If any agent fails, the overall result is a failure.

    Example:
        parallel = ParallelAgents([
            LintAgent(),
            TypeCheckAgent(),
            TestAgent(),
        ])
        result = await parallel.run(auto_approve=True)

    Note:
        Use with caution - parallel agents should not modify the same files.
    """

    def __init__(self, agents: AgentSequence):
        """Initialize with agents to run in parallel.

        Args:
            agents: Sequence of agents to execute concurrently
        """
        if not agents:
            raise ValueError("ParallelAgents requires at least one agent")
        self.agents = list(agents)  # Convert to list for internal use

    async def run(self, auto_approve: bool = False, **kwargs: Any) -> ComposedResult:
        """Run all agents in parallel.

        Args:
            auto_approve: If True, auto-approve all agents
            **kwargs: Additional arguments passed to all agents

        Returns:
            ComposedResult with all individual results
        """
        import asyncio

        # Run all agents concurrently
        tasks = [agent.run(auto_approve=auto_approve, **kwargs) for agent in self.agents]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results
        agent_results: list[AgentPhaseResult] = []
        failed_index: int | None = None
        error_msg: str | None = None

        for i, result in enumerate(results):
            if isinstance(result, BaseException):
                # Convert exception to failed result
                agent_result = AgentPhaseResult(
                    success=False,
                    error=f"Agent raised exception: {result}",
                )
                agent_results.append(agent_result)
                if failed_index is None:
                    failed_index = i
                    error_msg = (
                        f"Agent {i} ({self.agents[i].__class__.__name__}) "
                        f"raised exception: {result}"
                    )
            else:
                # Type narrowing: result is AgentPhaseResult
                agent_result = result
                agent_results.append(agent_result)
                if not agent_result.success and failed_index is None:
                    failed_index = i
                    error_msg = (
                        f"Agent {i} ({self.agents[i].__class__.__name__}) "
                        f"failed: {agent_result.error}"
                    )

        return ComposedResult(
            overall_success=failed_index is None,
            results=agent_results,
            failed_at_index=failed_index,
            error=error_msg,
        )

    def __len__(self) -> int:
        """Return number of agents."""
        return len(self.agents)
