"""Orchestrator — supervisor agent that delegates to sub-agents."""

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
    from koan.agents.router import Router

log = structlog.get_logger()


class Orchestrator:
    """Routes input to the best sub-agent using a configurable strategy.

    Usage::

        orchestrator = Orchestrator(
            agents={"weather": weather_agent, "code": code_agent},
            router=RuleRouter({"weather": "weather", "code": "code"}),
        )
        result = await orchestrator.run("What's the weather in NYC?")
    """

    def __init__(
        self,
        agents: dict[str, BaseAgent],
        router: Router,
        *,
        fallback: str | None = None,
    ) -> None:
        if not agents:
            raise ValueError("Orchestrator requires at least one agent")
        self.agents = agents
        self.router = router
        self.fallback = fallback

    async def run(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AgentResult:
        """Route input to the best agent and execute.

        Args:
            input_text: User input to route.
            context: Shared RunContext.

        Returns:
            AgentResult from the chosen agent, or error if no route found.
        """
        chosen = await self.router.route(input_text, self.agents, context=context)

        if chosen is None and self.fallback and self.fallback in self.agents:
            chosen = self.fallback

        if chosen is None or chosen not in self.agents:
            log.warning("orchestrator.no_route", input_length=len(input_text))
            return AgentResult(
                output="",
                agent_name="orchestrator",
                metadata={"error": "no_route_found"},
            )

        agent = self.agents[chosen]
        log.debug("orchestrator.route", chosen=chosen, agent=agent.name)

        result = await agent.run(input_text, context=context)

        return AgentResult(
            output=result.output,
            agent_name=f"orchestrator({chosen})",
            iterations=result.iterations,
            tool_calls_made=result.tool_calls_made,
            metadata={**result.metadata, "routed_to": chosen},
        )

    async def run_stream(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Route input to the best agent and stream its events."""
        chosen = await self.router.route(input_text, self.agents, context=context)

        if chosen is None and self.fallback and self.fallback in self.agents:
            chosen = self.fallback

        if chosen is None or chosen not in self.agents:
            log.warning("orchestrator.no_route", input_length=len(input_text))
            yield AgentError(error="no_route_found")
            return

        agent = self.agents[chosen]
        yield AgentStart(agent_name=chosen, input_text=input_text)

        run_stream = getattr(agent, "run_stream", None)
        if run_stream is not None:
            async for event in run_stream(input_text, context=context):
                yield event
        else:
            result = await agent.run(input_text, context=context)
            if result.is_error:
                yield AgentError(
                    error=result.metadata.get("error", "unknown"),
                    output=result.output,
                    iterations=result.iterations,
                    tool_calls_made=result.tool_calls_made,
                )
            else:
                yield AgentComplete(
                    output=result.output,
                    iterations=result.iterations,
                    tool_calls_made=result.tool_calls_made,
                    metadata={"routed_to": chosen},
                )
