from enum import Enum
import threading
import time
from server.data_common.rwlock import RWLock
from server.common.errors import DatasetAccessError
from server.common.data_locator import DataLocator
from contextlib import contextmanager
from http import HTTPStatus


class MatrixDataCacheItem(object):
    """This class provides access and caching for a dataset.  The first time a dataset is accessed, it is
    opened and cached.  Later accesses use the cached version.   It may also be deleted by the
    MatrixDataCacheManager to make room for another dataset.  While a dataset is actively being used
    (during the lifetime of a api request), a reader lock is locked.  During that time, the dataset cannot
    be removed."""

    def __init__(self, loader):
        self.loader = loader
        self.data_adaptor = None
        self.data_lock = RWLock()

    def acquire_existing(self):
        """If the data_adaptor exists, take a read lock and return it, else return None"""
        self.data_lock.r_acquire()
        if self.data_adaptor:
            return self.data_adaptor

        self.data_lock.r_release()
        return None

    def acquire_and_open(self, app_config, dataset_config=None):
        """returns the data_adaptor if cached.  opens the data_adaptor if not.
        In either case, the a reader lock is taken.  Must call release when
        the data_adaptor is no longer needed"""
        self.data_lock.r_acquire()
        if self.data_adaptor:
            return self.data_adaptor
        self.data_lock.r_release()

        self.data_lock.w_acquire()
        # the data may have been loaded while waiting on the lock
        if not self.data_adaptor:
            try:
                self.loader.pre_load_validation()
                self.data_adaptor = self.loader.open(app_config, dataset_config)
            except Exception as e:
                # necessary to hold the reader lock after an exception, since
                # the release will occur when the context exits.
                self.data_lock.w_demote()
                raise DatasetAccessError(str(e))

        # demote the write lock to a read lock.
        self.data_lock.w_demote()
        return self.data_adaptor

    def release(self):
        """Release the reader lock"""
        self.data_lock.r_release()

    def delete(self):
        """Clear resources used by this dataset"""
        with self.data_lock.w_locked():
            if self.data_adaptor:
                self.data_adaptor.cleanup()
                self.data_adaptor = None

    def attempt_delete(self):
        """Delete, but only if the write lock can be immediately locked.  Return True if the delete happened"""
        if self.data_lock.w_acquire_non_blocking():
            if self.data_adaptor:
                try:
                    self.data_adaptor.cleanup()
                    self.data_adaptor = None
                except Exception:
                    # catch all exceptions to ensure the lock is released
                    pass

                self.data_lock.w_release()
                return True
        else:
            return False


class MatrixDataCacheInfo(object):
    def __init__(self, cache_item, timestamp):
        # The MatrixDataCacheItem in the cache
        self.cache_item = cache_item
        # The last time the cache_item was accessed
        self.last_access = timestamp
        # The number of times the cache_item was accessed (used for testing)
        self.num_access = 1


class MatrixDataCacheManager(object):
    """A class to manage the cached datasets.   This is intended to be used as a context manager
    for handling api requests.  When the context is created, the data_adator is either loaded or
    retrieved from a cache.  In either case, the reader lock is taken during this time, and release
    when the context ends.  This class currently implements a simple least recently used cache,
    which can delete a dataset from the cache to make room for a new one.

    This is the intended usage pattern:

           m = MatrixDataCacheManager(max_cached=..., timelimmit_s = ...)
           with m.data_adaptor(location, app_config) as data_adaptor:
               # use the data_adaptor for some operation
    """

    # FIXME:   If the number of active datasets exceeds the max_cached, then each request could
    # lead to a dataset being deleted and a new only being opened: the cache will get thrashed.
    # In this case, we may need to send back a 503 (Server Unavailable), or some other error message.

    # NOTE:  If the actual dataset is changed.  E.g. a new set of datafiles replaces an existing set,
    # then the cache will not react to this, however once the cache time limit is reached, the dataset
    # will automatically be refreshed.

    def __init__(self, max_cached, timelimit_s=None):
        # key is tuple(url_dataroot, location), value is a MatrixDataCacheInfo
        self.datasets = {}

        # lock to protect the datasets
        self.lock = threading.Lock()

        #  The number of datasets to cache.  When max_cached is reached, the least recently used
        #  cache is replaced with the newly requested one.
        #  TODO:  This is very simple.  This can be improved by taking into account how much space is actually
        #         taken by each dataset, instead of arbitrarily picking a max datasets to cache.
        self.max_cached = max_cached

        # items are automatically removed from the cache once this time limit is reached
        self.timelimit_s = timelimit_s

    @contextmanager
    def data_adaptor(self, url_dataroot, location, app_config):
        # create a loader for to this location if it does not already exist

        delete_adaptor = None
        data_adaptor = None
        cache_item = None

        key = (url_dataroot, location)
        with self.lock:
            self.evict_old_datasets()
            info = self.datasets.get(key)
            if info is not None:
                info.last_access = time.time()
                info.num_access += 1
                self.datasets[key] = info
                data_adaptor = info.cache_item.acquire_existing()
                cache_item = info.cache_item

            if data_adaptor is None:
                while True:
                    if len(self.datasets) < self.max_cached:
                        break

                    items = list(self.datasets.items())
                    items = sorted(items, key=lambda x: x[1].last_access)
                    # close the least recently used loader
                    oldest = items[0]
                    oldest_cache = oldest[1].cache_item
                    oldest_key = oldest[0]
                    del self.datasets[oldest_key]
                    delete_adaptor = oldest_cache

                loader = MatrixDataLoader(location, app_config=app_config)
                cache_item = MatrixDataCacheItem(loader)
                item = MatrixDataCacheInfo(cache_item, time.time())
                self.datasets[key] = item

        try:
            assert cache_item
            if delete_adaptor:
                delete_adaptor.delete()
            if data_adaptor is None:
                dataset_config = app_config.get_dataset_config(url_dataroot)
                data_adaptor = cache_item.acquire_and_open(app_config, dataset_config)
            yield data_adaptor
        except DatasetAccessError:
            cache_item.release()
            with self.lock:
                del self.datasets[key]
                cache_item.delete()
            cache_item = None
            raise

        finally:
            if cache_item:
                cache_item.release()

    def evict_old_datasets(self):
        # must be called with the lock held
        if self.timelimit_s is None:
            return

        now = time.time()
        to_del = []
        for key, info in self.datasets.items():
            if (now - info.last_access) > self.timelimit_s:
                # remove the data_cache when if it has been in the cache too long
                to_del.append((key, info))

        for key, info in to_del:
            # try and get the write_lock for the dataset.
            # if this returns false, it means the dataset is being used, and should
            # not be removed.
            if info.cache_item.attempt_delete():
                del self.datasets[key]


class MatrixDataType(Enum):
    H5AD = "h5ad"
    CXG = "cxg"
    UNKNOWN = "unknown"


class MatrixDataLoader(object):
    def __init__(self, location, matrix_data_type=None, app_config=None):
        """ location can be a string or DataLocator """
        region_name = None if app_config is None else app_config.server_config.data_locator__s3__region_name
        self.location = DataLocator(location, region_name=region_name)
        if not self.location.exists():
            raise DatasetAccessError("Dataset does not exist.", HTTPStatus.NOT_FOUND)

        # matrix_data_type is an enum value of type MatrixDataType
        self.matrix_data_type = matrix_data_type
        # matrix_type is a DataAdaptor type, which corresonds to the matrix_data_type
        self.matrix_type = None

        if matrix_data_type is None:
            self.matrix_data_type = self.__matrix_data_type()

        if not self.__matrix_data_type_allowed(app_config):
            raise DatasetAccessError("Dataset does not have an allowed type.")

        if self.matrix_data_type == MatrixDataType.H5AD:
            from server.data_anndata.anndata_adaptor import AnndataAdaptor

            self.matrix_type = AnndataAdaptor
        elif self.matrix_data_type == MatrixDataType.CXG:
            from server.data_cxg.cxg_adaptor import CxgAdaptor

            self.matrix_type = CxgAdaptor

    def __matrix_data_type(self):
        if self.location.path.endswith(".h5ad"):
            return MatrixDataType.H5AD
        elif ".cxg" in self.location.path:
            return MatrixDataType.CXG
        else:
            return MatrixDataType.UNKNOWN

    def __matrix_data_type_allowed(self, app_config):
        if self.matrix_data_type == MatrixDataType.UNKNOWN:
            return False

        if not app_config:
            return True
        if not app_config.is_multi_dataset():
            return True
        if len(app_config.server_config.multi_dataset__allowed_matrix_types) == 0:
            return True

        for val in app_config.server_config.multi_dataset__allowed_matrix_types:
            try:
                if self.matrix_data_type == MatrixDataType(val):
                    return True
            except ValueError:
                # Check case where multi_dataset_allowed_matrix_type does not have a
                # valid MatrixDataType value.  TODO:  Add a feature to check
                # the AppConfig for errors on startup
                return False

        return False

    def pre_load_validation(self):
        if self.matrix_data_type == MatrixDataType.UNKNOWN:
            raise DatasetAccessError("Dataset does not have a recognized type: .h5ad or .cxg")
        self.matrix_type.pre_load_validation(self.location)

    def file_size(self):
        return self.matrix_type.file_size(self.location)

    def open(self, app_config, dataset_config=None):
        # create and return a DataAdaptor object
        return self.matrix_type.open(self.location, app_config, dataset_config)
