#!/usr/bin/env python3
"""09 — Streaming agent with tool calls.

Watch the agent think, call tools, and respond in real time.
Each event type is printed as it arrives.

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

from __future__ import annotations

import asyncio

from koan.agents import (
    AgentComplete,
    AgentConfig,
    IterationStart,
    StreamingReactAgent,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.llm.client import LLMClient
from koan.tools import ToolRegistry, tool


@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is 22°C and partly cloudy."


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


@tool
def search_docs(query: str) -> str:
    """Search internal documentation."""
    return f"Found 3 results for '{query}': [Doc A, Doc B, Doc C]"


async def main() -> None:
    registry = ToolRegistry()
    registry.add(get_weather)
    registry.add(add_numbers)
    registry.add(search_docs)

    async with LLMClient() as client:
        agent = StreamingReactAgent(
            "assistant",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt=(
                    "You are a helpful assistant with weather, calculator, and search tools. "
                    "Always use tools when relevant."
                ),
            ),
            tools=registry,
        )

        queries = [
            "What's the weather in Paris? and What is the sum of 15 and 27?",
            "What is 42 + 58?",
        ]

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

            async for event in agent.run_stream(query):
                if isinstance(event, IterationStart):
                    print(f"\n[Iteration {event.iteration}]")
                elif isinstance(event, TextDelta):
                    print(event.text, end="", flush=True)
                elif isinstance(event, ToolCallStart):
                    print(f"\n  -> Calling tool: {event.tool_name}({event.arguments})")
                elif isinstance(event, ToolCallEnd):
                    if event.error:
                        print(f"  <- Error: {event.error}")
                    else:
                        print(f"  <- Result: {event.result}")
                elif isinstance(event, AgentComplete):
                    print(f"\n[Done: {event.iterations} iterations, {event.tool_calls_made} tool calls]")

            print()


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