ciaovolo/flight-comparator/date_resolver.py

"""
Date resolution and seasonal scan logic for flight comparator.
"""

from datetime import date, timedelta
from dateutil.relativedelta import relativedelta
from typing import Optional

# Primary configuration constants
SEARCH_WINDOW_MONTHS = 6  # Default seasonal scan window
SAMPLE_DAY_OF_MONTH = 15  # Representative mid-month date for seasonal queries


def resolve_dates(date_arg: Optional[str], window: int) -> list[str]:
    """
    Resolve query dates based on CLI input.

    Returns a single date if --date is provided; otherwise generates one date
    per month across the window for seasonal scanning.

    Args:
        date_arg: Optional date string in YYYY-MM-DD format
        window: Number of months to scan (only used if date_arg is None)

    Returns:
        List of date strings in YYYY-MM-DD format
    """
    if date_arg:
        return [date_arg]

    today = date.today()
    dates = []

    for i in range(1, window + 1):
        # Generate dates starting from next month
        target_date = today + relativedelta(months=i)
        # Set to the sample day of month (default: 15th)
        try:
            target_date = target_date.replace(day=SAMPLE_DAY_OF_MONTH)
        except ValueError:
            # Handle months with fewer days (e.g., February with day=31)
            # Use the last day of the month instead
            target_date = target_date.replace(day=1) + relativedelta(months=1) - relativedelta(days=1)

        dates.append(target_date.strftime('%Y-%m-%d'))

    return dates


def resolve_dates_daily(
    start_date: Optional[str],
    end_date: Optional[str],
    window: int,
) -> list[str]:
    """
    Generate a list of ALL dates in a range, day-by-day (Monday-Sunday).

    This is used for comprehensive daily scans to catch flights that only operate
    on specific days of the week (e.g., Saturday-only routes).

    Args:
        start_date: Optional start date in YYYY-MM-DD format
        end_date: Optional end date in YYYY-MM-DD format
        window: Number of months to scan if dates not specified

    Returns:
        List of date strings in YYYY-MM-DD format, one per day

    Examples:
        # Scan next 3 months daily
        resolve_dates_daily(None, None, 3)

        # Scan specific range
        resolve_dates_daily("2026-04-01", "2026-06-30", 3)
    """
    today = date.today()

    # Determine start date
    if start_date:
        start = date.fromisoformat(start_date)
    else:
        # Start from tomorrow (or first day of next month)
        start = today + timedelta(days=1)

    # Determine end date
    if end_date:
        end = date.fromisoformat(end_date)
    else:
        # End after window months
        end = today + relativedelta(months=window)

    # Generate all dates in range
    dates = []
    current = start
    while current <= end:
        dates.append(current.strftime('%Y-%m-%d'))
        current += timedelta(days=1)

    return dates


def detect_new_connections(monthly_results: dict[str, list]) -> dict[str, str]:
    """
    Detect routes that appear for the first time in later months.

    Compares route sets month-over-month and tags routes that weren't present
    in any previous month.

    Args:
        monthly_results: Dict mapping month strings to lists of flight objects
                        Each flight must have 'origin' and 'destination' attributes

    Returns:
        Dict mapping route keys (e.g., "FRA->JFK") to the first month they appeared
    """
    seen = set()
    new_connections = {}

    # Process months in chronological order
    for month in sorted(monthly_results.keys()):
        flights = monthly_results[month]
        current = {f"{f['origin']}->{f['destination']}" for f in flights}

        # Find routes in current month that weren't in any previous month
        for route in current - seen:
            if seen:  # Only tag as NEW if it's not the very first month
                new_connections[route] = month

        # Add all current routes to seen set
        seen |= current

    return new_connections