"""Extensible callback system for augmenting objects and tracking iterations"""

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

# %% ../nbs/10_callback.ipynb 1
from __future__ import annotations

# %% auto 0
__all__ = ['Callback', 'run_cbs', 'EchoCB', 'FuncCB', 'PassCB', 'HasCallbacks', 'with_cbs', 'CollBack', 'trackback', 'process_']

# %% ../nbs/10_callback.ipynb
import collections
import time
from contextlib import contextmanager
from functools import partial
from operator import attrgetter
from operator import length_hint
from typing import Any
from typing import Callable
from typing import Iterable
from typing import Iterator
from typing import Sequence
from typing import Type
from typing import TypeVar

import fastcore.all as FC


# %% ../nbs/10_callback.ipynb
from .basic import AD
from .basic import empty
from .basic import EmptyT
from .basic import update_
from .common import Runner
from .common import setattrs


# %% ../nbs/10_callback.ipynb
class Callback:
    "Base class of callbacks."
    order = 0
    cbs: Sequence[Callback]  # if present, run callbacks in the list after running this callback

# %% ../nbs/10_callback.ipynb
def run_cbs(cbs: Iterable[Callback] | FC.L, method_nm:str, ctx=None, *args, **kwargs):
    "Run `method_nm(ctx, ...)` of each callback in `cbs` in order."
    for cb in sorted(cbs, key=attrgetter('order')):
        if f := getattr(cb, method_nm, None): Runner(f)(ctx, *args, **kwargs)
        if nested := getattr(cb, 'cbs', None): 
            run_cbs(filter(lambda x: x != cb, nested), method_nm, ctx, *args, **kwargs)

# %% ../nbs/10_callback.ipynb
class EchoCB(Callback):
    "Print the arguments."
    cbs=()
    def echo(self, ctx, *args, **kwargs): print(ctx, args, kwargs)
    def __getattr__(self, name): return partial(self.echo, name)

# %% ../nbs/10_callback.ipynb
class FuncCB(Callback):
    "Store functions that can be called as callbacks."
    def __init__(self, **kwargs): setattrs(self, kwargs)

# %% ../nbs/10_callback.ipynb
class PassCB(Callback):
    "A callback that does nothing."
    def noop(self, *args, **kwargs): pass
    def __getattr__(self, name): return self.noop

# %% ../nbs/10_callback.ipynb
class HasCallbacks:
    "Base for classes that can be augmented with callbacks."
    cbs: list[Callback]
    cbs_names: tuple[str,...] = ()  # this callbacks becomes self methods if a subclass overrides them

    def __new__(cls, *args, **kwargs):
        self = super().__new__(cls)
        self.cbs = []#PassCB()]
        return self

    def __init__(self, cbs:Sequence[Callback]=()): self.cbs = list(cbs)

    def with_cbs(self, cbs:Sequence[Callback], extend=False):
        if extend: self.cbs.extend(cbs)
        else: self.cbs = list(cbs)
        return self
    
    @contextmanager
    def this_cbs(self, cbs:Sequence[Callback]):
        "Use temporary `cbs` in `with` block."
        for cb in cbs: self.cbs.append(cb)
        try:
            yield        
        finally:
            for cb in cbs: self.cbs.remove(cb)

    def __getattr__(self, name):
        "Run `name` as a callback if it's in `self.cbs_names`."
        if name in self.cbs_names: return partial(self.callback, name)
        raise AttributeError(name)

    def callback(self, method_nm, *args, **kwargs): run_cbs(self.cbs, method_nm, self, *args, **kwargs)

# %% ../nbs/10_callback.ipynb
class with_cbs:
    "Decorator to add callbacks to a method."
    def __init__(self, nm:str|None=None): self.nm = nm
    def __call__(self, f):
        def _f(o, *args, **kwargs):
            nm = self.nm or f.__name__
            try:
                o.callback(f'before_{nm}')
                f(o, *args, **kwargs)
                o.callback(f'after_{nm}')
            except globals()[f'Cancel{nm.title()}Exception']: pass
            finally: o.callback(f'cleanup_{nm}')
        return _f

# %% ../nbs/10_callback.ipynb
def _get_total(total: int|None|Type[EmptyT], source) -> int|None:
    if total is empty:
        try: return len(source)
        except: return length_hint(source) or None
    if total is not None and (not isinstance(total, int) or total < 0): total = None
    return total

# %% ../nbs/10_callback.ipynb
class CollBack(HasCallbacks):
    "Track iterables and extend them with callbacks."

    cbs_names = ('before_iter', 'after_iter', 'on_iter', 'on_interrupt')

    n: int|None = None
    elapsed_time: float|None = None
    active, interrupted = False, False
    item: Any = empty
    def __init__(self, 
            source: Iterable[Any] = (), 
            total: int|None|Type[EmptyT] = empty, 
            context: Any = empty,
            **kwargs):
        self._source, self._total = source, total
        self.total: int|None = _get_total(self._total, self._source)
        self.context = context
        super().__init__(**kwargs)

    @property
    def progress(self):
        if self.total:
            return None if self.n is None else min(1., round((self.n+1)/float(self.total), 4))
        return None

    @property
    def state(self):
        return AD(update_(item=self.item, n=self.n, total=self.total, progress=self.progress, 
            context=self.context, 
            elapsed_time=self.elapsed_time, interrupted=self.interrupted or empty, empty_value=empty))

    def __repr__(self): return f'{self.__class__.__name__}#{self._source}, total={self._total}'

    def _start(self):
        self.total = _get_total(self._total, self._source)
        self.active, self.elapsed_time, self.n = True, None, None
        run_cbs(self.cbs, 'before_iter', self.state)

    def _stop(self):
        if self.total is None and self.n is not None: self.total = self.n + 1
        self.active, self.item = False, empty
        run_cbs(self.cbs, 'after_iter', self.state)
    
    def _interrupt(self): self.interrupted = True; run_cbs(self.cbs, 'on_interrupt', self.state)

    def __iter__(self) -> Iterator[Any]:
        if self._source is None: return
        try:
            start_time = time.time()
            self._start()
            for self.n, self.item in enumerate(self._source):
                if self.total is not None and self.n >= self.total: break
                yield self.item
                self.elapsed_time = time.time() - start_time
                run_cbs(self.cbs, 'on_iter', self.state, self.item)
                if self.total is not None and self.n >= self.total-1: break
        except Exception as e: self._interrupt(); raise e
        finally:
            self._stop()

    def trackback(self, cbs:Sequence[Callback]=()) -> Iterator[tuple[AD[Any], Any]]:
        with self.this_cbs(cbs):
            for elem in self:
                yield self.state, elem

# %% ../nbs/10_callback.ipynb
def trackback(source: Iterable[Any], total: int|None|Type[EmptyT]=empty, context: Any=empty, 
        cbs:Sequence[Callback]=()) -> Iterator[tuple[AD[Any], Any]]:
    return CollBack(source, total, context, cbs=cbs).trackback()

# %% ../nbs/10_callback.ipynb
_T = TypeVar('_T')

def process_(
        iterable:Iterable[_T], /, 
        cbs: Callback|Sequence[Callback]=(), 
        slc:slice|None=None, 
        pred:Callable[[_T], bool]|None=None, 
        context:Any=empty,
        **kwargs  # FuncCB kwargs
    ) -> tuple[Callback,...]:
    "Process a subset `slc` of `iterable` filtered by `pred` with callbacks from `cbs` and `FuncCB` `kwargs`"
    _cbs = FC.tuplify(cbs) + ((FuncCB(**kwargs),) if kwargs else ())
    items = FC.L(iterable)[slc or slice(None)].filter(pred)  # type: ignore
    collections.deque(CollBack(items, context=context,cbs=_cbs), maxlen=0)
    return _cbs  # type: ignore
