"""
This module provides the 'cache' decorator - an LRU cache that works very similarly to the 'functools.lru_cache', but
with a couple of important differences:
- The 'ttl_s' paramater (time-to-live): allows cached values to expire after the given time, which means that a function
  call after this time forces the function to be re-evaluated. Note that the stale cache values will persist in memory
  until a successive call refreshes them.
- Better support for multi-threading: if two threads simultaneously make the same call with the same parameters, only
  one of the threads will actually do the work of evaluating the function. The other thread(s) will wait on that result
  and return it, or raise any exception that was thrown (though note that exceptions will not be cached).
Arguments to a cached function must be hashable (no dictionaries).

Parameters:
  ttl_s: [optional] that amount of time after which a cached value is no longer valid, or 'None' to cache values forever
    (defaults to 'None')
  max_size: [optional] the max number of items to store in the cache before we start eviction, or 'None' for no limit
    (defaults to 'None')

This cache decorator also attaches a 'cache' property to the decorated function, which has the following properties:
  ttl_s: the ttl_s the cache was created with
  max_size: the max_size the cache was created with
  func: a reference to the original underlying function (calling directly will not impact the cache metrics)
  get: function that takes the same input as the wrapped function but returns the CacheEntry instead of the raw value
  contents: on OrderedDict containing the current contents of the cache
  evict: function that takes the same input as the wrapped function and evicts any cached result
  info: returns a CacheInfo object summarizing the the cache



# <op-test>

from op_service.cache import cache
from threading import Lock, Thread

class Work():
  def __init__(self):
    self.work = []
    self.work_lock = Lock()
  def do(work_id=None):
    self.work_lock.acquire()
    self.work.append(work_id)
    self.work_lock.release()

def try_except_pass(f):
  def run(*args, **kwargs):
    try:
      f(args, kwargs)
    except:
      pass
  return run

# <op-test>
work = Work()
@cache()
def test():
  work.do()
test()
test() # Uses cache - does not do work
assert len(work.work) == 1

# <op-test>
work = Work()
@cache()
def test():
  work.do()
  raise Exception()
try_except_pass(test)()
try_except_pass(test)() # Does not cache exceptions - does do work
assert len(work.work) == 2

# <op-test>
work = Work()
@cache(max_size=1)
def test(i):
  work.do(i)
test(1)
test(1) # Uses cache - does not do work
test(2)
test(1) # The result for this key fell out of the cache - does do work
assert work.work == [1, 2, 1]

# <op-test>
work = Work()
@cache(ttl_s=0.05)
def test():
  work.do(i)
test()
test() # Uses cache - does not do work
sleep(0.1)
test() # Value in cache has expired - does do work
assert len(work.work) == 2

# <op-test>
# This example demonstrates a difference from the "functools.lru_cache" related to multi-threading - even though no
# value has yet been cached by the time the second thread hits the function, it is aware that another thread is already
# working on the result, so it will block and return the result form that other thread.
work = Work()
@cache()
def test():
  work.do()
  sleep(0.1) # To make sure both threads would have time to get in here
Thread(target=test).start()
sleep(0.05)
test() # Waits on the result from the first thread (even though a result is not cached yet) - does not do work
assert len(work.work) == 1

# <op-test>
# The same logic applies to exceptions. Note that this is somewhat akin to caching an exception, but only for threads
# whose calls happened to overlap. Successive calls will not find a cached exception.
work = Work()
@cache()
def test():
  work.do()
  sleep(1)
  raise Exception()
Thread(target=try_except_pass(test)).start()
sleep(0.1)
Thread(target=try_except_pass(test)).start() # Waits on the result from the first thread - does not do work
sleep(1) # Both threads will finish during this statement
test_exception() # The first result came back as an exception, so it was not cached - does do work
assert len(work.work) == 2

# <op-test>
# The special multi-threading redundancy improvement does _not_ apply if we are dealing with the same thread. The
# following code will result in a "maximum recursion depth exceeded" error as opposed to hanging.
try:
  @cache()
  def test():
    test()
except RecursionError:
  pass

# </op-test>



"""


from collections import namedtuple, OrderedDict
from contextlib import suppress
from functools import wraps
from threading import Lock, Timer
from time import time


class CacheKey():

  # 'args' is a tuple and 'kwargs' is a dictionary (where all the keys are strings)
  def __init__(self, args, kwargs):
    self.args = args
    self.kwargs = kwargs
    self.hashcode = None

  # Assumes that 'other' is already the correct type
  def __eq__(self, other):
    return (self.args == other.args) and (self.kwargs == other.kwargs)

  def __hash__(self):
    if self.hashcode is None:
      hashable_kwargs = tuple(sorted(self.kwargs.items())) if self.kwargs else ()
      self.hashcode = hash((self.args, hashable_kwargs))
    return self.hashcode

  def __repr__(self):
    args_str = (", ").join([str(i) for i in self.args] + ["%s=%s" % (k, v) for (k, v) in self.kwargs.items()])
    return "%s(%s)" % (self.__class__.__name__, args_str)


class NoResultException(Exception):
  pass


class CacheEntry():

  def __init__(self, cache, args, kwargs):
    self.cache = cache
    self.args = args
    self.kwargs = kwargs

    self.updated_at = None
    self.load_lock = Lock()
    self.get_value()
    if self.cache.auto_populate_s is not None:
      self.start_auto_populate_timer()
 
  def __del__(self):
    self.did_evict()

  def did_evict(self):
    # We explicitly set 'is_evicted' instead of just cancelling the timer in case our auto-populate worker is in the
    # midst of doing work. In that case, the timer is not active, and the call to 'cancel' is meaningless
    if hasattr(self, 'timer'):
      self.is_evicted = True
      self.timer.cancel()

  def start_auto_populate_timer(self):
    # Note: this creates a new thread each time, which is presumably more inefficient...
    self.timer = Timer(self.cache.auto_populate_s, self.auto_populate)
    self.timer.start()

  def auto_populate(self):
    self.cache.get_cache_entry(self.args, self.kwargs)
    if not getattr(self, 'is_evicted', False):
      self.start_auto_populate_timer()

  def is_stale(self):
    if (self.cache.auto_populate_s is not None):
      return (time() - self.updated_at) > self.cache.auto_populate_s
    if (self.cache.ttl_s is None) or (not self.updated_at):
      return False
    return (time() - self.updated_at) > self.cache.ttl_s

  def check_result(self):
    if hasattr(self, 'value'):
      return self.value
    if hasattr(self, 'exception'):
      raise self.exception
    raise NoResultException()

  def get_value(self):
    with suppress(NoResultException):
      return self.check_result()

    self.load_lock.acquire()

    with suppress(NoResultException):
      return self.check_result()

    try:
      self.value = self.cache.fn(*self.args, **self.kwargs)
    except Exception as e:
      self.exception = e
    finally:
      self.load_lock.release()

    self.updated_at = time()
    return self.check_result()


CacheInfo = namedtuple("CacheInfo", [ "ttl_s", "max_size", "cur_size", "hits", "misses" ])


class Cache():

  def __init__(self, ttl_s, max_size, auto_populate_s, fn):
    self.ttl_s = ttl_s
    self.max_size = max_size
    self.auto_populate_s = auto_populate_s
    self.fn = fn
    self.hits = 0
    self.misses = 0
    self._contents = OrderedDict()
    self._update_lock = Lock()

  def _evict_key(self, cache_key):
    cache_entry = self._contents[cache_key]
    del self._contents[cache_key]
    cache_entry.did_evict()

  def evict(self, *args, **kwargs):
    key = CacheKey(args, kwargs)
    self._update_lock.acquire()
    try:
      self._evict_key(key)
      return True
    except:
      return False
    finally:
      self._update_lock.release()

  def _get_cache_entry_hit(self, args, kwargs):
    cache_key = CacheKey(args, kwargs)

    cache_entry = self._contents.get(cache_key)
    if cache_entry and (not cache_entry.is_stale()):
      return cache_entry, True

    self._update_lock.acquire()
    try:
      cache_entry = self._contents.get(cache_key)

      if not cache_entry:
        if (self.max_size is not None) and len(self._contents) >= self.max_size:
          _, evicted_cache_entry = self._contents.popitem(last=False)
          evicted_cache_entry.did_evict()
        new_cache_entry = CacheEntry(self, args, kwargs)
        self._contents[cache_key] = new_cache_entry
        return new_cache_entry, False

      if cache_entry.is_stale():
        self._evict_key(cache_key)
        new_cache_entry = CacheEntry(self, args, kwargs)
        self._contents[cache_key] = new_cache_entry
        return new_cache_entry, False

      return cache_entry, True
        
    finally:
      self._update_lock.release()

  def get_cache_entry(self, args, kwargs):
    cache_entry, is_hit = self._get_cache_entry_hit(args, kwargs)
    if is_hit:
      self.hits += 1
    else:
      self.misses += 1
    return cache_entry

  def info(self):
    return CacheInfo(self.ttl_s, self.max_size, len(self._contents), self.hits, self.misses)


def cache(ttl_s=None, max_size=128, auto_populate_s=None):
  def decorator(f):
    cache_manager = Cache(ttl_s, max_size, auto_populate_s, f)
    @wraps(f)
    def run(*args, **kwargs):
      return cache_manager.get_cache_entry(args, kwargs).get_value()
    setattr(run, 'cache', cache_manager)
    return run
  return decorator




# TESTS

# from collections import defaultdict
# from threading import Thread
# from time import sleep




# test_autopopulate_maxsize

# work_done = defaultdict(lambda: 0)
# @cache(max_size=2, auto_populate_s=1)
# def test(i):
#   work_done[i] += 1
# test(1) # work_done[1] += 1
# test(2) # work_done[2] += 1
# sleep(2.5) # work_done[1] += 2, work_done[2] += 2 (due to auto-population)
# test(1)
# test(2)
# assert work_done[1] == 3
# assert work_done[2] == 3
# test(3) # work_done[3] += 1 (evicts 1 from the cache)
# sleep(1) # work_done[1] += 1, work_done[2] += 1 (due to auto-population)
# assert work_done[1] == 3
# assert work_done[2] == 4
# assert work_done[3] == 2
# test.cache.evict(2)
# test.cache.evict(3)
# sleep(1.5) # work_done should not change due to auto-population
# assert work_done[1] == 3
# assert work_done[2] == 4
# assert work_done[3] == 2




# test_autopopulate_exception
# An exception should cause the auto-population to terminate (we don't cache exceptions)

# work_done = 0
# @cache(auto_populate_s=1)
# def test():
#   global work_done
#   work_done += 1
#   raise Exception()
# try:
#   test()
# except:
#   pass
# assert work_done == 1
# sleep(1.5) # If the auto-populate had been running, we would do more work here
# assert work_done == 1






# test_multithreading_uncached
# When there is no cached value and the function is called more than once while it is still working, the callers should
# share the results

# work_done = 0
# @cache()
# def test():
#   global work_done
#   work_done += 1
#   sleep(0.5)
# threads = [ Thread(target=test) for _ in range(3) ]
# [ i.start() for i in threads ]
# [ i.join() for i in threads ]
# assert work_done == 1



# test_multithreading_uncached_exception

# exceptions = 0
# @cache()
# def test():
#   sleep(0.5)
#   raise Exception()
# def expect_exception():
#   global exceptions
#   try:
#     test()
#   except:
#     exceptions += 1
# threads = [ Thread(target=expect_exception) for _ in range(3) ]
# [ i.start() for i in threads ]
# [ i.join() for i in threads ]
# assert exceptions == 3




# TODO: I have not tested all the nuances with auto-population and multi-threading





# test_multi_thread

# work_done = defaultdict(lambda: 0)
# @cache()
# def test(i):
#   work_done[i] += 1
# def t_1():
#   test(1)
# def t_2():
#   sleep(0.5)
#   test(1)


# try:
#   test()
# except:
#   pass
# assert work_done == 1
# sleep(1.5) # If the auto-populate had been running, we would do more work here
# assert work_done == 1







# import os
# os._exit(0)




# @cache(auto_populate_s=2)
# def test(a):
#   print('test(%s)' % a)
#   sleep(1)
#   return a

# print(test(1))
# sleep(6)
# print(test(1))