tools.py

import os
import json
from datetime import date
from typing import Optional
from pydantic import BaseModel
from serpapi import GoogleSearch

import asyncio

BRAZIL_AIRPORTS = {
    # São Paulo
    "são paulo": "GRU,CGH",
    "sao paulo": "GRU,CGH",
    "saopaulo": "GRU,CGH",
    "sãopaulo": "GRU,CGH",
    "sp": "GRU,CGH",
    "guarulhos": "GRU",
    "campinas": "VCP",
    "viracopos": "VCP",
    "congonhas": "CGH",

    # Sudeste
    "rio de janeiro": "GIG,SDU",
    "riodejaneiro": "GIG,SDU",
    "rj": "GIG,SDU",
    "belo horizonte": "CNF,PLU",
    "belohorizonte": "CNF,PLU",
    "bh": "CNF,PLU",
    "vitória": "VIX",
    "vitoria": "VIX",
    "juiz de fora": "IZA",
    "juizdefora": "IZA",
    "uberlândia": "UDI",
    "uberlandia": "UDI",
    "montes claros": "MOC",
    "montesclaros": "MOC",
    "governador valadares": "GVR",
    "governadorvaladares": "GVR",
    "uberaba": "UBA",

    # Nordeste
    "salvador": "SSA",
    "bahia": "SSA",
    "recife": "REC",
    "fortaleza": "FOR",
    "natal": "NAT",
    "maceió": "MCZ",
    "maceio": "MCZ",
    "joão pessoa": "JPA",
    "joao pessoa": "JPA",
    "joaopessoa": "JPA",
    "aracaju": "AJU",
    "são luís": "SLZ",
    "sao luis": "SLZ",
    "saoluis": "SLZ",
    "teresina": "THE",
    "ilhéus": "IOS",
    "ilheus": "IOS",
    "porto seguro": "BPS",
    "portoseguro": "BPS",
    "vitoria da conquista": "VDC",
    "vitória da conquista": "VDC",
    "vitóriadaconquista": "VDC",
    "feira de santana": "FEC",
    "feiraadesantana": "FEC",
    "petrolina": "PNZ",
    "paulo afonso": "PAV",
    "pauloafonso": "PAV",
    "juazeiro do norte": "JDO",
    "juazeirodonorte": "JDO",
    "campina grande": "CPV",
    "campinavgrande": "CPV",
    "mossoró": "MVF",
    "mossoro": "MVF",
    "imperatriz": "IMP",
    "barreiras": "BRA",

    # Sul
    "porto alegre": "POA",
    "portoalegre": "POA",
    "curitiba": "CWB",
    "florianópolis": "FLN",
    "florianopolis": "FLN",
    "navegantes": "NVT",
    "joinville": "JOI",
    "londrina": "LDB",
    "foz do iguaçu": "IGU",
    "fozdoiguacu": "IGU",
    "maringá": "MGF",
    "maringa": "MGF",
    "chapecó": "XAP",
    "chapeco": "XAP",
    "passo fundo": "PFB",
    "passofundo": "PFB",
    "caxias do sul": "CXJ",
    "caxiasdosul": "CXJ",
    "cascavel": "CAC",

    # Centro-Oeste
    "brasília": "BSB",
    "brasilia": "BSB",
    "goiânia": "GYN",
    "goiania": "GYN",
    "cuiabá": "CGB",
    "cuiaba": "CGB",
    "campo grande": "CGR",
    "campogrande": "CGR",
    "corumbá": "CMG",
    "corumba": "CMG",
    "rondonópolis": "ROO",
    "rondonopolis": "ROO",
    "sinop": "OPS",

    # Norte
    "manaus": "MAO",
    "belém": "BEL",
    "belem": "BEL",
    "porto velho": "PVH",
    "portovelho": "PVH",
    "boa vista": "BVB",
    "boavista": "BVB",
    "palmas": "PMW",
    "macapá": "MCP",
    "macapa": "MCP",
    "rio branco": "RBR",
    "riobranco": "RBR",
    "santarém": "STM",
    "santarem": "STM",
    "marabá": "MAB",
    "maraba": "MAB",
    "altamira": "ATM",
    "carajás": "CKS",
    "carajas": "CKS",
}



def city_to_iata(city: str) -> str:
    normalized = city.lower().strip()

    # 1. Match exato primeiro
    if normalized in BRAZIL_AIRPORTS:
        return BRAZIL_AIRPORTS[normalized]

    # 2. Busca parcial — pega a chave mais longa (mais específica)
    best_key = None
    best_iata = None
    for key, iata in BRAZIL_AIRPORTS.items():
        if key in normalized or normalized in key:
            if best_key is None or len(key) > len(best_key):
                best_key = key
                best_iata = iata

    if best_iata:
        return best_iata

    return city.upper()[:3]


class FlightQuery(BaseModel):
    origin_city: str
    destination_city: str
    departure_date: date
    return_date: Optional[date] = None
    adults: int = 1


def _search_flights_sync(query: FlightQuery) -> str:
    origin = city_to_iata(query.origin_city)
    dest = city_to_iata(query.destination_city)

    params = {
        "engine": "google_flights",
        "departure_id": origin,
        "arrival_id": dest,
        "outbound_date": query.departure_date.isoformat(),
        "currency": "BRL",
        "hl": "pt",
        "gl": "br",
        "adults": query.adults,
        "sort_by": "2",  # preço
        "api_key": os.getenv("SERPAPI_KEY"),
    }

    # ida e volta
    if query.return_date:
        params["type"] = "1"
        params["return_date"] = query.return_date.isoformat()
    else:
        params["type"] = "2"

    results = GoogleSearch(params).get_dict()

    all_flights = results.get("best_flights", []) + results.get("other_flights", [])

    if not all_flights:
        return json.dumps({"error": "Nenhum voo encontrado para essa rota/data."})

    flights_out = []
    for f in all_flights[:8]:
        legs = f["flights"]
        first, last = legs[0], legs[-1]
        h, m = divmod(f["total_duration"], 60)

        flights_out.append({
            "airline": first["airline"],
            "origin_airport": first["departure_airport"]["id"],
            "origin_name": first["departure_airport"].get("name", ""),
            "dest_airport": last["arrival_airport"]["id"],
            "dest_name": last["arrival_airport"].get("name", ""),
            "departure": first["departure_airport"]["time"],
            "arrival": last["arrival_airport"]["time"],
            "duration": f"{h}h{m:02d}m",
            "stops": len(legs) - 1,
            "stop_cities": [
                leg["arrival_airport"]["id"]
                for leg in legs[:-1]
            ] if len(legs) > 1 else [],
            "price_brl": f["price"],
        })

    insights = results.get("price_insights", {})

    return json.dumps({
        "route": f"{origin} → {dest}",
        "date": query.departure_date.isoformat(),
        "total_found": len(all_flights),
        "flights": flights_out,
        "price_insights": {
            "level": insights.get("price_level", ""),
            "lowest": insights.get("lowest_price"),
            "typical_range": insights.get("typical_price_range", []),
        },
    }, ensure_ascii=False, indent=2)


async def search_flights_serpapi(query: FlightQuery) -> str:
    """Wrapper assíncrono — roda a chamada bloqueante em thread pool."""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _search_flights_sync, query)