# Source code for propertyestimator.storage.localfile

    """
    A local file based storage backend.
    """
    import json
    import shutil
    from os import path, makedirs

    from propertyestimator.storage import StoredSimulationData
    from propertyestimator.storage.dataclasses import BaseStoredData
    from propertyestimator.substances import Substance
    from propertyestimator.utils.serialization import TypedJSONEncoder, TypedJSONDecoder
    from .storage import PropertyEstimatorStorage

    class LocalFileStorage(PropertyEstimatorStorage):
        """A storage backend which stores files in directories on the local
        disk.
        """

        @property
        def root_directory(self):
            """str: Returns the directory in which all stored objects are located."""
            return self.root_directory

        def __init__(self, root_directory='stored_data'):

            self._root_directory = root_directory

            if not path.isdir(root_directory):
                makedirs(root_directory)

            super().__init__()

        def _store_object(self, storage_key, object_to_store):

            file_path = path.join(self._root_directory, storage_key)

            # If the object to store is a simple string we write that
            # directly, otherwise we try and JSONify the object.
            if not isinstance(object_to_store, str):
                object_to_store = json.dumps(object_to_store, cls=TypedJSONEncoder)

            with open(file_path, 'w') as file:
                file.write(object_to_store)

            super(LocalFileStorage, self)._store_object(storage_key, object_to_store)

        def _retrieve_object(self, storage_key):

            if not self._has_object(storage_key):
                return None

            file_path = path.join(self._root_directory, storage_key)

            with open(file_path, 'r') as file:
                loaded_object_string = file.read()

            try:
                loaded_object = json.loads(loaded_object_string, cls=TypedJSONDecoder)
            except json.JSONDecodeError:
                loaded_object = loaded_object_string

            return loaded_object

        def _has_object(self, storage_key):

            file_path = path.join(self._root_directory, storage_key)

            if not path.isfile(file_path):
                return False

            return True

        def store_simulation_data(self, data_object, data_directory):

            unique_id = super(LocalFileStorage, self).store_simulation_data(data_object,
                                                                            data_directory)

            storage_directory_path = path.join(self._root_directory, f'{unique_id}_data')

            if path.isdir(storage_directory_path):
                shutil.rmtree(storage_directory_path, ignore_errors=True)

            shutil.move(data_directory, storage_directory_path)
            return unique_id

        def retrieve_simulation_data_by_id(self, unique_id):
            """Attempts to retrieve a storage piece of simulation data
            from it's unique id.

            Parameters
            ----------
            unique_id: str
                The unique id assigned to the data.

            Returns
            -------
            BaseStoredData
                The stored data object.
            str
                The path to the data's corresponding directory.
            """
            stored_object = self._retrieve_object(unique_id)

            # Make sure the stored object is a valid object.
            if not isinstance(stored_object, BaseStoredData):
                return None, None

            data_directory = path.join(self._root_directory, f'{unique_id}_data')
            return stored_object, data_directory

        def retrieve_simulation_data(self, substance, include_component_data=True,
                                     data_class=StoredSimulationData):

            substance_ids = {substance.identifier}

            # Find the substance identifiers of the substance components if
            # we should include component data.
            if isinstance(substance, Substance) and include_component_data is True:

                for component in substance.components:

                    component_substance = Substance()
                    component_substance.add_component(component, Substance.MoleFraction())

                    substance_ids.add(component_substance.identifier)

            return_data = {}

            for substance_id in substance_ids:

                if substance_id not in self._simulation_data_by_substance:
                    continue

                return_data[substance_id] = []

                for data_key in self._simulation_data_by_substance[substance_id]:

                    data_object, data_directory = self.retrieve_simulation_data_by_id(data_key)

                    if data_object is None:
                        continue

                    if not isinstance(data_object, data_class):
                        continue

                    return_data[substance_id].append((data_object, data_directory))

            return return_data