import base64
import functools
import pathlib
import pickle
import shelve
import weakref
from typing import Callable

import cachetools

_SHELVE_CACHES = weakref.WeakValueDictionary()


class DelegatingCache(cachetools.Cache):
    def __init__(self, get_cache: Callable[[], cachetools.Cache]):
        self.get_cache = get_cache

    @property
    def maxsize(self):
        return self.get_cache().maxsize

    @property
    def currsize(self):
        return self.get_cache().currsize

    def __getitem__(self, key):
        return self.get_cache()[key]

    def __setitem__(self, key, val):
        self.get_cache()[key] = val

    def __delitem__(self, key):
        del (self.get_cache())[key]


def delegating_cached(get_cache: Callable[[], cachetools.Cache]):
    cache = DelegatingCache(get_cache)
    return cachetools.cached(cache=cache, key=cachetools.keys.typedkey, info=True)


class ShelveCache(cachetools.Cache):
    """A simple file-based cachetools cache of unlimited size. Uses `shelve` to store the cache."""

    def __init__(self, fname, prefix=""):
        self.sh = shelve.open(fname)
        self.prefix = prefix

    @property
    def maxsize(self):
        return 1 << 30

    @property
    def currsize(self):
        return 0

    def db_key(self, key: any) -> str:
        return self.prefix + repr(key)

    def __getitem__(self, key):
        return self.sh[self.db_key(key)]

    def __setitem__(self, key, val):
        self.sh[self.db_key(key)] = val
        self.sh.sync()

    def __delitem__(self, key):
        del self.sh[self.db_key(key)]
        self.sh.sync()


def shelve_cached(fname, prefix=None):
    """
    Decorator for cached functions.

    Uses global table of caches to avoid opening the same file twice.
    The files are eventually closed when the caches are garbage collected (e.g. upon decorated function redefinition).
    """
    fname = pathlib.Path(fname).absolute()
    cache = _SHELVE_CACHES.get(fname)
    if cache is None:
        cache = ShelveCache(str(fname))
        _SHELVE_CACHES[fname] = cache
    keyfunc = cachetools.keys.typedkey
    if prefix is not None:
        keyfunc = functools.partial(keyfunc, prefix)
    return cachetools.cached(cache=cache, key=keyfunc, info=True)


class DynamoDBCache(cachetools.Cache):
    """A simple DynamoDB-based cachetools cache of unlimited size. Uses `pickle` to store the objects."""

    def __init__(self, client, table_name, prefix=""):
        self.client = client
        self.table_name = table_name
        self.prefix = prefix

    @property
    def maxsize(self):
        return 1 << 30

    @property
    def currsize(self):
        return 0

    def db_key(self, key: any) -> str:
        return self.prefix + repr(key)

    def __getitem__(self, key):
        r = self.client.get_item(TableName=self.table_name, Key={"key": {"S": self.db_key(key)}})
        if "Item" not in r:
            raise KeyError("Key not present in DynamoDB cache table")
        d = base64.b64decode(r["Item"]["value"]["B"])
        return pickle.loads(d)

    def __setitem__(self, key, val):
        d = base64.b64encode(pickle.dumps(val))
        self.client.put_item(TableName=self.table_name, Item={"key": {"S": self.db_key(key)}, "value": {"B": d}})

    def __delitem__(self, key):
        r = self.client.delete_item(TableName=self.table_name, Key={"key": {"S": self.db_key(key)}})


def dynamodb_cached(client, table_name, prefix=None):
    """
    Decorator for DynamoDBCache-cached functions.
    """
    cache = DynamoDBCache(client, table_name)
    keyfunc = cachetools.keys.typedkey
    if prefix is not None:
        keyfunc = functools.partial(keyfunc, prefix)
    return cachetools.cached(cache=cache, key=keyfunc, info=True)
