import datetime
import functools
from typing import Callable, Optional


def timed_cache(timedelta: datetime.timedelta,
                maxsize: Optional[int] = 128,
                typed: Optional[bool] = False) -> Callable:
    """Timeout is applied to the whole cache."""
    def decorator(func):
        # Apply @lru_cache to func
        func = functools.lru_cache(maxsize=maxsize, typed=typed)(func)

        next_update = datetime.datetime.utcnow() - timedelta

        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            nonlocal next_update

            now = datetime.datetime.utcnow()
            if now > next_update:
                func.cache_clear()
                next_update = now + timedelta
            return func(*args, **kwargs)

        wrapped.cache_info = func.cache_info
        wrapped.cache_clear = func.cache_clear

        return wrapped

    return decorator


def granular_timed_cache(timedelta: datetime.timedelta,
                         maxsize: Optional[int] = 128,
                         typed: Optional[bool] = False) -> Callable:
    """The timeout is applied to the cache of each individual result."""
    def decorator(func):
        _args_ = {}

        # Apply @lru_cache to new func
        @functools.lru_cache(maxsize=maxsize, typed=typed)
        def _wrapped(expiration, *args, **kwargs):
            return func(*args, **kwargs)

        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            hashable_args = tuple(args) + tuple((k, v)
                                                for k, v in kwargs.items())

            expiration = _args_.get(hashable_args, None)

            now = datetime.datetime.utcnow()
            if not expiration or now >= expiration:
                expiration = now + timedelta
                _args_[hashable_args] = expiration

            return _wrapped(expiration, *args, **kwargs)

        wrapped.cache_info = _wrapped.cache_info
        wrapped.cache_clear = _wrapped.cache_clear

        return wrapped

    return decorator


if __name__ == '__main__':
    import time

    for decorator in timed_cache, granular_timed_cache:
        @decorator(datetime.timedelta(seconds=5))
        def f(num):
            print(f'The function was called with argument {num}')
            return num

        for i in range(10):
            print(decorator.__name__, f(i % 2))
            time.sleep(1)