"""ParallelFanOut — run agents concurrently and collect results."""

from __future__ import annotations

import asyncio
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 ParallelFanOut:
    """Runs multiple agents concurrently on the same input.

    All agents receive the same input text and run in parallel via
    ``asyncio.gather``. Results are collected and merged.

    Usage::

        fan_out = ParallelFanOut([summarizer, classifier, extractor])
        result = await fan_out.run("Analyze this document...")
        # result.output contains combined outputs
        # result.metadata["results"] has individual AgentResults
    """

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

    async def run(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AgentResult:
        """Execute all agents concurrently on the same input.

        Args:
            input_text: Input for all agents.
            context: Shared RunContext for all agents.

        Returns:
            Merged AgentResult with combined output and metadata.
        """
        tasks = [agent.run(input_text, context=context) for agent in self.agents]

        if self.timeout is not None:
            results = await asyncio.gather(
                *[asyncio.wait_for(t, timeout=self.timeout) for t in tasks],
                return_exceptions=True,
            )
        else:
            results = await asyncio.gather(*tasks, return_exceptions=True)

        agent_results: list[AgentResult] = []
        errors: list[str] = []
        total_iterations = 0
        total_tool_calls = 0
        outputs: list[str] = []

        for agent, result in zip(self.agents, results, strict=True):
            if isinstance(result, BaseException):
                log.warning("parallel.error", agent=agent.name, error=str(result))
                errors.append(f"{agent.name}: {result}")
                agent_results.append(
                    AgentResult(
                        output="",
                        agent_name=agent.name,
                        metadata={"error": str(result)},
                    )
                )
            else:
                agent_results.append(result)
                total_iterations += result.iterations
                total_tool_calls += result.tool_calls_made
                outputs.append(result.output)

        metadata: dict[str, object] = {
            "results": agent_results,
        }
        if errors:
            metadata["errors"] = errors

        return AgentResult(
            output="\n\n".join(outputs),
            agent_name=f"parallel({','.join(a.name for a in self.agents)})",
            iterations=total_iterations,
            tool_calls_made=total_tool_calls,
            metadata=metadata,
        )

    async def run_stream(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Execute all agents concurrently, yielding events from each.

        Events from different agents will interleave. Each agent's events
        are preceded by an ``AgentStart`` event.
        """
        queue: asyncio.Queue[AgentEvent | None] = asyncio.Queue()
        agent_count = len(self.agents)

        async def _run_agent(agent: BaseAgent) -> None:
            await queue.put(AgentStart(agent_name=agent.name, input_text=input_text))
            run_stream = getattr(agent, "run_stream", None)
            try:
                if run_stream is not None:
                    async for event in run_stream(input_text, context=context):
                        await queue.put(event)
                else:
                    result = await agent.run(input_text, context=context)
                    await queue.put(
                        AgentComplete(
                            output=result.output,
                            iterations=result.iterations,
                            tool_calls_made=result.tool_calls_made,
                        )
                    )
            except Exception as e:
                log.warning("parallel.stream_error", agent=agent.name, error=str(e))
                await queue.put(AgentError(error=str(e)))
            finally:
                await queue.put(None)  # Sentinel

        tasks = [asyncio.create_task(_run_agent(a)) for a in self.agents]
        finished = 0

        while finished < agent_count:
            event = await queue.get()
            if event is None:
                finished += 1
            else:
                yield event

        # Ensure all tasks are done
        await asyncio.gather(*tasks, return_exceptions=True)
