import sys
import time
from typing import Optional, Union


class ProgressLoader:
    """
    Enhanced progress bar with timing, ETA, and customizable display options.

    This class provides a more robust progress tracking system compared to the
    original bs.loader function, with features like elapsed time, ETA estimation,
    custom messages, and automatic cleanup.
    """

    def __init__(self, total: int, desc: str = "Progress", width: int = 50,
                 show_eta: bool = True, show_rate: bool = True):
        """
        Initialize the progress loader.

        Parameters:
        -----------
        total : int
            Total number of iterations expected
        desc : str
            Description to show before the progress bar
        width : int
            Width of the progress bar in characters
        show_eta : bool
            Whether to show estimated time of arrival
        show_rate : bool
            Whether to show processing rate (items/sec)
        """
        self.total = total
        self.desc = desc
        self.width = width
        self.show_eta = show_eta
        self.show_rate = show_rate
        self.start_time = time.time()
        self.current = 0
        self.last_update_time = self.start_time

    def update(self, current: Optional[int] = None, message: str = ""):
        """
        Update the progress bar.

        Parameters:
        -----------
        current : int, optional
            Current progress value. If None, increments by 1
        message : str, optional
            Additional message to display
        """
        if current is not None:
            self.current = current
        else:
            self.current += 1

        # Limit update frequency to avoid performance issues
        current_time = time.time()
        if current_time - self.last_update_time < 0.1 and self.current < self.total:
            return
        self.last_update_time = current_time

        self._display(message)

    def _display(self, message: str = ""):
        """Display the progress bar."""
        if self.total == 0:
            return

        progress = min(self.current / self.total, 1.0)
        filled_width = int(self.width * progress)

        # Create the bar
        bar = '█' * filled_width + '░' * (self.width - filled_width)

        # Calculate timing information
        elapsed = time.time() - self.start_time

        # Build the display string
        display_parts = [
            f"\r{self.desc}: ",
            f"{bar} ",
            f"{self.current}/{self.total} ",
            f"({progress*100:.1f}%) "
        ]

        # Add elapsed time
        display_parts.append(f"[{self._format_time(elapsed)}]")

        # Add ETA if requested and we have enough data
        if self.show_eta and self.current > 0 and self.current < self.total:
            eta = (elapsed / self.current) * (self.total - self.current)
            display_parts.append(f" ETA: {self._format_time(eta)}")

        # Add rate if requested
        if self.show_rate and elapsed > 0:
            rate = self.current / elapsed
            display_parts.append(f" Rate: {rate:.1f}/s")

        # Add custom message
        if message:
            display_parts.append(f" | {message}")

        # Write to stdout
        sys.stdout.write(''.join(display_parts))
        sys.stdout.flush()

        # Add newline when complete
        if self.current >= self.total:
            sys.stdout.write('\n')

    def _format_time(self, seconds: float) -> str:
        """Format time in a human-readable way."""
        if seconds < 60:
            return f"{seconds:.1f}s"
        elif seconds < 3600:
            minutes = int(seconds // 60)
            secs = int(seconds % 60)
            return f"{minutes}m{secs:02d}s"
        else:
            hours = int(seconds // 3600)
            minutes = int((seconds % 3600) // 60)
            return f"{hours}h{minutes:02d}m"

    def finish(self, message: str = "Complete"):
        """Finish the progress bar with a final message."""
        self.current = self.total
        self._display(message)


def loader(current: int, total: int, desc: str = "Progress",
           width: int = 30, show_details: bool = True) -> None:
    """
    Enhanced version of the original bs.loader function with better formatting and timing.

    This function provides a drop-in replacement for the original loader with improved
    visual appearance, timing information, and ETA estimation while maintaining
    backward compatibility.

    Parameters:
    -----------
    current : int
        Current iteration number
    total : int  
        Total number of iterations
    desc : str, optional
        Description text to show (default: "Progress")
    width : int, optional
        Width of progress bar in characters (default: 30)
    show_details : bool, optional
        Whether to show timing and rate information (default: True)
    """
    if total == 0:
        return

    progress = min(current / total, 1.0)
    filled_width = int(width * progress)
    
    if current == 1 and total == 1:
        return

    # Create progress bar with better visual elements
    bar = '█' * filled_width + '░' * (width - filled_width)

    # Basic display
    display = f"\r{desc}: {bar} {current}/{total} ({progress*100:.1f}%)"

    # Add timing information if requested
    if show_details and hasattr(loader, '_start_time'):
        elapsed = time.time() - loader._start_time
        if current > 0:
            eta = (elapsed / current) * \
                (total - current) if current < total else 0
            rate = current / elapsed if elapsed > 0 else 0
            display += f" | {elapsed:.1f}s elapsed"
            display += f" | ETA: {eta:.1f}s"
            if rate > 0:
                display += f" | {rate:.1f}/s"
    elif show_details and current == 0:
        # Initialize timing on first call
        loader._start_time = time.time()

    sys.stdout.write(display)
    sys.stdout.flush()

    # Add newline when complete
    if current >= total:
        sys.stdout.write('\n')
        if hasattr(loader, '_start_time'):
            delattr(loader, '_start_time')


# Context manager version for cleaner usage
class progress_loader:
    """
    Context manager version of the progress loader for cleaner usage patterns.

    Usage:
    ------
    with progress_loader(total_items, "Processing") as pbar:
        for i, item in enumerate(items):
            # Do work here
            pbar.update(i + 1, f"Processing {item}")
    """

    def __init__(self, total: int, desc: str = "Progress", **kwargs):
        self.loader = ProgressLoader(total, desc, **kwargs)

    def __enter__(self):
        return self.loader

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.loader.finish()
        else:
            self.loader.finish("Failed")
