#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   progress.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Custom progress bar with integrated console redirection.
"""

import sys
import time
from typing import Any

from rich.console import Console
from rich.progress import Progress
from vi.utils.console_utils import redirect_console


def is_jupyter() -> bool:
    """Detect if running in Jupyter notebook/lab environment.

    Returns:
        True if running in Jupyter, False otherwise.

    """
    try:
        # Check for IPython shell
        shell = get_ipython().__class__.__name__  # type: ignore[name-defined]
        # ZMQInteractiveShell is used by Jupyter notebooks and lab
        return shell == "ZMQInteractiveShell"
    except NameError:
        # get_ipython is not defined, not in IPython/Jupyter
        return False


class ViProgress(Progress):
    """Enhanced progress bar that automatically redirects console output to Rich formatting.

    This class combines Rich's Progress functionality with automatic console redirection,
    ensuring that all stdout, stderr, and logging output during progress operations
    is captured and displayed using Rich formatting.

    When used as a context manager, captured output from external libraries is displayed
    in real-time on new lines without disrupting the progress spinner.

    Console redirection uses a non-blocking background thread to process output,
    minimizing performance impact on the main thread operations.
    """

    def __init__(
        self,
        *args: Any,
        redirect_console_output: bool = False,
        **kwargs: Any,
    ):
        """Initialize ViProgress with optional console redirection.

        Args:
            *args: Arguments passed to Rich Progress
            redirect_console_output: Whether to enable console redirection (default: False).
                When True, stdout/stderr/logging output is captured and displayed
                asynchronously in a background thread without blocking the main thread.
                This is useful during model loading but typically not needed during inference.
            **kwargs: Keyword arguments passed to Rich Progress

        Note:
            Console redirection is disabled by default for performance. Enable it
            only when you need to capture output from external libraries (e.g.,
            during model loading with transformers).

        """
        if "console" not in kwargs:
            # Force terminal mode for inline ANSI output in Jupyter (no white background)
            kwargs["console"] = Console(
                file=sys.stderr,
                force_terminal=True,
                force_jupyter=False,
                legacy_windows=False,
            )

        self._redirect_console = redirect_console_output
        self._console_context = None
        # Store transient parameter for later use in __exit__
        self._is_transient = kwargs.get("transient", False)
        super().__init__(*args, **kwargs)

    def __enter__(self):
        """Enter context manager and set up display."""
        # Start the progress display
        progress_instance = super().__enter__()

        # Set up console redirection with non-blocking async output to our console
        if self._redirect_console:
            self._console_context = redirect_console(rich_console=self.console)
            self._console_context.__enter__()

        return progress_instance

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager and clean up console redirection."""
        # Clean up console redirection first (will flush remaining output)
        if self._redirect_console and self._console_context is not None:
            self._console_context.__exit__(exc_type, exc_val, exc_tb)

        # Check if this is a transient progress bar
        is_transient = self._is_transient

        # Then stop the progress display
        result = super().__exit__(exc_type, exc_val, exc_tb)

        # Explicitly flush the console to ensure clean output
        # This is especially important for transient progress bars
        # and helps prevent overlapping in Jupyter notebooks
        self.console.file.flush()

        # In Jupyter notebooks, transient progress bars need a small delay
        # to properly clear before the next progress bar starts
        if is_transient and is_jupyter():
            time.sleep(0.1)

        return result
