"""Provides the patching mechanism to intercept tqdm and emit callbacks with that progress information"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/patch_tqdm.ipynb.

# %% auto 0
__all__ = ['patch_tqdm']

# %% ../nbs/patch_tqdm.ipynb 3
from contextlib import contextmanager
from importlib import import_module
import time
from typing import Optional, Callable
from .progress_info import ProgressInfo

# %% ../nbs/patch_tqdm.ipynb 4
def _make_callback_class(
    BaseTqdm,  # Base tqdm class to extend with callback functionality
    default_cb: Optional[Callable[[ProgressInfo], None]],
    min_update_interval: float = 0.1,  # Minimum time between callback invocations (seconds)
    min_delta_pct: float = 1.0,      # emit only if pct moves by >= this
    emit_initial: bool = False       # whether to emit at 0%
):
    "Create a tqdm subclass that emits progress callbacks during iteration"
    class CallbackTqdm(BaseTqdm):
        "Extended tqdm class that invokes callbacks with progress information"
        def __init__(self, *args, **kwargs):
            "Initialize CallbackTqdm with callback configuration"
            user_cb = kwargs.pop('progress_callback', None)
            self._progress_callback = user_cb if user_cb is not None else default_cb
            self._min_update_interval = kwargs.pop('min_update_interval', min_update_interval)
            self._min_delta_pct = kwargs.pop('min_delta_pct', min_delta_pct)
            self._emit_initial = kwargs.pop('emit_initial', emit_initial)
            self._last_cb_time = 0.0
            self._last_pct = -1.0
            self._last_n = -1 # remember last emitted n
            self._done_emitted = False
            super().__init__(*args, **kwargs)
    
        def _emit(
            self,
            force: bool = False  # Bypass throttling and emit callback immediately
        ):
            "Emit progress callback with current state if conditions are met"
            cb = self._progress_callback
            if cb is None:
                return
    
            total = self.total or 0
            pct = (self.n / total * 100.0) if total else 0.0
    
            # skip very first 0% unless requested
            if not force and not self._emit_initial and self.n == 0:
                return
    
            # de-dupe same step; also don't re-emit once done
            if not force:
                if self._done_emitted and pct >= 100.0:
                    return
                if self.n == self._last_n:
                    return
    
                now = time.time()
                if (now - self._last_cb_time) < self._min_update_interval and abs(pct - self._last_pct) < self._min_delta_pct:
                    return
                self._last_cb_time = now
                self._last_pct = pct
    
            fd = getattr(self, 'format_dict', {})
            fmt = getattr(self, 'format_interval', None)
            elapsed   = fmt(fd.get("elapsed"))   if fmt and fd.get("elapsed")   is not None else None
            remaining = fmt(fd.get("remaining")) if fmt and fd.get("remaining") is not None else None
            rate = None
            if fd.get("rate") is not None:
                unit = fd.get("unit", "it")
                rate = f"{fd['rate']:.2f} {unit}/s"
    
            cb(ProgressInfo(
                progress=pct,
                current=self.n,
                total=self.total or None,
                rate=rate,
                elapsed=elapsed,
                remaining=remaining,
                description=self.desc or None,
                raw_output=""
            ))
    
            self._last_n = self.n # mark last emitted step
            if pct >= 100.0:
                self._done_emitted = True
    
        def display(self, *a, **k):
            "Update display and emit progress callback"
            out = super().display(*a, **k)
            self._emit()
            return out
    
        def close(self):
            "Close progress bar and emit final callback if needed"
            try:
                super().close()
            finally:
                # only force if we didn't already emit done
                if not self._done_emitted:
                    self._emit(force=True)

    return CallbackTqdm

# %% ../nbs/patch_tqdm.ipynb 5
@contextmanager
def patch_tqdm(
    progress_callback: Optional[Callable[[ProgressInfo], None]],  # Function to call with progress updates
    min_update_interval: float = 0.1,  # Minimum time between callback invocations (seconds)
    min_delta_pct: float = 10.0,   # e.g., only every ~10%
    emit_initial: bool = False  # Whether to emit callback at 0% progress
):
    "Context manager that patches tqdm to emit progress callbacks"
    modules = ['tqdm', 'tqdm.std', 'tqdm.auto', 'tqdm.notebook', 'tqdm.asyncio', 'tqdm.gui']
    originals, patched = {}, []
    try:
        for name in modules:
            try:
                mod = import_module(name)
                if hasattr(mod, 'tqdm'):
                    originals[name] = mod.tqdm
                    mod.tqdm = _make_callback_class(
                        mod.tqdm, progress_callback,
                        min_update_interval=min_update_interval,
                        min_delta_pct=min_delta_pct,
                        emit_initial=emit_initial
                    )
                    patched.append(name)
            except Exception:
                pass
        yield
    finally:
        for name in patched:
            try:
                import_module(name).tqdm = originals[name]
            except Exception:
                pass
