"""Supervisor — LLM-powered orchestrator that decomposes, delegates, and synthesizes."""

from __future__ import annotations

from typing import TYPE_CHECKING

import orjson
import structlog

from koan.agents.events import (
    AgentComplete,
    AgentError,
    AgentEvent,
    AgentStart,
    TextDelta,
)
from koan.agents.types import AgentResult, RunContext
from koan.types import Message

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.agents.base import BaseAgent
    from koan.agents.router import Router
    from koan.llm.client import LLMClient
    from koan.llm.config import LLMConfig

log = structlog.get_logger()

_DECOMPOSE_SYSTEM = (
    "You are a task planner. Given a user request and a list of available specialist agents, "
    "decompose the request into independent sub-tasks.\n\n"
    "Available agents:\n{agent_list}\n\n"
    "Rules:\n"
    "- If the request is a single task, return a one-element array.\n"
    "- Each sub-task must be a self-contained question or instruction.\n"
    "- Output ONLY a JSON array of strings. No other text."
)

_SYNTHESIZE_SYSTEM = (
    "You are a helpful assistant. The user asked a question and multiple specialist agents "
    "produced answers for different parts of the request. Combine the results into a single, "
    "coherent response. Be concise and natural."
)


class Supervisor:
    """LLM-powered orchestrator that decomposes, delegates, and synthesizes.

    Unlike the simple ``Orchestrator`` which routes to a single agent,
    the ``Supervisor`` can handle compound queries by:

    1. **Decomposing** the input into independent sub-tasks using an LLM
    2. **Delegating** each sub-task to the best agent via the router
    3. **Synthesizing** a unified response from all sub-agent results

    Usage::

        supervisor = Supervisor(
            agents={"weather": weather_agent, "math": math_agent},
            router=RuleRouter({"weather": "weather", "add": "math"}),
            llm_client=client,
            llm_config=LLMConfig(provider="openai", model="gpt-4o-mini"),
        )
        result = await supervisor.run("What's the weather in Tokyo and what is 12 + 34?")
    """

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

    async def _decompose(self, input_text: str) -> list[str]:
        """Split a compound query into independent sub-tasks."""
        agent_list = "\n".join(f"- {name}" for name in self.agents)
        messages = [
            Message(
                role="system",
                content=_DECOMPOSE_SYSTEM.format(agent_list=agent_list),
            ),
            Message(role="user", content=input_text),
        ]

        response = await self._client.complete(self._config, messages)

        try:
            tasks = orjson.loads(response.text)
            if isinstance(tasks, list) and all(isinstance(t, str) for t in tasks):
                log.debug("supervisor.decompose", task_count=len(tasks))
                return tasks
        except Exception:
            log.debug("supervisor.decompose_fallback", raw=response.text)

        return [input_text]

    async def _route_and_run(
        self, sub_task: str, *, context: RunContext | None = None
    ) -> AgentResult:
        """Route a single sub-task to the best agent and run it."""
        chosen = await self.router.route(sub_task, 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("supervisor.no_route", sub_task=sub_task[:80])
            return AgentResult(
                output="",
                agent_name="supervisor",
                metadata={"error": "no_route_found", "sub_task": sub_task},
            )

        agent = self.agents[chosen]
        log.debug("supervisor.delegate", chosen=chosen, sub_task=sub_task[:80])
        result = await agent.run(sub_task, context=context)

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

    async def _synthesize(self, query: str, results: list[AgentResult]) -> str:
        """Combine multiple sub-agent results into a single response."""
        parts = []
        for r in results:
            sub_task = r.metadata.get("sub_task", "")
            parts.append(f"Sub-task: {sub_task}\nAgent: {r.agent_name}\nAnswer: {r.output}")

        messages = [
            Message(role="system", content=_SYNTHESIZE_SYSTEM),
            Message(
                role="user",
                content=(
                    f"Original question: {query}\n\n"
                    f"Agent results:\n\n" + "\n\n".join(parts)
                ),
            ),
        ]

        response = await self._client.complete(self._config, messages)
        return response.text

    async def run(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AgentResult:
        """Decompose, delegate to sub-agents, and synthesize a final answer."""
        sub_tasks = await self._decompose(input_text)

        # Single task — skip synthesis overhead
        if len(sub_tasks) == 1:
            return await self._route_and_run(sub_tasks[0], context=context)

        # Multiple sub-tasks — delegate each, then synthesize
        results: list[AgentResult] = []
        total_iterations = 0
        total_tool_calls = 0

        for sub_task in sub_tasks:
            result = await self._route_and_run(sub_task, context=context)
            results.append(result)
            total_iterations += result.iterations
            total_tool_calls += result.tool_calls_made

        synthesized = await self._synthesize(input_text, results)

        return AgentResult(
            output=synthesized,
            agent_name="supervisor",
            iterations=total_iterations,
            tool_calls_made=total_tool_calls,
            metadata={
                "sub_tasks": [r.metadata for r in results],
            },
        )

    async def run_stream(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Decompose, delegate with streaming, and synthesize."""
        sub_tasks = await self._decompose(input_text)

        results: list[AgentResult] = []
        total_iterations = 0
        total_tool_calls = 0

        for sub_task in sub_tasks:
            chosen = await self.router.route(sub_task, 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:
                yield AgentError(error="no_route_found")
                continue

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

            # Stream sub-agent events
            agent_output = ""
            agent_iterations = 0
            agent_tool_calls = 0
            run_stream = getattr(agent, "run_stream", None)

            if run_stream is not None:
                async for event in run_stream(sub_task, context=context):
                    yield event
                    if isinstance(event, (AgentComplete, AgentError)):
                        agent_output = event.output
                        agent_iterations = event.iterations
                        agent_tool_calls = event.tool_calls_made
            else:
                result = await agent.run(sub_task, context=context)
                agent_output = result.output
                agent_iterations = result.iterations
                agent_tool_calls = result.tool_calls_made
                yield AgentComplete(
                    output=result.output,
                    iterations=result.iterations,
                    tool_calls_made=result.tool_calls_made,
                )

            total_iterations += agent_iterations
            total_tool_calls += agent_tool_calls
            results.append(AgentResult(
                output=agent_output,
                agent_name=chosen,
                iterations=agent_iterations,
                tool_calls_made=agent_tool_calls,
                metadata={"routed_to": chosen, "sub_task": sub_task},
            ))

        # Synthesize if multiple sub-tasks
        if len(results) > 1:
            synthesized = await self._synthesize(input_text, results)
            yield TextDelta(text="\n")
            yield AgentComplete(
                output=synthesized,
                iterations=total_iterations,
                tool_calls_made=total_tool_calls,
                metadata={"synthesized": True},
            )
        elif len(results) == 1 and not any(
            isinstance(e, AgentComplete) for e in []
        ):
            # Single task already yielded its AgentComplete via streaming
            pass
