"""
Defines the base API for the openff-evaluator storage backend.
"""
import abc
import uuid
from collections import defaultdict
from threading import RLock
from typing import Dict

from openff.evaluator.attributes import Attribute
from openff.evaluator.storage.data import (
    BaseStoredData,
    ForceFieldData,
    HashableStoredData,
    ReplaceableData,
)


class StorageBackend(abc.ABC):
    """An abstract base representation of how the openff-evaluator will
    interact with and store simulation data.

    Notes
    -----
    When implementing this class, only private methods should be overridden
    as the public methods only mainly implement thread locks, while their
    private version perform their actual function.
    """

    class _ObjectKeyData(BaseStoredData):
        """An object which keeps track of the items in
        the storage system.
        """

        object_keys = Attribute(
            docstring="The unique keys of the objects stored in a `StorageBackend`.",
            type_hint=dict,
            default_value=dict(),
        )

        @classmethod
        def has_ancillary_data(cls):
            return False

        def to_storage_query(self):
            # This should never be called so doesn't need an
            # implementation.
            raise NotImplementedError()

    def __init__(self):
        """Constructs a new StorageBackend object.
        """
        self._stored_object_keys = None
        self._stored_object_keys_id = "object_keys"

        # Store a map between the unique id of a stored object,
        # and its hash value for easy comparision.
        self._object_hashes: Dict[int, str] = dict()

        # Create a thread lock to prevent concurrent
        # thread access.
        self._lock = RLock()

        self._load_stored_object_keys()

    def _load_stored_object_keys(self):
        """Load the unique key to each object stored in the
        storage system.
        """
        keys_object, _ = self._retrieve_object(self._stored_object_keys_id)

        if keys_object is None:
            keys_object = StorageBackend._ObjectKeyData()

        assert isinstance(keys_object, StorageBackend._ObjectKeyData)

        stored_object_keys = keys_object.object_keys
        self._stored_object_keys = defaultdict(list)

        all_object_keys = set()

        for data_type in stored_object_keys:

            for unique_key in stored_object_keys[data_type]:

                if not self._object_exists(unique_key):
                    # The stored entry key does not exist in the system,
                    # so skip the entry. This may happen when the local
                    # file do not exist on disk any more for example.
                    continue

                if unique_key in all_object_keys:

                    raise KeyError(
                        "Two objects with the same unique key have been found."
                    )

                stored_object, _ = self.retrieve_object(unique_key)

                # Make sure the data matches the expected type and is valid.
                assert stored_object.__class__.__name__ == data_type
                stored_object.validate()

                if isinstance(stored_object, HashableStoredData):
                    self._object_hashes[hash(stored_object)] = unique_key

                self._stored_object_keys[data_type].append(unique_key)
                all_object_keys.add(unique_key)

        # Store a fresh copy of the key dictionary so that only entries
        # that exist in the system actually referenced.
        self._save_stored_object_keys()

    def _save_stored_object_keys(self):
        """Save the unique key of each of the objects stored in the storage system.
        """
        keys_object = StorageBackend._ObjectKeyData()
        keys_object.object_keys = self._stored_object_keys

        self._store_object(keys_object, self._stored_object_keys_id)

    @abc.abstractmethod
    def _object_exists(self, storage_key):
        """Check whether an object with the specified key exists in the
        storage system.

        Parameters
        ----------
        storage_key: str
            A unique key that describes where the stored object can be found
            within the storage system.

        Returns
        -------
        True if the object is within the storage system.
        """
        raise NotImplementedError()

    def _is_key_unique(self, storage_key):
        """Checks whether a given key is already in the storage system.

        Parameters
        ----------
        storage_key: str
            The key to check for.

        Returns
        -------
        bool
            `True` if the key exists in the system, `False`
            otherwise.
        """
        # Make sure the key in unique.
        return not any(
            storage_key in self._stored_object_keys[data_type]
            for data_type in self._stored_object_keys
        )

    @abc.abstractmethod
    def _store_object(
        self, object_to_store, storage_key=None, ancillary_data_path=None
    ):
        """The internal implementation of the `store_object` method.
        It is safe to assume here that all object and key validation
        has already been performed, and that this method is called under
        a thread lock.

        Notes
        -----
        This method should overwrite any existing data with the same key.

        Parameters
        ----------
        object_to_store: BaseStoredData
            The object to store.
        storage_key: str, optional
            A unique key to associate with the stored object. If `None`,
            one will be randomly generated
        ancillary_data_path: str, optional
            The data path to the ancillary directory-like
            data to store alongside the object if the data
            type requires one.
        """
        raise NotImplementedError()

    def store_object(self, object_to_store, ancillary_data_path=None):
        """Store an object in the storage system, returning the key
        of the stored object. This may be different to `storage_key`
        depending on whether the same or a similar object was already
        present in the system.

        Parameters
        ----------
        object_to_store: BaseStoredData
            The object to store.
        ancillary_data_path: str, optional
            The data path to the ancillary directory-like
            data to store alongside the object if the data
            type requires one.

        Returns
        -------
        str
            The unique key assigned to the stored object.
        """

        # Make sure the object is valid.
        if object_to_store is None:
            raise ValueError("The object to store cannot be None.")

        object_to_store.validate()

        # Make sure the object is a supported type.
        if not isinstance(object_to_store, BaseStoredData):

            raise ValueError(
                "Only objects inheriting from `BaseStoredData` can "
                "be stored in the storage system."
            )

        # Make sure we have ancillary data if required.
        object_class = object_to_store.__class__

        if object_class.has_ancillary_data() and ancillary_data_path is None:
            raise ValueError("This object requires ancillary data.")

        # Check whether the exact same object already exists within
        # the storage system based on its hash.
        storage_key = self.has_object(object_to_store)

        if storage_key is not None:

            if not isinstance(object_to_store, ReplaceableData):
                # Handle the case where the existing data
                # should be returned, rather than storing
                # the passed object.
                return storage_key

            existing_object, _ = self.retrieve_object(storage_key, ReplaceableData)

            # noinspection PyTypeChecker
            object_to_store = object_to_store.most_information(
                existing_object, object_to_store
            )

            if object_to_store is None:

                raise ValueError(
                    "Something went wrong when trying to "
                    "determine whether the object trying to "
                    "be stored is redundant."
                )

            elif object_to_store == existing_object:
                # Don't try to re-store the existing object.
                return storage_key

        else:

            # Generate a unique id for this object.
            while storage_key is None or not self._is_key_unique(storage_key):
                storage_key = str(uuid.uuid4()).replace("-", "")

        # Hash this object if appropriate
        if isinstance(object_to_store, HashableStoredData):
            self._object_hashes[hash(object_to_store)] = storage_key

        # Save the object into the storage system with the given key.
        with self._lock:
            self._store_object(object_to_store, storage_key, ancillary_data_path)

        # Register the key in the storage system.
        if (
            not isinstance(object_to_store, StorageBackend._ObjectKeyData)
            and storage_key not in self._stored_object_keys[object_class.__name__]
        ):

            self._stored_object_keys[object_class.__name__].append(storage_key)
            self._save_stored_object_keys()

        return storage_key

    def store_force_field(self, force_field):
        """A convenience method for storing `ForceFieldSource` objects.

        Parameters
        ----------
        force_field: ForceFieldSource
            The force field to store.

        Returns
        -------
        str
            The unique id of the stored force field.
        """
        force_field_data = ForceFieldData()
        force_field_data.force_field_source = force_field

        return self.store_object(force_field_data)

    @abc.abstractmethod
    def _retrieve_object(self, storage_key, expected_type=None):
        """The internal implementation of the `retrieve_object` method.
        It is safe to assume that this method is called under a thread lock.

        Parameters
        ----------
        storage_key: str
            A unique key that describes where the stored object can be found
            within the storage system.
        expected_type: type of BaseStoredData, optional
            The expected data type. An exception is raised if
            the retrieved data doesn't match the type.

        Returns
        -------
        BaseStoredData, optional
            The stored object if the object key is found, otherwise None.
        str, optional
            The path to the ancillary data if present.
        """
        raise NotImplementedError()

    def retrieve_object(self, storage_key, expected_type=None):
        """Retrieves a stored object for the estimators storage system.

        Parameters
        ----------
        storage_key: str
            A unique key that describes where the stored object can be found
            within the storage system.
        expected_type: type of BaseStoredData, optional
            The expected data type. An exception is raised if
            the retrieved data doesn't match the type.

        Returns
        -------
        BaseStoredData, optional
            The stored object if the object key is found, otherwise None.
        str, optional
            The path to the ancillary data if present.
        """
        with self._lock:
            return self._retrieve_object(storage_key, expected_type)

    def retrieve_force_field(self, storage_key):
        """A convenience method for retrieving `ForceFieldSource` objects.

        Parameters
        ----------
        storage_key: str
            The key of the force field to retrieve.

        Returns
        -------
        ForceFieldSource
            The retrieved force field source.
        """
        force_field_data, _ = self.retrieve_object(storage_key, ForceFieldData)

        if force_field_data is None:

            raise KeyError(
                f"The force field with id {storage_key} does not exist "
                f"in the storage system."
            )

        return force_field_data.force_field_source

    def _has_object(self, storage_object):
        """The internal implementation of the `has_object` method.
        It is safe to assume that this method is called under a
        thread lock.

        Parameters
        ----------
        storage_object: BaseStoredData
            The object to check for.

        Returns
        -------
        str, optional
            The unique key of the object if it is in the system, `None` otherwise.
        """

        if isinstance(storage_object, HashableStoredData):

            hash_key = hash(storage_object)
            return self._object_hashes.get(hash_key, None)

        data_query = storage_object.to_storage_query()
        query_results = self.query(data_query)

        if len(query_results) == 0:
            return None

        if len(query_results) > 1 or len(query_results[0]) > 1:

            raise ValueError(
                "The backend contains multiple copies of the "
                "same piece of data. This should not be possible."
            )

        storage_key, _, _ = next(iter(query_results.values()))[0]
        return storage_key

    def has_object(self, storage_object):
        """Checks whether a given hashable object exists in the
        storage system.

        Parameters
        ----------
        storage_object: BaseStoredData
            The object to check for.

        Returns
        -------
        str, optional
            The unique key of the object if it is in the system, `None` otherwise.
        """
        with self._lock:
            return self._has_object(storage_object)

    def has_force_field(self, force_field):
        """A convenience method for checking whether the specified
        `ForceFieldSource` object is stored in the backend.

        Parameters
        ----------
        force_field: ForceFieldSource
            The force field to look for.

        Returns
        -------
        str, optional
            The unique key of the object if it is in the system, `None` otherwise.
        """
        force_field_data = ForceFieldData()
        force_field_data.force_field_source = force_field

        return self.has_object(force_field_data)

    def _query(self, data_query):
        """The internal implementation of the `query` method.
        It is safe to assume that this method is called under a
        thread lock.

        Parameters
        ----------
        data_query: BaseDataQuery
            The query to perform.

        Returns
        -------
        dict of tuple and list of tuple of str, BaseStoredData and str
            The data that matches the query partitioned by the
            matched values. The list values take the form
            (storage_key, data_object, data_directory_path).
        """

        data_class = data_query.data_class()
        results = defaultdict(list)

        if len(self._stored_object_keys.get(data_class.__name__, [])) == 0:
            # Exit early of there are no objects of the correct type.
            return results

        for unique_key in self._stored_object_keys[data_class.__name__]:

            if not self._object_exists(unique_key):
                # Make sure the object is still in the system.
                continue

            stored_object, stored_directory = self.retrieve_object(
                unique_key, data_class
            )

            matches = data_query.apply(stored_object)

            if matches is None:
                continue

            results[matches].append((unique_key, stored_object, stored_directory))

        return results

    def query(self, data_query):
        """Query the storage backend for data matching the
        query criteria.

        Parameters
        ----------
        data_query: BaseDataQuery
            The query to perform.

        Returns
        -------
        dict of tuple and list of tuple of str, BaseStoredData and str
            The data that matches the query partitioned by the
            matched values. The list values take the form
            (storage_key, data_object, data_directory_path).
        """
        with self._lock:
            return self._query(data_query)
