backend/orchestrator.py

"""
orchestrator.py — Tripplanner search orchestrator.

Coordinates Delta, IHG, and Resy subagents sequentially over a single shared
BrowserSession, then merges their results into ranked TripResult objects.
"""

from __future__ import annotations

import os
import asyncio
from typing import Callable

from .models import (
    TripSearchRequest,
    FlightOption,
    HotelOption,
    RestaurantOption,
    TripResult,
    SearchProgress,
)
from .tools.browser import BrowserSession
from .agents.delta_agent import search_companion_cert
from .agents.ihg_agent import search_points_availability
from .agents.resy_agent import find_resy_restaurants


# ---------------------------------------------------------------------------
# Public helpers
# ---------------------------------------------------------------------------

def calculate_cash_out_of_pocket(
    flight: FlightOption,
    hotel: HotelOption | None,
) -> float:
    """
    Return the total cash the traveler pays.

    The companion's ticket is covered by the certificate; they only owe the
    companion_taxes.  Hotel is redeemed with points so costs $0 cash.
    Restaurant dining credit is separate from this calculation.
    """
    # Primary ticket + companion taxes (companion cert covers the fare itself)
    return flight.base_price + flight.companion_taxes


def collect_benefits(
    flight: FlightOption,
    hotel: HotelOption | None,
    restaurant: RestaurantOption | None,
) -> list[str]:
    """Return human-readable benefit strings for a trip combination."""
    benefits: list[str] = []

    if flight.eligible_for_companion_cert:
        benefits.append("Delta companion certificate used")

    if hotel is not None:
        if hotel.nights_free > 0:
            benefits.append(
                f"IHG 4th night free ({hotel.nights_free} night{'s' if hotel.nights_free != 1 else ''} free)"
            )
        else:
            benefits.append(f"IHG points redemption at {hotel.property_name}")

    if restaurant is not None:
        if restaurant.resy_credit_eligible:
            benefits.append("Resy $20 dining credit eligible")
        if restaurant.global_dining_access:
            benefits.append(f"Global Dining Access reservation at {restaurant.name}")

    return benefits


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _score_trip(
    flight: FlightOption,
    hotel: HotelOption | None,
    restaurant: RestaurantOption | None,
) -> float:
    """
    Score a trip combination on a 0–100 scale.

    Breakdown:
      +40  companion cert eligible
      +30  hotel found with points availability
      +15  restaurant found
      +15  value score (lower cash out of pocket relative to flight base price)
    """
    score = 0.0

    if flight.eligible_for_companion_cert:
        score += 40.0

    if hotel is not None:
        score += 30.0

    if restaurant is not None:
        score += 15.0

    # Value component: full +15 when cash_out == companion_taxes only (i.e.
    # traveler saved the most), scaling down as cash_out approaches base_price*2.
    cash_out = calculate_cash_out_of_pocket(flight, hotel)
    max_possible_cash = flight.base_price * 2  # both tickets at full price
    if max_possible_cash > 0:
        savings_ratio = 1.0 - (cash_out / max_possible_cash)
        # Clamp to [0, 1] so edge cases don't break the scale.
        savings_ratio = max(0.0, min(1.0, savings_ratio))
        score += savings_ratio * 15.0

    return round(score, 2)


def _extract_date_pairs(
    flights: list[FlightOption],
) -> list[tuple[str, str]]:
    """Return (check_in_iso, check_out_iso) pairs from flight results."""
    return [
        (flight.outbound_date.isoformat(), flight.return_date.isoformat())
        for flight in flights
    ]


def _extract_destinations(flights: list[FlightOption]) -> list[str]:
    return list(dict.fromkeys(flight.destination for flight in flights))


def _extract_cities_and_arrival_dates(
    flights: list[FlightOption],
) -> tuple[list[str], list[str]]:
    cities = list(dict.fromkeys(flight.destination for flight in flights))
    arrival_dates = list(
        dict.fromkeys(flight.outbound_date.isoformat() for flight in flights)
    )
    return cities, arrival_dates


def _find_matching_hotel(
    flight: FlightOption,
    hotels: list[HotelOption],
) -> HotelOption | None:
    """Return the first HotelOption whose destination and dates overlap the flight."""
    for hotel in hotels:
        if hotel.destination != flight.destination:
            continue
        # Dates overlap when check_in <= return_date AND check_out >= outbound_date
        if hotel.check_in <= flight.return_date and hotel.check_out >= flight.outbound_date:
            return hotel
    return None


def _find_matching_restaurant(
    flight: FlightOption,
    restaurants: list[RestaurantOption],
) -> RestaurantOption | None:
    """Return the first RestaurantOption matching the flight city and arrival date."""
    for restaurant in restaurants:
        if (
            restaurant.city == flight.destination
            and restaurant.reservation_date == flight.outbound_date
        ):
            return restaurant
    return None


def _emit(
    callback: Callable[[SearchProgress], None],
    step: str,
    message: str,
    progress: int,
) -> None:
    callback(SearchProgress(step=step, message=message, progress=progress))


# ---------------------------------------------------------------------------
# Main orchestrator
# ---------------------------------------------------------------------------

async def run_trip_search(
    request: TripSearchRequest,
    progress_callback: Callable[[SearchProgress], None],
) -> list[TripResult]:
    """
    Coordinate Delta, IHG, and Resy subagents and return ranked TripResult list.

    All three agents share a single BrowserSession so that login state is
    preserved across site transitions.
    """
    user_data_dir = os.environ.get("BROWSER_USER_DATA_DIR", "./browser_data")

    async with BrowserSession(user_data_dir=user_data_dir) as browser:

        # --- Step 1: Delta companion certificate search ---
        _emit(
            progress_callback,
            step="delta",
            message="Checking companion certificate availability...",
            progress=10,
        )
        def delta_progress(msg: str) -> None:
            _emit(progress_callback, step="delta", message=msg, progress=20)

        flights: list[FlightOption] = await search_companion_cert(
            request, browser, delta_progress
        )

        # --- Step 2: IHG points availability ---
        _emit(
            progress_callback,
            step="ihg",
            message="Checking IHG points availability...",
            progress=40,
        )
        destinations = _extract_destinations(flights)
        date_pairs = _extract_date_pairs(flights)

        def ihg_progress(msg: str) -> None:
            _emit(progress_callback, step="ihg", message=msg, progress=55)

        hotels: list[HotelOption] = await search_points_availability(
            destinations, date_pairs, browser,
            request.ihg_brands, request.trip_duration_nights, ihg_progress
        )

        # --- Step 3: Resy dining options ---
        _emit(
            progress_callback,
            step="resy",
            message="Finding Resy dining options...",
            progress=70,
        )
        cities, arrival_dates = _extract_cities_and_arrival_dates(flights)

        def resy_progress(msg: str) -> None:
            _emit(progress_callback, step="resy", message=msg, progress=80)

        restaurants: list[RestaurantOption] = await find_resy_restaurants(
            cities, arrival_dates, request.party_size, browser, resy_progress
        )

    # BrowserSession closed — merge results in memory.

    # --- Step 4: Merge and rank ---
    _emit(
        progress_callback,
        step="merging",
        message="Ranking trip combinations...",
        progress=90,
    )

    results: list[TripResult] = []
    for flight in flights:
        hotel = _find_matching_hotel(flight, hotels)
        restaurant = _find_matching_restaurant(flight, restaurants)

        cash_out = calculate_cash_out_of_pocket(flight, hotel)
        points_required = hotel.total_points if hotel is not None else 0
        benefits = collect_benefits(flight, hotel, restaurant)
        score = _score_trip(flight, hotel, restaurant)

        results.append(
            TripResult(
                destination=flight.destination,
                flight=flight,
                hotel=hotel,
                restaurant=restaurant,
                total_cash_out_of_pocket=cash_out,
                total_points_required=points_required,
                benefits_captured=benefits,
                score=score,
            )
        )

    results.sort(key=lambda r: r.score, reverse=True)

    # --- Step 5: Done ---
    _emit(
        progress_callback,
        step="done",
        message="Search complete!",
        progress=100,
    )

    return results