"""OpenAI conversational agent with tool access, instrumented by Flowlines.

A simple interactive agent that can look up city information using a tool.
All LLM calls are automatically traced and exported to the Flowlines backend.

Usage:
    1. Copy ``.env.example`` to ``.env`` and fill in your API keys.
    2. Run: ``uv run python examples/openai/agent.py``
"""

from __future__ import annotations

import json
import os
import uuid

from dotenv import load_dotenv
from openai import OpenAI

import flowlines

# ---------------------------------------------------------------------------
# City data & tool definition
# ---------------------------------------------------------------------------

CITIES: dict[str, dict[str, str | int]] = {
    "paris": {
        "name": "Paris",
        "country": "France",
        "population": 2_161_000,
        "known_for": "Eiffel Tower, Louvre Museum, croissants, fashion",
        "fun_fact": "Paris has only one stop sign in the entire city.",
    },
    "tokyo": {
        "name": "Tokyo",
        "country": "Japan",
        "population": 13_960_000,
        "known_for": "Shibuya Crossing, cherry blossoms, sushi, anime",
        "fun_fact": "Tokyo has more Michelin-starred restaurants than any other city.",
    },
    "new york": {
        "name": "New York",
        "country": "United States",
        "population": 8_336_000,
        "known_for": "Statue of Liberty, Central Park, Broadway, pizza",
        "fun_fact": "New York City's subway runs 24/7, one of the few in the world.",
    },
    "sydney": {
        "name": "Sydney",
        "country": "Australia",
        "population": 5_312_000,
        "known_for": "Opera House, Harbour Bridge, Bondi Beach, surfing",
        "fun_fact": "Sydney's Opera House roof is covered with over 1 million tiles.",
    },
    "nairobi": {
        "name": "Nairobi",
        "country": "Kenya",
        "population": 4_397_000,
        "known_for": "Nairobi National Park, vibrant tech scene, Great Rift Valley",
        "fun_fact": "Nairobi is the only capital city with a national park inside it.",
    },
}


def get_city_info(city: str) -> str:
    """Look up information about a city."""
    info = CITIES.get(city.lower())
    if info is None:
        return json.dumps({"error": f"No information available for '{city}'."})
    return json.dumps(info)


TOOLS: list[dict[str, object]] = [
    {
        "type": "function",
        "name": "get_city_info",
        "description": "Get information about a city including population, etc.",
        "parameters": {
            "type": "object",
            "properties": {
                "city": {
                    "type": "string",
                    "description": "The city name to look up (e.g. 'Paris', 'Tokyo').",
                },
            },
            "required": ["city"],
        },
    },
]

SYSTEM_PROMPT = (
    "You are a helpful travel assistant. When the user asks about a city, "
    "use the get_city_info tool to retrieve details. Summarise the result "
    "in a friendly, conversational way."
)

# ---------------------------------------------------------------------------
# Agent turn
# ---------------------------------------------------------------------------


def agent_turn(
    client: OpenAI,
    conversation: list[dict[str, object]],
    user_input: str,
) -> str:
    """Run one agent turn: send the user message, handle tool calls, return reply."""
    conversation.append({"role": "user", "content": user_input})

    response = client.responses.create(
        model="gpt-4.1-nano",
        instructions=SYSTEM_PROMPT,
        input=conversation,
        tools=TOOLS,
    )

    # Handle tool-call loop
    while True:
        tool_calls = [item for item in response.output if item.type == "function_call"]
        if not tool_calls:
            break

        for tool_call in tool_calls:
            args = json.loads(tool_call.arguments)
            result = get_city_info(args["city"])

            conversation.append(tool_call)
            conversation.append(
                {
                    "type": "function_call_output",
                    "call_id": tool_call.call_id,
                    "output": result,
                }
            )

        response = client.responses.create(
            model="gpt-4.1-nano",
            instructions=SYSTEM_PROMPT,
            input=conversation,
            tools=TOOLS,
        )

    # Append assistant output to conversation for multi-turn context
    for item in response.output:
        conversation.append(item)

    return response.output_text


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    """Run the interactive agent loop."""
    load_dotenv()

    flowlines_api_key = os.environ.get("FLOWLINES_API_KEY", "")
    if not flowlines_api_key:
        print("Error: FLOWLINES_API_KEY is not set. See .env.example.")
        return

    if not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY is not set. See .env.example.")
        return

    # Initialize Flowlines *before* creating the OpenAI client so that
    # auto-instrumentation hooks are in place.
    flowlines.init(api_key=flowlines_api_key, verbose=True)

    client = OpenAI()
    conversation: list[dict[str, object]] = []
    user_id = f"example-user-{uuid.uuid4().hex[:8]}"
    session_id = f"sess-{uuid.uuid4().hex[:8]}"

    print("City Info Agent (type 'quit' to exit)")
    print(f"  user_id: {user_id}")
    print(f"  session_id: {session_id}")
    print()

    try:
        with flowlines.context(user_id=user_id, session_id=session_id):
            while True:
                try:
                    user_input = input("You: ").strip()
                except (EOFError, KeyboardInterrupt):
                    print()
                    break

                if not user_input:
                    continue
                if user_input.lower() in ("quit", "exit"):
                    break

                reply = agent_turn(client, conversation, user_input)
                print(f"Agent: {reply}")
                print()
    finally:
        flowlines.shutdown()


if __name__ == "__main__":
    main()
