from .commoncache import HashCollisionError

from collections.abc import MutableMapping, Iterator
from os import makedirs, remove, listdir
from os.path import exists, join


class Cache:
    """Dictionary-like object for saving function outputs to disk

    This cache, which can be used by the `persist` decorator in `persist.py`,
    stores computed values on disk in a specified directory so that they can be
    restored later using a key.  Like a dictionary, a key-value pair can be
    added using `cache[key] = val`, looked up using `cache[key]`, and removed
    using `del cache[key]`.  The number of values stored can be found using
    `len(cache)`.

    A disk cache might not store its keys, and therefore we cannot iterate
    through its keys as we can with a dictionary.  However, see
    `CacheWithKeys`.

    Parameters
    ----------
    func : persist_wrapper
        Memoised function whose results this is caching.  Options which are not
        specific to local disk storage, such as the key, hash, and pickle
        functions, are taken from this.
    dir : str
        Directory into which to save results.  The same directory can be used
        for several different functions, since a subdirectory will be created
        for each function based on its `funcname`.

    """

    def __init__(self, func, dir):
        self._func = func
        self._dir = join(dir, self._func._funcname)
        if not exists(self._dir):
            makedirs(self._dir)

    def __getitem__(self, key):
        fname = self._key_to_fname(key, OUT)
        if self._func._unhash:
            storedkey = self._fname_to_key(fname)
            if storedkey != key:
                raise HashCollisionError(storedkey, key)
        if exists(fname):
            if self._func._storekey:
                keyfname = self._key_to_fname(key, KEY)
                assert(exists(fname))
                keyfile = open(keyfname, 'r')
                keystring = keyfile.read()
                keyfile.close()
                storedkey = self._func._unpickle(keystring)
                if storedkey != key:
                    raise HashCollisionError(storedkey, key)
            file = open(fname, 'r')
            val = self._func._unpickle(file.read())
            file.close()
        else:
            raise KeyError(key)
        return val

    def __setitem__(self, key, val):
        if self._func._storekey:
            keyfname = self._key_to_fname(key, KEY)
            keyfile = open(keyfname, 'w')
            keyfile.write(self._func._pickle(key))
            keyfile.close()
        fname = self._key_to_fname(key, OUT)
        file = open(fname, 'w')
        file.write(self._func._pickle(val))
        file.close()

    def __delitem__(self, key):
        fname = self._key_to_fname(key, OUT)
        if exists(fname):
            remove(fname)
        else:
            raise KeyError(key)

    def __len__(self):
        # Number of files ending with '.out'
        return sum(fname.endswith(OUT) for fname in listdir(self._dir))

    def clear(self):
        """Delete all the results stored in this cache."""
        for f in listdir(self._dir):
            path = join(self._dir, f)
            # TODO: safety checks?
            remove(path)

    def _key_to_fname(self, key, ext):
        h = self._func._hash(key)
        return join(self._dir, h + ext)

    def _fname_to_key(self, fname):
        if fname.startswith(self._dir):
            fname = fname[len(self._dir + '/'):]  # remove directory
        h = fname[:fname.rfind('.')]  # remove extension
        return self._func._unhash(h)


class CacheWithKeys(Cache, MutableMapping):
    """Mutable mapping for saving function outputs to disk

    This subclass of `Cache` can be used in place of `Cache` whenever
    `storekey` is True or `unhash` is set, to implement the `MutableMapping`
    abstract base class.  This allows the cache to be used exactly like a
    dictionary, including the ability to iterate through all keys in the cache.

    """

    def __iter__(self):
        return self.KeysIter(self)

    class KeysIter(Iterator):
        """Iterator class for the keys of a `CacheWithKeys` object"""
        def __init__(self, cache):
            self._cache = cache
            self._pos = 0
            self._files = [fname
                           for fname in listdir(self._cache._dir)
                           if fname.endswith(OUT)]

        def __next__(self):
            if self._pos >= len(self._files):
                raise StopIteration
            fname = self._files[self._pos]
            if self._cache._func._unhash:
                # Unhash from filename
                key = self._cache._fname_to_key(fname)
            else:
                assert(self._cache._func._storekey)
                # Read key from file
                path = join(self._cache._dir, fname)
                path = path[:-len(OUT)] + KEY
                file = open(path, 'r')
                string = file.read()
                key = self._cache._func._unpickle(string)
                file.close()
            self._pos += 1
            return key


# Filename extensions
OUT = '.out'
KEY = '.key'
META = '.meta'
