planmate/flights.py

from typing import List, Dict, Tuple
from .clients.amadeus_client import AmadeusClient
from .utils import format_price


def resolve_city_to_iata(city_name: str) -> Tuple[str, str]:
    """Resolve a human city name to an IATA city code using Amadeus.
    Returns (iata_code, display_label). Raises RuntimeError if not found.
    """
    client = AmadeusClient()
    res = client.locations(city_name, subtypes="CITY")
    data = res.get("data", [])
    if not data:
        # try airports as fallback and take the cityName + iataCode of airport
        res = client.locations(city_name, subtypes="AIRPORT")
        data = res.get("data", [])
    if not data:
        raise RuntimeError(f"Could not resolve city to IATA: {city_name}")

    # Prefer exact name match then highest relevance (first)
    city_lower = city_name.strip().lower()
    chosen = None
    for d in data:
        name = (d.get("name") or "").lower()
        address_city = (d.get("address", {}).get("cityName") or "").lower()
        if name == city_lower or address_city == city_lower:
            chosen = d
            break
    if not chosen:
        chosen = data[0]

    code = chosen.get("iataCode")
    label = f"{chosen.get('name')} ({code})"
    if not code:
        raise RuntimeError(f"No IATA code found for: {city_name}")
    return code, label


def search_airports(keyword: str) -> List[Dict]:
    client = AmadeusClient()
    res = client.locations(keyword)
    items = []
    for d in res.get("data", []):
        items.append(
            {
                "name": d.get("name"),
                "iataCode": d.get("iataCode"),
                "subType": d.get("subType"),
                "address": d.get("address", {}).get("cityName")
                or d.get("address", {}).get("countryName"),
            }
        )
    return items


def search_flights(
    origin_code: str,
    dest_code: str,
    depart: str,
    ret: str,
    adults=1,
    currency="PKR",
    non_stop: bool = False,
) -> Dict:
    client = AmadeusClient()
    res = client.flight_offers(
        origin_code, dest_code, depart, ret, adults, currency, non_stop
    )
    dicts = res.get("dictionaries", {})
    carriers = dicts.get("carriers", {})
    result = []
    for offer in res.get("data", []):
        price = offer.get("price", {}).get("total")
        itineraries = offer.get("itineraries", [])
        legs = []
        for it in itineraries:
            segs = []
            for s in it.get("segments", []):
                carrier = s.get("carrierCode")
                segs.append(
                    {
                        "from": s.get("departure", {}).get("iataCode"),
                        "to": s.get("arrival", {}).get("iataCode"),
                        "dep": s.get("departure", {}).get("at"),
                        "arr": s.get("arrival", {}).get("at"),
                        "carrier": carriers.get(carrier, carrier),
                        "number": s.get("number"),
                        "duration": it.get("duration", ""),
                    }
                )
            legs.append(segs)
        result.append(
            {
                "price": price,
                "price_label": format_price(price, currency),
                "legs": legs,
                "oneWay": (ret is None or ret == ""),
            }
        )
    return {"flights": result, "carriers": carriers}