"""
OpenAI Agents SDK (Python) instrumented with Cascade SDK

Demonstrates:
 - instrument_openai_agents() auto-instrumentation via trace processor
 - Agent with tools, handoffs
 - All LLM calls, tool executions, agent spans captured automatically

Prerequisites:
    pip install openai-agents cascade-sdk

Run:
    python tests/test_openai_agents_sdk.py
"""

import os
import asyncio
import json
import random
from dotenv import load_dotenv

load_dotenv()

from cascade import init_tracing, trace_run
from cascade.integrations import instrument_openai_agents

from agents import Agent, Runner, function_tool

# ---------------------------------------------------------------------------
# 1. Cascade tracing + auto-instrument the OpenAI Agents SDK
# ---------------------------------------------------------------------------

init_tracing(project="openai-agents-python-test")
instrument_openai_agents()

# ---------------------------------------------------------------------------
# 2. Tools
# ---------------------------------------------------------------------------


@function_tool
def search_flights(origin: str, destination: str, departure_date: str) -> str:
    """Search for available flights between two cities."""
    airlines = ["United Airlines", "Delta", "American Airlines"]
    results = []
    for airline in airlines:
        price = 380 + random.randint(0, 200)
        hours = 5 + random.randint(0, 7)
        mins = random.randint(0, 59)
        results.append(f"  {airline} - ${price}, {hours}h {mins}m")
    return f"Flights from {origin} to {destination} on {departure_date}:\n" + "\n".join(results)


@function_tool
def search_hotels(city: str, check_in: str, check_out: str) -> str:
    """Search for available hotels in a city."""
    hotels = [
        ("Grand Plaza Hotel", 5, 250),
        ("City Center Inn", 4, 150),
        ("Budget Stay Hostel", 3, 80),
        ("Luxury Suites", 5, 400),
        ("Comfort Lodge", 3, 120),
    ]
    lines = []
    for name, stars, price in hotels:
        lines.append(f"  {name} ({'*' * stars}) - ${price}/night")
    return f"Hotels in {city} ({check_in} to {check_out}):\n" + "\n".join(lines)


@function_tool
def get_weather_forecast(city: str, date: str) -> str:
    """Get weather forecast for a city on a specific date."""
    conditions = ["Sunny", "Partly Cloudy", "Cloudy", "Clear"]
    condition = random.choice(conditions)
    low = 15 + random.randint(0, 5)
    high = low + 5 + random.randint(0, 5)
    humidity = 40 + random.randint(0, 30)
    return f"Weather for {city} on {date}: {condition}, {low}-{high}°C, Humidity: {humidity}%"


@function_tool
def get_travel_recommendations(city: str) -> str:
    """Get travel recommendations for a city."""
    db = {
        "Paris": {
            "attractions": ["Eiffel Tower", "Louvre Museum", "Notre-Dame", "Champs-Elysees", "Montmartre"],
            "food": ["Croissants", "French Onion Soup", "Coq au Vin", "Macarons"],
            "activities": ["River Seine Cruise", "Versailles Day Trip", "Art Gallery Tours"],
        },
        "Tokyo": {
            "attractions": ["Senso-ji Temple", "Shibuya Crossing", "Tokyo Skytree"],
            "food": ["Sushi", "Ramen", "Tempura", "Takoyaki"],
            "activities": ["Cherry Blossom Viewing", "Onsen Experience"],
        },
    }
    data = db.get(city, {"attractions": ["City Center"], "food": ["Local Cuisine"], "activities": ["City Tour"]})
    lines = [
        f"Recommendations for {city}:",
        f"  Attractions: {', '.join(data['attractions'])}",
        f"  Food: {', '.join(data['food'])}",
        f"  Activities: {', '.join(data['activities'])}",
    ]
    return "\n".join(lines)


@function_tool
def get_currency_info(city: str) -> str:
    """Get currency and exchange rate info for a city."""
    currencies = {
        "Paris": ("EUR", 0.92),
        "Tokyo": ("JPY", 150.0),
        "London": ("GBP", 0.79),
    }
    currency, rate = currencies.get(city, ("USD", 1.0))
    return f"Currency for {city}: {currency} (1 USD = {rate} {currency})"


# ---------------------------------------------------------------------------
# 3. Agent — single agent with all tools (like the Python travel agent)
# ---------------------------------------------------------------------------

travel_agent = Agent(
    name="Travel Agent",
    instructions="""You are a helpful travel agent assistant. Your job is to help users plan their trips by:
1. Searching for flights
2. Finding hotels
3. Checking weather forecasts
4. Providing travel recommendations
5. Providing currency information

Use ALL available tools to gather complete information before giving your final response.
Be friendly and thorough.""",
    model="gpt-4.1-mini",
    tools=[search_flights, search_hotels, get_weather_forecast, get_travel_recommendations, get_currency_info],
)

# ---------------------------------------------------------------------------
# 4. Main
# ---------------------------------------------------------------------------


async def main():
    print("=== OpenAI Agents SDK (Python) + Cascade ===\n")

    user_request = (
        "I want to plan a 5-day trip to Paris from New York. "
        "I'm interested in museums, food, and art. "
        "Departure: 2026-06-15, Return: 2026-06-20. "
        "Please find flights, hotels, check the weather, "
        "get recommendations, and provide currency information."
    )

    print(f"Customer: {user_request}\n")

    with trace_run("TravelPlanningSession", metadata={"destination": "Paris"}):
        result = await Runner.run(travel_agent, user_request)
        print(f"\nAgent: {result.final_output}\n")

    print("=== Done ===")
    print("Project: openai-agents-python-test")
    print("Check your Cascade dashboard for the full trace tree.")
    print()
    print("What you should see:")
    print("  - Root span: TravelPlanningSession")
    print("  - Agent span: Travel Agent")
    print("  - LLM spans: GenerationSpanData (with model, tokens, input/output)")
    print("  - Tool spans: FunctionSpanData (search_flights, search_hotels, etc.)")


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