tools.py

import os
import dateparser
from amadeus import Client, ResponseError
from langchain.tools import Tool

# Initialize Amadeus API Client
AMADEUS_API_KEY = os.getenv("AMADEUS_API_KEY")
AMADEUS_API_SECRET = os.getenv("AMADEUS_API_SECRET")

amadeus = Client(
    client_id=AMADEUS_API_KEY,
    client_secret=AMADEUS_API_SECRET
)

# 🌍 Convert City Name to IATA Code
def get_airport_code(city_name: str):
    """Find the IATA airport code for a given city."""
    try:
        response = amadeus.reference_data.locations.get(
            keyword=city_name,
            subType="AIRPORT,CITY"
        )
        if response.data:
            return response.data[0]["iataCode"]
        return None
    except ResponseError:
        return None

# ✈️ Get Full Airline Name from Code
def get_airline_name(airline_code: str):
    """Get full airline name using Amadeus API."""
    try:
        response = amadeus.reference_data.airlines.get(airlineCodes=airline_code)
        if response.data:
            return response.data[0]["businessName"]
        return airline_code  # Fallback to code if not found
    except ResponseError:
        return airline_code

# 🕑 Format Flight Duration
def format_duration(duration: str):
    """Convert ISO 8601 duration (PT20H25M) into readable format (20h 25m)."""
    duration = duration.replace("PT", "").replace("H", "h ").replace("M", "m")
    return duration.strip()

# 🔗 Generate Booking Link
def generate_booking_link(from_iata: str, to_iata: str, departure_date: str):
    """Generate a booking link using Google Flights."""
    return f"https://www.google.com/flights?hl=en#flt={from_iata}.{to_iata}.{departure_date};c:USD;e:1;s:0;sd:1;t:f"

# 🛫 Flight Search Tool (Now Gradio-Compatible)
def search_flights(query: str):
    """Search for flights using Amadeus API and return styled Markdown output for Gradio."""
    try:
        words = query.lower().split()
        from_city, to_city, date_phrase = None, None, None

        # Extract "from", "to", and date information
        if "from" in words and "to" in words:
            from_index = words.index("from") + 1
            to_index = words.index("to") + 1
            from_city = " ".join(words[from_index:to_index - 1]).title()
            to_city = " ".join(words[to_index:words.index("in")]) if "in" in words else " ".join(words[to_index:]).title()

        date_phrase = " ".join(words[words.index("in") + 1:]) if "in" in words else None

        # Validate extracted details
        if not from_city or not to_city:
            return "❌ Could not detect valid departure and destination cities. Please use 'from <city> to <city>'."

        # Convert city names to IATA codes
        from_iata = get_airport_code(from_city)
        to_iata = get_airport_code(to_city)

        if not from_iata or not to_iata:
            return f"❌ Could not find airport codes for {from_city} or {to_city}. Please check spelling."

        # Convert date phrase to YYYY-MM-DD
        departure_date = dateparser.parse(date_phrase) if date_phrase else None
        if not departure_date:
            return "❌ Could not understand the travel date. Use formats like 'next week' or 'on May 15'."

        departure_date_str = departure_date.strftime("%Y-%m-%d")

        # Fetch flight offers from Amadeus
        response = amadeus.shopping.flight_offers_search.get(
            originLocationCode=from_iata,
            destinationLocationCode=to_iata,
            departureDate=departure_date_str,
            adults=1,
            max=5
        )

        flights = response.data
        if not flights:
            return f"❌ No flights found from {from_city} to {to_city} on {departure_date_str}."

        # 🏆 IMPROVED OUTPUT FOR GRADIO MARKDOWN
        result = f"### ✈️ Flights from {from_city} ({from_iata}) to {to_city} ({to_iata})\n📅 **Date:** {departure_date_str}\n\n"

        for flight in flights:
            airline_code = flight["validatingAirlineCodes"][0]
            airline_name = get_airline_name(airline_code)
            price = flight["price"]["total"]
            duration = format_duration(flight["itineraries"][0]["duration"])
            departure_time = flight["itineraries"][0]["segments"][0]["departure"]["at"]
            arrival_time = flight["itineraries"][0]["segments"][-1]["arrival"]["at"]
            stops = len(flight["itineraries"][0]["segments"]) - 1

            booking_link = generate_booking_link(from_iata, to_iata, departure_date_str)

            # 🛫 Styled Travel Card with Markdown
            result += f"""
---
🚀 **{airline_name}**  
💰 **Price:** `${price}`  
⏳ **Duration:** {duration}  
📅 **Departure:** {departure_time}  
🛬 **Arrival:** {arrival_time}  
🛑 **Stops:** {stops}  
🔗 **[Book Now]({booking_link})**
---
"""

        # "Show More Flights" Button
        more_flights_link = generate_booking_link(from_iata, to_iata, departure_date_str)
        result += f"\n🔍 **[Show More Flights]({more_flights_link})**"

        return result

    except ResponseError as error:
        return str(error)

# ✅ Register as a Tool
tools = [
    Tool(
        name="Flight Booking",
        func=search_flights,
        description="Find and book flights using natural language. Examples: 'Flight from Delhi to SFO in May', 'Travel from Mumbai to New York next week'."
    )
]