#!/usr/bin/env python3
"""04 — Agent with custom tools.

Create a ReAct agent that can use tools to answer questions.

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

from __future__ import annotations

import asyncio
import random

from koan.agents import AgentConfig, ReactAgent
from koan.llm.client import LLMClient
from koan.tools import ToolRegistry, tool


# Define tools using the @tool decorator — KOAN auto-generates JSON schemas
@tool
def get_weather(city: str, units: str = "celsius") -> str:
    """Get the current weather for a city."""
    # In a real app, this would call a weather API
    temp = random.randint(15, 30)  # noqa: S311
    return f"The weather in {city} is {temp} degrees {units} and sunny."


@tool
def calculate(expression: str) -> str:
    """Evaluate a mathematical expression."""
    try:
        result = eval(expression, {"__builtins__": {}})  # noqa: S307
        return f"Result: {result}"
    except Exception as e:
        return f"Error: {e}"


async def main() -> None:
    # Register tools
    registry = ToolRegistry()
    registry.add(get_weather)
    registry.add(calculate)

    # Create a ReAct agent with tools
    async with LLMClient() as client:
        agent = ReactAgent(
            "assistant",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a helpful assistant with access to weather and calculator tools.",
            ),
            tools=registry,
        )

        # Run the agent — it will decide when to use tools
        result = await agent.run("What's the weather in Tokyo?")

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


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