ciaovolo/flight-comparator/progress.py

"""
Live progress display for flight searches.

Shows a real-time table of search progress with cache hits, API calls, and results.
"""

from datetime import datetime
from collections import defaultdict

try:
    from rich.live import Live
    from rich.table import Table
    from rich.console import Console
    HAS_RICH = True
except ImportError:
    HAS_RICH = False


class SearchProgress:
    """Track and display search progress in real-time."""

    def __init__(self, total_routes: int, show_progress: bool = True):
        self.total_routes = total_routes
        self.show_progress = show_progress
        self.completed = 0
        self.cache_hits = 0
        self.api_calls = 0
        self.errors = 0
        self.flights_found = 0
        self.start_time = datetime.now()

        # Track results by origin airport
        self.results_by_origin = defaultdict(int)

        # For live display
        self.live = None
        self.console = Console() if HAS_RICH else None

    def __enter__(self):
        """Start live display."""
        if self.show_progress and HAS_RICH:
            self.live = Live(self._generate_table(), refresh_per_second=4, console=self.console)
            self.live.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop live display."""
        if self.live:
            self.live.__exit__(exc_type, exc_val, exc_tb)
            # Print final summary
            self._print_summary()

    def update(self, origin: str, destination: str, date: str, status: str, flights_count: int, error: str = None):
        """
        Update progress with search result.

        Args:
            origin: Origin airport code
            destination: Destination airport code
            date: Search date
            status: One of 'cache_hit', 'api_success', 'error'
            flights_count: Number of flights found
            error: Error message if status is 'error'
        """
        self.completed += 1

        if status == "cache_hit":
            self.cache_hits += 1
        elif status == "api_success":
            self.api_calls += 1
        elif status == "error":
            self.errors += 1

        if flights_count > 0:
            self.flights_found += flights_count
            self.results_by_origin[origin] += flights_count

        # Update live display
        if self.live:
            self.live.update(self._generate_table())
        elif self.show_progress and not HAS_RICH:
            # Fallback to simple text progress
            self._print_simple_progress()

    def _generate_table(self) -> Table:
        """Generate Rich table with current progress."""
        table = Table(title="🔍 Flight Search Progress", title_style="bold cyan")

        table.add_column("Metric", style="cyan", no_wrap=True)
        table.add_column("Value", style="green", justify="right")

        # Progress
        progress_pct = (self.completed / self.total_routes * 100) if self.total_routes > 0 else 0
        table.add_row("Progress", f"{self.completed}/{self.total_routes} ({progress_pct:.1f}%)")

        # Cache performance
        table.add_row("💾 Cache Hits", str(self.cache_hits))
        table.add_row("🌐 API Calls", str(self.api_calls))
        if self.errors > 0:
            table.add_row("⚠️  Errors", str(self.errors), style="yellow")

        # Results
        table.add_row("✈️  Flights Found", str(self.flights_found), style="bold green")
        airports_with_flights = len([c for c in self.results_by_origin.values() if c > 0])
        table.add_row("🛫 Airports w/ Flights", str(airports_with_flights))

        # Timing
        elapsed = (datetime.now() - self.start_time).total_seconds()
        table.add_row("⏱️  Elapsed", f"{elapsed:.1f}s")

        if self.completed > 0 and elapsed > 0:
            rate = self.completed / elapsed
            remaining = (self.total_routes - self.completed) / rate if rate > 0 else 0
            table.add_row("⏳ Est. Remaining", f"{remaining:.0f}s")

        return table

    def _print_simple_progress(self):
        """Simple text progress for when Rich is not available."""
        print(f"\rProgress: {self.completed}/{self.total_routes} | "
              f"Cache: {self.cache_hits} | API: {self.api_calls} | "
              f"Flights: {self.flights_found}", end="", flush=True)

    def _print_summary(self):
        """Print final summary."""
        elapsed = (datetime.now() - self.start_time).total_seconds()

        print("\n")
        print("="*60)
        print("SEARCH SUMMARY")
        print("="*60)
        print(f"Total Routes:      {self.total_routes}")
        print(f"Completed:         {self.completed}")
        print(f"Cache Hits:        {self.cache_hits} ({self.cache_hits/self.total_routes*100:.1f}%)")
        print(f"API Calls:         {self.api_calls}")
        print(f"Errors:            {self.errors}")
        print(f"Flights Found:     {self.flights_found}")

        airports_with_flights = len([c for c in self.results_by_origin.values() if c > 0])
        print(f"Airports w/ Flights: {airports_with_flights}")

        print(f"Time Elapsed:      {elapsed:.1f}s")
        print("="*60)
        print()