import time
from functools import wraps


def memoize(f):
    @wraps(f)
    def wrapper(obj, *args, **kwargs):
        return obj.memoizer.get(f, *args, **kwargs)
    return wrapper


class MemoizableMixin(object):
    def __init__(self, *args, **kwargs):
        super(MemoizableMixin, self).__init__(*args, **kwargs)
        self.memoizer = MemoizeController(self)


class MemoContext(object):
    def __init__(self, obj):
        self.obj = obj
        self.state = None

    def __enter__(self):
        pass  # pragma: no cover

    def __exit__(self, *args, **kwargs):
        pass  # pragma: no cover


class ClearCachePriorContext(MemoContext):
    __slots__ = ('obj', 'state',)

    def __enter__(self):
        self.obj.clear_cache()


class IgnoreCacheContext(MemoContext):
    __slots__ = ('obj', 'state',)

    def __enter__(self):
        self.state = self.obj.caching
        self.obj.caching = False

    def __exit__(self, *args, **kwargs):
        self.obj.caching = self.state


class ClearCachePostContext(MemoContext):
    __slots__ = ('obj',)

    def __exit__(self, *args, **kwargs):
        self.obj.clear_cache()


class MemoizeController(object):
    __slots__ = ('funcs', 'parent', 'caching',
                 '_mangler', 'clearfirst',
                 'clearafter', 'ignorecache')

    def __init__(self, parent, memoize=True):
        self.funcs = set()
        self.parent = parent
        self.caching = memoize
        self._mangler = '__mem__'

        # helper memoize context managers
        self.clearfirst = ClearCachePriorContext(self)
        self.ignorecache = IgnoreCacheContext(self)
        self.clearafter = ClearCachePostContext(self)

    def get(self, f, *args, **kwargs):
        if self.caching:
            key = self._get_key(f)
            try:
                return self.parent.__dict__[key]
            except KeyError:
                self.funcs.add(f)
                val = f(self.parent, *args, **kwargs)
                self.parent.__dict__[key] = val
            return val
        return f(self.parent, *args, **kwargs)

    def clear_cache(self):
        for key in (self._get_key(f) for f in self.funcs):
            self.parent.__dict__.pop(key, None)

    def _get_key(self, f):
        return "{0}{1}".format(self._mangler, f.__name__)