"""SequentialChain — pipe output between agents in sequence."""

from __future__ import annotations

from typing import TYPE_CHECKING

import structlog

from koan.agents.events import AgentComplete, AgentError, AgentEvent, AgentStart
from koan.agents.types import AgentResult, RunContext

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.agents.base import BaseAgent

log = structlog.get_logger()


class SequentialChain:
    """Runs agents sequentially, piping each agent's output as the next agent's input.

    Usage::

        chain = SequentialChain([researcher, writer, reviewer])
        result = await chain.run("Analyze the market trends")
        # researcher output → writer input → reviewer input
    """

    def __init__(self, agents: list[BaseAgent]) -> None:
        if not agents:
            raise ValueError("SequentialChain requires at least one agent")
        self.agents = agents

    async def run(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
        stop_on_error: bool = True,
    ) -> AgentResult:
        """Execute all agents in sequence.

        Args:
            input_text: Initial input for the first agent.
            context: Shared RunContext for all agents.
            stop_on_error: If True, stop on first agent error.

        Returns:
            AgentResult from the last agent that ran.
        """
        current_input = input_text
        total_iterations = 0
        total_tool_calls = 0
        last_result: AgentResult | None = None

        for agent in self.agents:
            log.debug("chain.step", agent=agent.name, input_length=len(current_input))

            result = await agent.run(current_input, context=context)
            total_iterations += result.iterations
            total_tool_calls += result.tool_calls_made
            last_result = result

            if result.is_error and stop_on_error:
                log.warning("chain.error", agent=agent.name, error=result.metadata.get("error"))
                return AgentResult(
                    output=result.output,
                    agent_name=f"chain({agent.name})",
                    iterations=total_iterations,
                    tool_calls_made=total_tool_calls,
                    metadata={"error": f"Agent '{agent.name}' failed", "failed_agent": agent.name},
                )

            current_input = result.output

        assert last_result is not None
        return AgentResult(
            output=last_result.output,
            agent_name=f"chain({','.join(a.name for a in self.agents)})",
            iterations=total_iterations,
            tool_calls_made=total_tool_calls,
        )

    async def run_stream(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
        stop_on_error: bool = True,
    ) -> AsyncIterator[AgentEvent]:
        """Execute all agents in sequence, streaming events from each."""
        current_input = input_text
        total_iterations = 0
        total_tool_calls = 0

        for agent in self.agents:
            yield AgentStart(agent_name=agent.name, input_text=current_input)

            agent_output = ""
            run_stream = getattr(agent, "run_stream", None)

            if run_stream is not None:
                async for event in run_stream(current_input, context=context):
                    yield event
                    if isinstance(event, AgentComplete):
                        agent_output = event.output
                        total_iterations += event.iterations
                        total_tool_calls += event.tool_calls_made
                    elif isinstance(event, AgentError):
                        total_iterations += event.iterations
                        total_tool_calls += event.tool_calls_made
                        if stop_on_error:
                            return
                        agent_output = event.output
            else:
                result = await agent.run(current_input, context=context)
                total_iterations += result.iterations
                total_tool_calls += result.tool_calls_made
                agent_output = result.output

                if result.is_error and stop_on_error:
                    yield AgentError(
                        error=result.metadata.get("error", "unknown"),
                        output=result.output,
                        iterations=total_iterations,
                        tool_calls_made=total_tool_calls,
                    )
                    return

                yield AgentComplete(
                    output=result.output,
                    iterations=result.iterations,
                    tool_calls_made=result.tool_calls_made,
                )

            current_input = agent_output
