#!/usr/bin/env python3
"""10 — Supervisor: decompose, delegate, and synthesize.

The Supervisor handles compound queries by:
1. Decomposing the input into independent sub-tasks using an LLM
2. Routing each sub-task to the best specialist agent
3. Synthesizing a unified final response from all results

Requirements:
    pip install koan[llm]
    export OPENAI_API_KEY=sk-...
"""

from __future__ import annotations

import asyncio

from koan.agents import (
    AgentComplete,
    AgentConfig,
    AgentError,
    AgentStart,
    RuleRouter,
    StreamingReactAgent,
    Supervisor,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.llm.client import LLMClient
from koan.llm.config import LLMConfig
from koan.tools import ToolRegistry, tool

# --- Weather tools ---


@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    forecasts = {
        "london": "14°C, rainy",
        "tokyo": "28°C, humid",
        "paris": "22°C, sunny",
        "new york": "18°C, cloudy",
    }
    return forecasts.get(city.lower(), f"{city}: 20°C, clear skies")


# --- Math tools ---


@tool
def add_numbers(a: int, b: int) -> int:
    """Add two numbers together."""
    return a + b


@tool
def multiply_numbers(a: int, b: int) -> int:
    """Multiply two numbers together."""
    return a * b


async def main() -> None:
    weather_tools = ToolRegistry()
    weather_tools.add(get_weather)

    math_tools = ToolRegistry()
    math_tools.add(add_numbers)
    math_tools.add(multiply_numbers)

    async with LLMClient() as client:
        weather_agent = StreamingReactAgent(
            "weather",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the get_weather tool to answer weather questions. Be concise.",
            ),
            tools=weather_tools,
        )

        math_agent = StreamingReactAgent(
            "math",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt=(
                    "Use the math tools (add_numbers, multiply_numbers) to answer. Be concise."
                ),
            ),
            tools=math_tools,
        )

        general_agent = StreamingReactAgent(
            "general",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a helpful general assistant. Be concise.",
            ),
        )

        # The Supervisor handles decomposition, delegation, and synthesis
        supervisor = Supervisor(
            agents={"weather": weather_agent, "math": math_agent, "general": general_agent},
            router=RuleRouter(
                rules={
                    "weather": "weather",
                    "temperature": "weather",
                    "forecast": "weather",
                    "add": "math",
                    "multiply": "math",
                    "sum": "math",
                    "product": "math",
                    "calculate": "math",
                },
                default="general",
            ),
            llm_client=client,
            llm_config=LLMConfig(provider="openai", model="gpt-4o-mini"),
            fallback="general",
        )

        queries = [
            (
                "What's the weather in Tokyo? Also tell me a fun fact "
                "about penguins and calculate the product of 12 and 34."
            ),
            "Calculate the sum of 123 and 456",
            "Tell me a fun fact about penguins",
        ]

        for query in queries:
            print(f"\n{'='*60}")
            print(f"Query: {query}")
            print("=" * 60)

            # --- Streaming mode ---
            async for event in supervisor.run_stream(query):
                if isinstance(event, AgentStart):
                    print(f"\n[Delegated to: {event.agent_name}] {event.input_text}")
                elif isinstance(event, TextDelta):
                    print(event.text, end="", flush=True)
                elif isinstance(event, ToolCallStart):
                    print(f"\n  -> {event.tool_name}({event.arguments})")
                elif isinstance(event, ToolCallEnd):
                    print(f"  <- {event.result}")
                elif isinstance(event, AgentComplete):
                    if event.metadata.get("synthesized"):
                        print(f"\n[Supervisor] Final answer:\n{event.output}")
                    else:
                        print(f"\n[Done: {event.iterations} iters, {event.tool_calls_made} tool calls]")
                elif isinstance(event, AgentError):
                    print(f"\n[Error: {event.error}]")

            print()


if __name__ == "__main__":
    asyncio.run(main())
