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

Demonstrates both non-streaming and streaming modes of the Supervisor.

The Supervisor automatically:
1. Decomposes compound queries into sub-tasks
2. Delegates each sub-task to the best specialist agent
3. Synthesizes a unified final response

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

from __future__ import annotations

import asyncio

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

# --- 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")


@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


def build_supervisor(client: LLMClient, *, streaming: bool) -> Supervisor:
    """Build a Supervisor with weather, math, and general agents."""
    weather_tools = ToolRegistry()
    weather_tools.add(get_weather)

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

    agent_cls = StreamingReactAgent if streaming else ReactAgent

    weather_agent = agent_cls(
        "weather",
        client=client,
        config=AgentConfig(
            provider="openai",
            model="gpt-4o-mini",
            system_prompt="Use the get_weather tool to answer. Be concise.",
        ),
        tools=weather_tools,
    )

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

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

    return Supervisor(
        agents={"weather": weather_agent, "math": math_agent, "general": general_agent},
        router=RuleRouter(
            rules={
                "weather": "weather",
                "temperature": "weather",
                "forecast": "weather",
                "add": "math",
                "sum": "math",
                "multiply": "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."
    ),
    "What is the sum of 99 and 1?",
]


async def demo_no_streaming() -> None:
    """Run the Supervisor without streaming — just get final results."""
    print("=" * 60)
    print("MODE: No Streaming (supervisor.run)")
    print("=" * 60)

    async with LLMClient() as client:
        supervisor = build_supervisor(client, streaming=False)

        for query in QUERIES:
            print(f"\nQuery: {query}")
            print("-" * 40)

            result = await supervisor.run(query)

            print(f"Agent: {result.agent_name}")
            print(f"Output: {result.output}")
            print(f"Iterations: {result.iterations}")
            print(f"Tool calls: {result.tool_calls_made}")


async def demo_streaming() -> None:
    """Run the Supervisor with streaming — see sub-agent work live."""
    print("\n" + "=" * 60)
    print("MODE: Streaming (supervisor.run_stream)")
    print("=" * 60)

    async with LLMClient() as client:
        supervisor = build_supervisor(client, streaming=True)

        for query in QUERIES:
            print(f"\nQuery: {query}")
            print("-" * 40)

            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("\n  [Supervisor final answer]")
                        print(f"  {event.output}")
                    else:
                        print(
                            f"\n  [Done: {event.iterations} iters, "
                            f"{event.tool_calls_made} tool calls]"
                        )
                elif isinstance(event, AgentError):
                    print(f"\n  [Error: {event.error}]")

            print()


async def main() -> None:
    await demo_no_streaming()
    await demo_streaming()


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