# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

from __future__ import annotations

import os
import sys
import time
from functools import lru_cache
from typing import IO, Any


@lru_cache(maxsize=1)
def is_noninteractive_console() -> bool:
    """Check for known non-interactive console environments."""
    return "GITHUB_ACTIONS" in os.environ or "RUNPOD_POD_ID" in os.environ


class TQDM:
    """
    Lightweight zero-dependency progress bar for Ultralytics.

    Provides clean, rich-style progress bars suitable for various environments including Weights & Biases,
    console outputs, and other logging systems. Features zero external dependencies, clean single-line output,
    rich-style progress bars with Unicode block characters, context manager support, iterator protocol support,
    and dynamic description updates.

    Attributes:
        iterable (object): Iterable to wrap with progress bar.
        desc (str): Prefix description for the progress bar.
        total (int): Expected number of iterations.
        disable (bool): Whether to disable the progress bar.
        unit (str): String for units of iteration.
        unit_scale (bool): Auto-scale units flag.
        unit_divisor (int): Divisor for unit scaling.
        leave (bool): Whether to leave the progress bar after completion.
        mininterval (float): Minimum time interval between updates.
        initial (int): Initial counter value.
        n (int): Current iteration count.
        closed (bool): Whether the progress bar is closed.
        bar_format (str): Custom bar format string.
        file (object): Output file stream.

    Methods:
        update: Update progress by n steps.
        set_description: Set or update the description.
        set_postfix: Set postfix for the progress bar.
        close: Close the progress bar and clean up.
        refresh: Refresh the progress bar display.
        clear: Clear the progress bar from display.
        write: Write a message without breaking the progress bar.

    Examples:
        Basic usage with iterator:
        >>> for i in TQDM(range(100)):
        ...     time.sleep(0.01)

        With custom description:
        >>> pbar = TQDM(range(100), desc="Processing")
        >>> for i in pbar:
        ...     pbar.set_description(f"Processing item {i}")

        Context manager usage:
        >>> with TQDM(total=100, unit="B", unit_scale=True) as pbar:
        ...     for i in range(100):
        ...         pbar.update(1)

        Manual updates:
        >>> pbar = TQDM(total=100, desc="Training")
        >>> for epoch in range(100):
        ...     # Do work
        ...     pbar.update(1)
        >>> pbar.close()
    """

    # Constants
    MIN_RATE_CALC_INTERVAL = 0.01  # Minimum time interval for rate calculation
    RATE_SMOOTHING_FACTOR = 0.3  # Factor for exponential smoothing of rates
    MAX_SMOOTHED_RATE = 1000000  # Maximum rate to apply smoothing to
    NONINTERACTIVE_MIN_INTERVAL = 60.0  # Minimum interval for non-interactive environments

    def __init__(
        self,
        iterable: Any = None,
        desc: str | None = None,
        total: int | None = None,
        leave: bool = True,
        file: IO[str] | None = None,
        mininterval: float = 0.1,
        disable: bool | None = None,
        unit: str = "it",
        unit_scale: bool = True,
        unit_divisor: int = 1000,
        bar_format: str | None = None,  # kept for API compatibility; not used for formatting
        initial: int = 0,
        **kwargs,
    ) -> None:
        """
        Initialize the TQDM progress bar with specified configuration options.

        Args:
            iterable (object, optional): Iterable to wrap with progress bar.
            desc (str, optional): Prefix description for the progress bar.
            total (int, optional): Expected number of iterations.
            leave (bool, optional): Whether to leave the progress bar after completion.
            file (object, optional): Output file stream for progress display.
            mininterval (float, optional): Minimum time interval between updates (default 0.1s, 60s in GitHub Actions).
            disable (bool, optional): Whether to disable the progress bar. Auto-detected if None.
            unit (str, optional): String for units of iteration (default "it" for items).
            unit_scale (bool, optional): Auto-scale units for bytes/data units.
            unit_divisor (int, optional): Divisor for unit scaling (default 1000).
            bar_format (str, optional): Custom bar format string.
            initial (int, optional): Initial counter value.
            **kwargs (Any): Additional keyword arguments for compatibility (ignored).

        Examples:
            >>> pbar = TQDM(range(100), desc="Processing")
            >>> with TQDM(total=1000, unit="B", unit_scale=True) as pbar:
            ...     pbar.update(1024)  # Updates by 1KB
        """
        # Disable if not verbose
        if disable is None:
            try:
                from ultralytics.utils import LOGGER, VERBOSE

                disable = not VERBOSE or LOGGER.getEffectiveLevel() > 20
            except ImportError:
                disable = False

        self.iterable = iterable
        self.desc = desc or ""
        self.total = total or (len(iterable) if hasattr(iterable, "__len__") else None) or None  # prevent total=0
        self.disable = disable
        self.unit = unit
        self.unit_scale = unit_scale
        self.unit_divisor = unit_divisor
        self.leave = leave
        self.noninteractive = is_noninteractive_console()
        self.mininterval = max(mininterval, self.NONINTERACTIVE_MIN_INTERVAL) if self.noninteractive else mininterval
        self.initial = initial

        # Kept for API compatibility (unused for f-string formatting)
        self.bar_format = bar_format

        self.file = file or sys.stdout

        # Internal state
        self.n = self.initial
        self.last_print_n = self.initial
        self.last_print_t = time.time()
        self.start_t = time.time()
        self.last_rate = 0.0
        self.closed = False
        self.is_bytes = unit_scale and unit in ("B", "bytes")
        self.scales = (
            [(1073741824, "GB/s"), (1048576, "MB/s"), (1024, "KB/s")]
            if self.is_bytes
            else [(1e9, f"G{self.unit}/s"), (1e6, f"M{self.unit}/s"), (1e3, f"K{self.unit}/s")]
        )

        if not self.disable and self.total and not self.noninteractive:
            self._display()

    def _format_rate(self, rate: float) -> str:
        """Format rate with units."""
        if rate <= 0:
            return ""
        fallback = f"{rate:.1f}B/s" if self.is_bytes else f"{rate:.1f}{self.unit}/s"
        return next((f"{rate / t:.1f}{u}" for t, u in self.scales if rate >= t), fallback)

    def _format_num(self, num: int | float) -> str:
        """Format number with optional unit scaling."""
        if not self.unit_scale or not self.is_bytes:
            return str(num)

        for unit in ("", "K", "M", "G", "T"):
            if abs(num) < self.unit_divisor:
                return f"{num:3.1f}{unit}B" if unit else f"{num:.0f}B"
            num /= self.unit_divisor
        return f"{num:.1f}PB"

    def _format_time(self, seconds: float) -> str:
        """Format time duration."""
        if seconds < 60:
            return f"{seconds:.1f}s"
        elif seconds < 3600:
            return f"{int(seconds // 60)}:{seconds % 60:02.0f}"
        else:
            h, m = int(seconds // 3600), int((seconds % 3600) // 60)
            return f"{h}:{m:02d}:{seconds % 60:02.0f}"

    def _generate_bar(self, width: int = 12) -> str:
        """Generate progress bar."""
        if self.total is None:
            return "━" * width if self.closed else "─" * width

        frac = min(1.0, self.n / self.total)
        filled = int(frac * width)
        bar = "━" * filled + "─" * (width - filled)
        if filled < width and frac * width - filled > 0.5:
            bar = f"{bar[:filled]}╸{bar[filled + 1 :]}"
        return bar

    def _should_update(self, dt: float, dn: int) -> bool:
        """Check if display should update."""
        if self.noninteractive:
            return False
        return (self.total is not None and self.n >= self.total) or (dt >= self.mininterval)

    def _display(self, final: bool = False) -> None:
        """Display progress bar."""
        if self.disable or (self.closed and not final):
            return

        current_time = time.time()
        dt = current_time - self.last_print_t
        dn = self.n - self.last_print_n

        if not final and not self._should_update(dt, dn):
            return

        # Calculate rate (avoid crazy numbers)
        if dt > self.MIN_RATE_CALC_INTERVAL:
            rate = dn / dt if dt else 0.0
            # Smooth rate for reasonable values, use raw rate for very high values
            if rate < self.MAX_SMOOTHED_RATE:
                self.last_rate = self.RATE_SMOOTHING_FACTOR * rate + (1 - self.RATE_SMOOTHING_FACTOR) * self.last_rate
                rate = self.last_rate
        else:
            rate = self.last_rate

        # At completion, use overall rate
        if self.total and self.n >= self.total:
            overall_elapsed = current_time - self.start_t
            if overall_elapsed > 0:
                rate = self.n / overall_elapsed

        # Update counters
        self.last_print_n = self.n
        self.last_print_t = current_time
        elapsed = current_time - self.start_t

        # Remaining time
        remaining_str = ""
        if self.total and 0 < self.n < self.total and elapsed > 0:
            est_rate = rate or (self.n / elapsed)
            remaining_str = f"<{self._format_time((self.total - self.n) / est_rate)}"

        # Numbers and percent
        if self.total:
            percent = (self.n / self.total) * 100
            n_str = self._format_num(self.n)
            t_str = self._format_num(self.total)
            if self.is_bytes:
                # Collapse suffix only when identical (e.g. "5.4/5.4MB")
                if n_str[-2] == t_str[-2]:
                    n_str = n_str.rstrip("KMGTPB")  # Remove unit suffix from current if different than total
        else:
            percent = 0.0
            n_str, t_str = self._format_num(self.n), "?"

        elapsed_str = self._format_time(elapsed)
        rate_str = self._format_rate(rate) or (self._format_rate(self.n / elapsed) if elapsed > 0 else "")

        bar = self._generate_bar()

        # Compose progress line via f-strings (two shapes: with/without total)
        if self.total:
            if self.is_bytes and self.n >= self.total:
                # Completed bytes: show only final size
                progress_str = f"{self.desc}: {percent:.0f}% {bar} {t_str} {rate_str} {elapsed_str}"
            else:
                progress_str = (
                    f"{self.desc}: {percent:.0f}% {bar} {n_str}/{t_str} {rate_str} {elapsed_str}{remaining_str}"
                )
        else:
            progress_str = f"{self.desc}: {bar} {n_str} {rate_str} {elapsed_str}"

        # Write to output
        try:
            if self.noninteractive:
                # In non-interactive environments, avoid carriage return which creates empty lines
                self.file.write(progress_str)
            else:
                # In interactive terminals, use carriage return and clear line for updating display
                self.file.write(f"\r\033[K{progress_str}")
            self.file.flush()
        except Exception:
            pass

    def update(self, n: int = 1) -> None:
        """Update progress by n steps."""
        if not self.disable and not self.closed:
            self.n += n
            self._display()

    def set_description(self, desc: str | None) -> None:
        """Set description."""
        self.desc = desc or ""
        if not self.disable:
            self._display()

    def set_postfix(self, **kwargs: Any) -> None:
        """Set postfix (appends to description)."""
        if kwargs:
            postfix = ", ".join(f"{k}={v}" for k, v in kwargs.items())
            base_desc = self.desc.split(" | ")[0] if " | " in self.desc else self.desc
            self.set_description(f"{base_desc} | {postfix}")

    def close(self) -> None:
        """Close progress bar."""
        if self.closed:
            return

        self.closed = True

        if not self.disable:
            # Final display
            if self.total and self.n >= self.total:
                self.n = self.total
            self._display(final=True)

            # Cleanup
            if self.leave:
                self.file.write("\n")
            else:
                self.file.write("\r\033[K")

            try:
                self.file.flush()
            except Exception:
                pass

    def __enter__(self) -> TQDM:
        """Enter context manager."""
        return self

    def __exit__(self, *args: Any) -> None:
        """Exit context manager and close progress bar."""
        self.close()

    def __iter__(self) -> Any:
        """Iterate over the wrapped iterable with progress updates."""
        if self.iterable is None:
            raise TypeError("'NoneType' object is not iterable")

        try:
            for item in self.iterable:
                yield item
                self.update(1)
        finally:
            self.close()

    def __del__(self) -> None:
        """Destructor to ensure cleanup."""
        try:
            self.close()
        except Exception:
            pass

    def refresh(self) -> None:
        """Refresh display."""
        if not self.disable:
            self._display()

    def clear(self) -> None:
        """Clear progress bar."""
        if not self.disable:
            try:
                self.file.write("\r\033[K")
                self.file.flush()
            except Exception:
                pass

    @staticmethod
    def write(s: str, file: IO[str] | None = None, end: str = "\n") -> None:
        """Static method to write without breaking progress bar."""
        file = file or sys.stdout
        try:
            file.write(s + end)
            file.flush()
        except Exception:
            pass


if __name__ == "__main__":
    import time

    print("1. Basic progress bar with known total:")
    for i in TQDM(range(3), desc="Known total"):
        time.sleep(0.05)

    print("\n2. Manual updates with known total:")
    pbar = TQDM(total=300, desc="Manual updates", unit="files")
    for i in range(300):
        time.sleep(0.03)
        pbar.update(1)
        if i % 10 == 9:
            pbar.set_description(f"Processing batch {i // 10 + 1}")
    pbar.close()

    print("\n3. Progress bar with unknown total:")
    pbar = TQDM(desc="Unknown total", unit="items")
    for i in range(25):
        time.sleep(0.08)
        pbar.update(1)
        if i % 5 == 4:
            pbar.set_postfix(processed=i + 1, status="OK")
    pbar.close()

    print("\n4. Context manager with unknown total:")
    with TQDM(desc="Processing stream", unit="B", unit_scale=True, unit_divisor=1024) as pbar:
        for i in range(30):
            time.sleep(0.1)
            pbar.update(1024 * 1024 * i)  # Simulate processing MB of data

    print("\n5. Iterator with unknown length:")

    def data_stream():
        """Simulate a data stream of unknown length."""
        import random

        for i in range(random.randint(10, 20)):
            yield f"data_chunk_{i}"

    for chunk in TQDM(data_stream(), desc="Stream processing", unit="chunks"):
        time.sleep(0.1)

    print("\n6. File processing simulation (unknown size):")

    def process_files():
        """Simulate processing files of unknown count."""
        return [f"file_{i}.txt" for i in range(18)]

    pbar = TQDM(desc="Scanning files", unit="files")
    files = process_files()
    for i, filename in enumerate(files):
        time.sleep(0.06)
        pbar.update(1)
        pbar.set_description(f"Processing {filename}")
    pbar.close()