"""
Defines the base API for defining new property estimator estimation layers.
"""
import json
import logging
import traceback
from os import path

from propertyestimator.storage.dataclasses import StoredDataCollection
from propertyestimator.utils.exceptions import PropertyEstimatorException
from propertyestimator.utils.serialization import TypedJSONDecoder

available_layers = {}


def register_calculation_layer():
    """A decorator which registers a class as being a calculation layer
    which may be used in property calculations.

    See Also
    --------
    TODO: add documentation for plugin support
    """

    def decorator(cls):

        if cls.__name__ in available_layers:
            raise ValueError('The {} layer is already registered.'.format(cls.__name__))

        available_layers[cls.__name__] = cls
        return cls

    return decorator


def return_args(*args, **kwargs):
    return args


class CalculationLayerResult:
    """The output returned from attempting to calculate a property on
    a PropertyCalculationLayer.

    Attributes
    ----------
    property_id: str
        The unique id of the original physical property that this
        calculation layer attempts to estimate.
    calculated_property: PhysicalProperty, optional
        The property which was estimated by this layer. The will
        be `None` if the layer could not estimate the property.
    exception: PropertyEstimatorException, optional
        The exception which was raised when estimating the property
        of interest, if any.
    data_to_store: list of tuple of str and str
        A list of pairs of a path to a JSON serialized `BaseStoredData`
        object, and the path to the corresponding data directory.
    """

    def __init__(self):
        """Constructs a new CalculationLayerResult object.
        """
        self.property_id = None

        self.calculated_property = None
        self.exception = None

        self.data_to_store = []

    def __getstate__(self):

        return {
            'property_id': self.property_id,
            'calculated_property': self.calculated_property,
            'exception': self.exception,
            'data_to_store': self.data_to_store
        }

    def __setstate__(self, state):

        self.property_id = state['property_id']

        self.calculated_property = state['calculated_property']
        self.exception = state['exception']

        self.data_to_store = state['data_to_store']


class PropertyCalculationLayer:
    """An abstract representation of a calculation layer whose goal is
    to estimate a set of physical properties using a single approach,
    such as a layer which employs direct simulations to estimate properties,
    or one which reweights cached simulation data to the same end.

    Notes
    -----
    Calculation layers must inherit from this class, and must override the
    `schedule_calculation` method.

    See Also
    --------
    TODO: Link to a general page outlining what calculation layers are
          and how they are used.
    """

    @staticmethod
    def _await_results(calculation_backend, storage_backend, layer_directory, server_request,
                       callback, submitted_futures, synchronous=False):
        """A helper method to handle passing the results of this layer back to
        the main thread.

        Parameters
        ----------
        calculation_backend: PropertyEstimatorBackend
            The backend to the submit the calculations to.
        storage_backend: PropertyEstimatorStorage
            The backend used to store / retrieve data from previous calculations.
        layer_directory: str
            The local directory in which to store all local, temporary calculation data from this layer.
        server_request: PropertyEstimatorServer.ServerEstimationRequest
            The request object which spawned the awaited results.
        callback: function
            The function to call when the backend returns the results (or an error).
        submitted_futures: list(dask.distributed.Future)
            A list of the futures returned by the backed when submitting the calculation.
        synchronous: bool
            If true, this function will block until the calculation has completed.
        """

        callback_future = calculation_backend.submit_task(return_args,
                                                          *submitted_futures,
                                                          key=f'return_{server_request.id}')

        def callback_wrapper(results_future):
            PropertyCalculationLayer._process_results(results_future, server_request, storage_backend, callback)

        if synchronous:
            callback_wrapper(callback_future)
        else:
            callback_future.add_done_callback(callback_wrapper)

    @staticmethod
    def _store_cached_output(server_request, returned_output, storage_backend):
        """Stores any cached pieces of simulation data using a storage backend.

        Parameters
        ----------
        server_request: PropertyEstimatorServer.ServerEstimationRequest
            The request which generated the cached data.
        returned_output: CalculationLayerResult
            The layer result which contains the cached data.
        storage_backend: PropertyEstimatorStorage
            The backend to use to store the cached data.
        """

        for data_tuple in returned_output.data_to_store:

            data_object_path, data_directory_path = data_tuple

            # Make sure the data directory / file to store actually exists
            if not path.isdir(data_directory_path) or not path.isfile(data_object_path):
                logging.info(f'Invalid data directory ({data_directory_path}) / '
                             f'file ({data_object_path})')
                continue

            # Attach any extra metadata which is missing.
            with open(data_object_path, 'r') as file:

                data_object = json.load(file, cls=TypedJSONDecoder)

                if data_object.force_field_id is None:
                    data_object.force_field_id = server_request.force_field_id

                if isinstance(data_object, StoredDataCollection):

                    for inner_data_object in data_object.data.values():

                        if inner_data_object.force_field_id is None:
                            inner_data_object.force_field_id = server_request.force_field_id

            storage_backend.store_simulation_data(data_object, data_directory_path)

    @staticmethod
    def _process_results(results_future, server_request, storage_backend, callback):
        """Processes the results of a calculation layer, updates the server request,
        then passes it back to the callback ready for propagation to the next layer
        in the stack.

        Parameters
        ----------
        results_future: distributed.Future
            The future object which will hold the results.
        server_request: PropertyEstimatorServer.ServerEstimationRequest
            The request object which spawned the awaited results.
        storage_backend: PropertyEstimatorStorage
            The backend used to store / retrieve data from previous calculations.
        callback: function
            The function to call when the backend returns the results (or an error).
        """

        # Wrap everything in a try catch to make sure the whole calculation backend /
        # server doesn't go down when an unexpected exception occurs.
        try:

            results = list(results_future.result())
            results_future.release()

            for returned_output in results:

                if returned_output is None:
                    # Indicates the layer could not calculate this
                    # particular property.
                    continue

                if not isinstance(returned_output, CalculationLayerResult):

                    # Make sure we are actually dealing with the object we expect.
                    raise ValueError('The output of the calculation was not '
                                     'a CalculationLayerResult as expected.')

                if returned_output.exception is not None:
                    # If an exception was raised, make sure to add it to the list.
                    server_request.exceptions.append(returned_output.exception)

                    logging.info(f'An exception was raised: '
                                 f'{returned_output.exception.directory} - '
                                 f'{returned_output.exception.message}')

                else:

                    # Make sure to store any important calculation data if no exceptions
                    # were thrown.
                    if (returned_output.data_to_store is not None and
                        returned_output.calculated_property is not None):

                        PropertyCalculationLayer._store_cached_output(server_request, returned_output, storage_backend)

                matches = [x for x in server_request.queued_properties if x.id == returned_output.property_id]

                if len(matches) > 1:
                    raise ValueError(f'A property id ({returned_output.property_id}) conflict occurred.')

                elif len(matches) == 0:

                    logging.info('A calculation layer returned results for a property not in the '
                                 'queue. This sometimes and expectedly occurs when using queue based '
                                 'calculation backends, but should be investigated.')

                    continue

                if returned_output.calculated_property is None:

                    if returned_output.exception is None:

                        logging.info('A calculation layer did not return an estimated property nor did it '
                                     'raise an Exception. This sometimes and expectedly occurs when using '
                                     'queue based calculation backends, but should be investigated.')

                    continue

                if returned_output.exception is not None:
                    continue

                for match in matches:
                    server_request.queued_properties.remove(match)

                substance_id = returned_output.calculated_property.substance.identifier

                if substance_id not in server_request.estimated_properties:
                    server_request.estimated_properties[substance_id] = []

                server_request.estimated_properties[substance_id].append(returned_output.calculated_property)

        except Exception as e:

            logging.info(f'Error processing layer results for request {server_request.id}')

            formatted_exception = traceback.format_exception(None, e, e.__traceback__)

            exception = PropertyEstimatorException(message='An unhandled internal exception '
                                                           'occurred: {}'.format(formatted_exception))

            server_request.exceptions.append(exception)

        callback(server_request)

    @staticmethod
    def schedule_calculation(calculation_backend, storage_backend, layer_directory,
                             data_model, callback, synchronous=False):
        """Submit the proposed calculation to the backend of choice.

        Parameters
        ----------
        calculation_backend: PropertyEstimatorBackend
            The backend to the submit the calculations to.
        storage_backend: PropertyEstimatorStorage
            The backend used to store / retrieve data from previous calculations.
        layer_directory: str
            The local directory in which to store all local, temporary calculation data from this layer.
        data_model: PropertyEstimatorServer.ServerEstimationRequest
            The data model encoding the proposed calculation.
        callback: function
            The function to call when the backend returns the results (or an error).
        synchronous: bool
            If true, this function will block until the calculation has completed.
            This is mainly intended for debugging purposes.
        """
        raise NotImplementedError()
