backend/models.py

"""Pydantic models for API request/response contracts."""

from __future__ import annotations

from datetime import date, datetime
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


class CabinClass(str, Enum):
    economy = "economy"
    premium_economy = "premium_economy"
    business = "business"
    first = "first"


class TripType(str, Enum):
    one_way = "one_way"
    round_trip = "round_trip"
    multi_city = "multi_city"


class SortBy(str, Enum):
    best = "best"
    cheapest = "cheapest"
    fastest = "fastest"


# --- Airport ---

class AirportInfo(BaseModel):
    iata: str
    name: str
    city_name: str
    country: str
    country_code: str
    continent: str
    latitude: float
    longitude: float
    timezone: str
    hub_score: float = 0.0
    route_count: int = 0


# --- Flight segment ---

class FlightSegment(BaseModel):
    airline_code: str
    airline_name: str
    flight_number: str
    aircraft: str
    origin: str
    origin_city: str
    destination: str
    destination_city: str
    departure: datetime
    arrival: datetime
    duration_minutes: int
    # Amenities
    legroom_inches: int = 31
    has_wifi: bool = False
    wifi_type: Optional[str] = None       # "Free Wi-Fi", "Wi-Fi for a fee"
    has_power: bool = False
    has_usb: bool = False
    has_video: bool = False
    video_type: Optional[str] = None      # "On-demand video", "Live TV", "Seatback screen"


# --- Flight offer (may have multiple segments) ---

class FlightOffer(BaseModel):
    id: str
    segments: list[FlightSegment]
    total_duration_minutes: int
    stops: int
    price_usd: float
    cabin_class: CabinClass
    origin: str
    destination: str
    departure: datetime
    arrival: datetime
    emissions_kg: int = 0
    is_best: bool = False


# --- Search request ---

class SearchLeg(BaseModel):
    origin: str = Field(..., min_length=3, max_length=3, description="IATA code")
    destination: str = Field(..., min_length=3, max_length=3, description="IATA code")
    date: date


class Passengers(BaseModel):
    adults: int = Field(1, ge=1, le=9)
    children: int = Field(0, ge=0, le=9)
    infants: int = Field(0, ge=0, le=4)

    @property
    def total(self) -> int:
        return self.adults + self.children + self.infants


class Filters(BaseModel):
    max_stops: Optional[int] = None
    max_price: Optional[float] = None
    max_duration_minutes: Optional[int] = None
    airlines: Optional[list[str]] = None  # IATA codes to include
    departure_time_min: Optional[str] = None  # "06:00"
    departure_time_max: Optional[str] = None  # "18:00"


class SearchRequest(BaseModel):
    trip_type: TripType = TripType.round_trip
    legs: list[SearchLeg] = Field(..., min_length=1, max_length=6)
    passengers: Passengers = Passengers()
    cabin_class: CabinClass = CabinClass.economy
    filters: Filters = Filters()
    sort_by: SortBy = SortBy.best


class SearchResponse(BaseModel):
    outbound_flights: list[FlightOffer]
    return_flights: list[FlightOffer] = []
    search_id: str
    origin: str
    destination: str
    same_airline_discount: float = 1.0


# --- Calendar ---

class CalendarDay(BaseModel):
    date: date
    cheapest_price: Optional[float] = None


class CalendarResponse(BaseModel):
    origin: str
    destination: str
    year: int
    month: int
    days: list[CalendarDay]


# --- Autocomplete ---

class AutocompleteResult(BaseModel):
    iata: str
    name: str
    city_name: str
    country: str
    display_name: str
    hub_score: float = 0.0


# --- Booking ---

class PassengerInfo(BaseModel):
    first_name: str
    last_name: str
    email: str
    phone: str


class PaymentInfo(BaseModel):
    card_number: str
    expiry_date: str  # MM/YY
    cvv: str
    cardholder_name: str


class BookingFlightInfo(BaseModel):
    id: str
    origin: str
    destination: str
    departure: datetime
    arrival: datetime
    total_duration_minutes: int
    stops: int
    price_usd: float
    cabin_class: CabinClass
    segments: list[FlightSegment]


class BookingRequest(BaseModel):
    passenger: PassengerInfo
    payment: PaymentInfo
    outbound_flight: BookingFlightInfo
    return_flight: Optional[BookingFlightInfo] = None


class BookingResponse(BaseModel):
    booking_id: str
    status: str
    booked_at: str
    passenger: PassengerInfo
    payment_summary: dict
    outbound_flight: BookingFlightInfo
    return_flight: Optional[BookingFlightInfo] = None
    total_price: float