import gradio as gr
from dotenv import load_dotenv
from openai import OpenAI
import pandas as pd

# Load environment variables and initialize OpenAI client
load_dotenv()
openai = OpenAI()
MODEL = "gpt-4o-mini"

flights_df = pd.read_csv('flightai_fictitious_flights.csv')

# System message guiding the assistant
system_message = """
You are a helpful assistant for an Airline called FlightAI.
You have access to a tool called "query_flights" that provides information about flights.

Guidelines:
- Give short, courteous answers, no more than 1-2 sentences when possible.
- Always use the query_flights tool to search for flights when users ask about flights.
- If no flights match the criteria, politely suggest alternative options.
- If you don't know the answer, say so honestly.
- Present flight information clearly, highlighting the best options.

Examples:
- "Show me flights from New York to London" → Use query_flights with origin_city="New York" and destination_city="London"
- "Cheapest business class to Paris" → Use query_flights with destination_city="Paris" and cabin_class="business"
"""

def query_flights(  origin_city = None,
                    destination_city = None,
                    departure_date = None,
                    max_price = None,
                    cabin_class = None,
                    airline = None,
                    limit = 5) :
    
    df = flights_df.copy()
    # Filter by origin
    if origin_city:
        df = df[df["origin_city"].str.lower() == origin_city.lower()]
    
    # Filter by destination
    if destination_city:
        df = df[df["destination_city"].str.lower() == destination_city.lower()]

    # Filter by date if provided
    if departure_date:
        df = df[df["departure_date"] == departure_date]

    # Filter by cabin class
    if cabin_class:
        df = df[df["cabin_class"].str.lower() == cabin_class.lower()]
    
    # Filter by airline
    if airline:
        df = df[df["airline"].str.lower() == airline.lower()]
    
    # Filter by max price
    if max_price and max_price > 0:
        df = df[df["price_usd"] <= max_price]
    
    if df.empty:
        return "No flights match your criteria."

    # Sort by price and limit results
    df = df.sort_values("price_usd").head(limit)
    df["price_usd"] = df["price_usd"].round(2)
    
    results = df.to_dict(orient="records")
    return results

# Define the tool for the LLM
query_function = {
    "name": "query_flights",
    "description": "Query the FlightAI database for flights using flexible filters. Use this whenever a user asks about flights, prices, or travel options.",
    "parameters": {
        "type": "object",
        "properties": {
            "origin_city": {
                "type": "string",
                "description": "The departure city name (e.g., 'New York', 'London')"
            },
            "destination_city": {
                "type": "string",
                "description": "The arrival city name (e.g., 'Paris', 'Tokyo')"
            },
            "departure_date": {
                "type": "string",
                "description": "Departure date in flexible format (e.g., 'tomorrow', '2024-12-25', 'next Monday')"
            },
            "max_price": {
                "type": "number",
                "description": "Maximum price in USD"
            },
            "cabin_class": {
                "type": "string",
                "description": "Cabin class: 'economy', 'business', or 'first'"
            },
            "airline": {
                "type": "string",
                "description": "Airline name"
            },
            "limit": {
                "type": "integer",
                "description": "Maximum number of results to return (default 5, max 20)",
                "default": 5
            }
        },
        "required": []
    }
}

tools = [query_function]


def chat(message, history):
    history = [{"role": h["role"], "content": h["content"]} for h in history]
    messages = [{"role": "system", "content": system_message}] + history + [{"role": "user", "content": message}]
    
    response = openai.chat.completions.create(
        model=MODEL,
        messages=messages,
        functions=tools,
        function_call="auto"
    )
    
    choice = response.choices[0].message
    # If LLM wants to call a function
    if choice.get("function_call"):
        func_name = choice["function_call"]["name"]
        args = eval(choice["function_call"]["arguments"])  # convert string to dict
        if func_name == "query_flights":
            result = query_flights(**args)
            # Add LLM response after function execution
            messages.append({"role": "function", "name": func_name, "content": str(result)})
            followup = openai.chat.completions.create(model=MODEL, messages=messages)
            return followup.choices[0].message.content
    else:
        return choice.content



gr.ChatInterface(fn=chat,
title="OpenAI Coder Chat",
).launch()