from typing import Callable, Any, List, TYPE_CHECKING
if TYPE_CHECKING:
    from ding.data.buffer.buffer import Buffer


def staleness_check(buffer_: 'Buffer', max_staleness: int = float("inf")) -> Callable:
    """
    Overview:
        This middleware aims to check staleness before each sample operation,
        staleness = train_iter_sample_data - train_iter_data_collected, means how old/off-policy the data is,
        If data's staleness is greater(>) than max_staleness, this data will be removed from buffer as soon as possible.
    Arguments:
        - max_staleness (:obj:`int`): The maximum legal span between the time of collecting and time of sampling.
    """

    def push(next: Callable, data: Any, *args, **kwargs) -> Any:
        assert 'meta' in kwargs and 'train_iter_data_collected' in kwargs[
            'meta'], "staleness_check middleware must push data with meta={'train_iter_data_collected': <iter>}"
        return next(data, *args, **kwargs)

    def sample(next: Callable, train_iter_sample_data: int, *args, **kwargs) -> List[Any]:
        delete_index = []
        for i, item in enumerate(buffer_.storage):
            index, meta = item.index, item.meta
            staleness = train_iter_sample_data - meta['train_iter_data_collected']
            meta['staleness'] = staleness
            if staleness > max_staleness:
                delete_index.append(index)
        for index in delete_index:
            buffer_.delete(index)
        data = next(*args, **kwargs)
        return data

    def _staleness_check(action: str, next: Callable, *args, **kwargs) -> Any:
        if action == "push":
            return push(next, *args, **kwargs)
        elif action == "sample":
            return sample(next, *args, **kwargs)
        return next(*args, **kwargs)

    return _staleness_check