"""
Property calculator 'server' side API.
"""
import copy
import json
import logging
import uuid
from os import path, makedirs

from tornado.ioloop import IOLoop
from tornado.iostream import StreamClosedError
from tornado.tcpserver import TCPServer

from propertyestimator.client import PropertyEstimatorSubmission, PropertyEstimatorResult, PropertyEstimatorOptions
from propertyestimator.layers import available_layers
from propertyestimator.utils.exceptions import PropertyEstimatorException
from propertyestimator.utils.serialization import TypedBaseModel
from propertyestimator.utils.tcp import PropertyEstimatorMessageTypes, pack_int, unpack_int


class PropertyEstimatorServer(TCPServer):
    """The object responsible for coordinating all properties estimations to to
    be ran using the property estimator, in addition to deciding at which fidelity
    a property will be calculated.

    It acts as a server, which receives submitted jobs from clients
    launched via the property estimator.

    Warnings
    --------
    This class is still heavily under development and is subject to rapid changes.

    Notes
    -----
    Methods to handle the TCP messages are based on the StackOverflow response from
    A. Jesse Jiryu Davis: https://stackoverflow.com/a/40257248

    Examples
    --------
    Setting up a general server instance using a dask LocalCluster backend:

    >>> # Create the backend which will be responsible for distributing the calculations
    >>> from propertyestimator.backends import DaskLocalCluster, ComputeResources
    >>> calculation_backend = DaskLocalCluster(1)
    >>>
    >>> # Calculate the backend which will be responsible for storing and retrieving
    >>> # the data from previous calculations
    >>> from propertyestimator.storage import LocalFileStorage
    >>> storage_backend = LocalFileStorage()
    >>>
    >>> # Create the server to which all estimation requests will be submitted
    >>> from propertyestimator.server import PropertyEstimatorServer
    >>> property_server = PropertyEstimatorServer(calculation_backend, storage_backend)
    >>>
    >>> # Instruct the server to listen for incoming requests
    >>> property_server.start_listening_loop()
    """

    class ServerEstimationRequest(TypedBaseModel):
        """Represents a request for the server to estimate a set of properties. Such requests
        are expected to only estimate properties for a single system (e.g. fixed components
        in a fixed ratio)
        """

        def __init__(self, estimation_id='', queued_properties=None, options=None,
                     force_field_id=None, parameter_gradient_keys=None):
            """Constructs a new ServerEstimationRequest object.

            Parameters
            ----------
            estimation_id: str
                A unique id assigned to this estimation request.
            queued_properties: list of PhysicalProperty, optional
                A list of physical properties waiting to be estimated.
            options: PropertyEstimatorOptions, optional
                The options used to estimate the properties.
            force_field_id: str
                The unique server side id of the force field parameters used to estimate the properties.
            parameter_gradient_keys: list of ParameterGradientKey
                A list of references to all of the parameters which all observables
                should be differentiated with respect to.
            """
            self.id = estimation_id

            self.queued_properties = queued_properties or []

            self.estimated_properties = {}
            self.unsuccessful_properties = {}

            self.exceptions = []

            self.options = options

            self.force_field_id = force_field_id

            self.parameter_gradient_keys = parameter_gradient_keys

        def __getstate__(self):
            return {
                'id': self.id,

                'queued_properties': self.queued_properties,

                'estimated_properties': self.estimated_properties,
                'unsuccessful_properties': self.unsuccessful_properties,

                'exceptions': self.exceptions,

                'options': self.options,

                'force_field_id': self.force_field_id,

                'parameter_gradient_keys': self.parameter_gradient_keys
            }

        def __setstate__(self, state):
            self.id = state['id']

            self.queued_properties = state['queued_properties']

            self.estimated_properties = state['estimated_properties']
            self.unsuccessful_properties = state['unsuccessful_properties']

            self.exceptions = state['exceptions']

            self.options = state['options']

            self.force_field_id = state['force_field_id']

            self.parameter_gradient_keys = state['parameter_gradient_keys']

    def __init__(self, calculation_backend, storage_backend,
                 port=8000, working_directory='working-data'):
        """Constructs a new PropertyEstimatorServer object.

        Parameters
        ----------
        calculation_backend: PropertyEstimatorBackend
            The backend to use for executing calculations.
        storage_backend: PropertyEstimatorStorage
            The backend to use for storing information from any calculations.
        port: int
            The port on which to listen for incoming client requests.
        working_directory: str
            The local directory in which to store all local, temporary calculation data.
        """

        assert calculation_backend is not None and storage_backend is not None

        self._calculation_backend = calculation_backend
        self._storage_backend = storage_backend

        self._port = port

        self._working_directory = working_directory

        if not path.isdir(self._working_directory):
            makedirs(self._working_directory)

        self._queued_calculations = {}
        self._finished_calculations = {}

        # Each client request id (i.e an id relating to a client requesting
        # that an entire data set of properties is estimated) is matched to
        # a set of server set request ids.
        #
        # The main difference is that on the server, a request to estimate
        # an entire data set is split into multiple requests to estimate
        # properties per substance.
        self._server_request_ids_per_client_id = {}

        super().__init__()

        self.bind(self._port)
        self.start(1)

        calculation_backend.start()

    async def _handle_job_submission(self, stream, address, message_length):
        """An asynchronous routine for handling the receiving and processing
        of job submissions from a client.

        Parameters
        ----------
        stream: IOStream
            An IO stream used to pass messages between the
            server and client.
        address: str
            The address from which the request came.
        message_length: int
            The length of the message being received.
        """

        logging.info('Received estimation request from {}'.format(address))

        # Read the incoming request from the server. The first four bytes
        # of the response should be the length of the message being sent.

        # Decode the client submission json.
        encoded_json = await stream.read_bytes(message_length)
        json_model = encoded_json.decode()

        # TODO: Add exception handling so the server can gracefully reject bad json.
        client_data_model = PropertyEstimatorSubmission.parse_json(json_model)

        client_request_id = str(uuid.uuid4())

        while client_request_id in self._server_request_ids_per_client_id:
            client_request_id = str(uuid.uuid4())

        self._server_request_ids_per_client_id[client_request_id] = []

        # Pass the ids of the submitted requests back to the
        # client.
        encoded_job_ids = json.dumps(client_request_id).encode()
        length = pack_int(len(encoded_job_ids))

        await stream.write(length + encoded_job_ids)

        logging.info('Request id sent to the client ({}): {}'.format(address, client_request_id))

        server_requests, request_ids_to_launch = self._prepare_server_requests(client_data_model,
                                                                               client_request_id)

        # Keep track of which server request ids belong to which client
        # request id.
        for request_id in request_ids_to_launch:
            self._schedule_server_request(server_requests[request_id])

    async def _handle_job_query(self, stream, message_length):
        """An asynchronous routine for handling the receiving and processing
        of job queries from a client

        Parameters
        ----------
        stream: IOStream
            An IO stream used to pass messages between the
            server and client.
        message_length: int
            The length of the message being received.
        """

        encoded_request_id = await stream.read_bytes(message_length)
        client_request_id = encoded_request_id.decode()

        response = None

        if client_request_id not in self._server_request_ids_per_client_id:

            response = PropertyEstimatorException(directory='',
                                                  message=f'The {client_request_id} request id was not found '
                                                          f'on the server.')

        else:
            response = self._query_client_request_status(client_request_id)

        encoded_response = response.json().encode()
        length = pack_int(len(encoded_response))

        await stream.write(length + encoded_response)

    async def handle_stream(self, stream, address):
        """A routine to handle incoming requests from
        a property estimator TCP client.

        Notes
        -----
        This method is based on the StackOverflow response from
        A. Jesse Jiryu Davis: https://stackoverflow.com/a/40257248

        Parameters
        ----------
        stream: IOStream
            An IO stream used to pass messages between the
            server and client.
        address: str
            The address from which the request came.
        """
        # logging.info("Incoming connection from {}".format(address))

        try:
            while True:

                # Receive an introductory message with the message type.
                packed_message_type = await stream.read_bytes(4)
                message_type_int = unpack_int(packed_message_type)[0]

                packed_message_length = await stream.read_bytes(4)
                message_length = unpack_int(packed_message_length)[0]

                # logging.info('Introductory packet recieved: {} {}'.format(message_type_int, message_length))

                message_type = None

                try:
                    message_type = PropertyEstimatorMessageTypes(message_type_int)
                    # logging.info('Message type: {}'.format(message_type))

                except ValueError as e:

                    logging.info('Bad message type recieved: {}'.format(e))

                    # Discard the unrecognised message.
                    if message_length > 0:
                        await stream.read_bytes(message_length)

                    continue

                if message_type is PropertyEstimatorMessageTypes.Submission:
                    await self._handle_job_submission(stream, address, message_length)
                elif message_type is PropertyEstimatorMessageTypes.Query:
                    await self._handle_job_query(stream, message_length)

        except StreamClosedError:

            # Handle client disconnections gracefully.
            # logging.info("Lost connection to {}:{} : {}.".format(address, self._port, e))
            pass

    def _find_server_estimation_request(self, request):
        """Checks whether the server is currently, or has previously completed
        a request to estimate a set of properties for a particular substance
        using the same force field parameters and estimation options.

        Parameters
        ----------
        request: PropertyEstimatorServer.ServerEstimationRequest
            The request to check for.

        Returns
        -------
        str, optional
            The id of the existing request if one exists, otherwise None.
        """

        cached_request_id = request.id

        for existing_id in self._queued_calculations:

            request.id = existing_id

            if request.json() != self._queued_calculations[existing_id].json():
                continue

            request.id = cached_request_id
            return existing_id

        for existing_id in self._finished_calculations:

            request.id = existing_id

            if request.json() != self._finished_calculations[existing_id].json():
                continue

            request.id = cached_request_id
            return existing_id

        request.id = cached_request_id
        return None

    def _prepare_server_requests(self, client_data_model, client_request_id):
        """Turns a client estimation submission request into a form more useful
        to the server, namely a list of properties to estimate separated by
        system composition.

        Parameters
        ----------
        client_data_model: PropertyEstimatorSubmission
            The client data model.
        client_request_id: str
            The id that was assigned to the client request.

        Returns
        -------
        dict of str and PropertyEstimatorServer.ServerEstimationRequest
            A list of the requests to be calculated by the server.
        list of str
            The ids of the requests which haven't already been launched by
            the server.
        """

        force_field_source = client_data_model.force_field_source
        force_field_id = self._storage_backend.has_force_field(force_field_source)

        if force_field_id is None:
            force_field_id = self._storage_backend.store_force_field(force_field_source)

        server_requests = {}

        # Split the full list of properties into lists partitioned by
        # substance.
        properties_by_substance = {}

        for physical_property in client_data_model.properties:

            if physical_property.substance.identifier not in properties_by_substance:
                properties_by_substance[physical_property.substance.identifier] = []

            properties_by_substance[physical_property.substance.identifier].append(physical_property)

        for substance_identifier in properties_by_substance:

            calculation_id = str(uuid.uuid4())

            # Make sure we don't somehow generate the same uuid
            # twice (although this is very unlikely to ever happen).
            while (calculation_id in self._queued_calculations or
                   calculation_id in self._finished_calculations):

                calculation_id = str(uuid.uuid4())

            properties_to_estimate = properties_by_substance[substance_identifier]

            options_copy = PropertyEstimatorOptions.parse_json(client_data_model.options.json())

            parameter_gradient_keys = copy.deepcopy(client_data_model.parameter_gradient_keys)

            request = self.ServerEstimationRequest(estimation_id=calculation_id,
                                                   queued_properties=properties_to_estimate,
                                                   options=options_copy,
                                                   force_field_id=force_field_id,
                                                   parameter_gradient_keys=parameter_gradient_keys)

            server_requests[calculation_id] = request

        request_ids_to_launch = []

        # Make sure this request is not already in the queue / has
        # already been completed, and if not add it to the list of
        # things to be queued.
        for server_request_id in server_requests:

            server_request = server_requests[server_request_id]
            existing_id = self._find_server_estimation_request(server_request)

            if existing_id is None:

                request_ids_to_launch.append(server_request_id)
                existing_id = server_request_id

                self._queued_calculations[server_request_id] = server_request

            self._server_request_ids_per_client_id[client_request_id].append(existing_id)

        return server_requests, request_ids_to_launch

    def _query_client_request_status(self, client_request_id):
        """Queries the current status of a client request by querying
        the state of the individual server requests it was split into.

        Parameters
        ----------
        client_request_id: str
            The id of the client request to query.

        Returns
        -------
        PropertyEstimatorResult
            The current results of the client request.
        """

        request_results = PropertyEstimatorResult(result_id=client_request_id)

        for server_request_id in self._server_request_ids_per_client_id[client_request_id]:

            server_request = None

            if server_request_id in self._queued_calculations:
                server_request = self._queued_calculations[server_request_id]

            elif server_request_id in self._finished_calculations:

                server_request = self._finished_calculations[server_request_id]

                if len(server_request.queued_properties) > 0:

                    return PropertyEstimatorException(message=f'An internal error occurred - the {server_request_id} '
                                                              f'was prematurely marked us finished.')

            else:

                return PropertyEstimatorException(message=f'An internal error occurred - the {server_request_id} '
                                                          f'request was not found on the server.')

            for physical_property in server_request.queued_properties:

                substance_id = physical_property.substance.identifier

                if substance_id not in request_results.queued_properties:
                    request_results.queued_properties[substance_id] = []

                request_results.queued_properties[substance_id].append(physical_property)

            for substance_id in server_request.unsuccessful_properties:

                physical_property = server_request.unsuccessful_properties[substance_id]

                if substance_id not in request_results.unsuccessful_properties:
                    request_results.unsuccessful_properties[substance_id] = []

                request_results.unsuccessful_properties[substance_id].append(physical_property)

            for substance_id in server_request.estimated_properties:

                physical_properties = server_request.estimated_properties[substance_id]

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

                request_results.estimated_properties[substance_id].extend(physical_properties)

            request_results.exceptions.extend(server_request.exceptions)

        return request_results

    def _schedule_server_request(self, server_request):
        """Schedules the estimation of the requested properties.

        This method will recursively cascade through all allowed calculation
        layers or until all properties have been calculated.

        Parameters
        ----------
        server_request : PropertyEstimatorServer.ServerEstimationRequest
            The object containing instructions about which calculations
            should be performed.
        """

        if len(server_request.options.allowed_calculation_layers) == 0 or \
           len(server_request.queued_properties) == 0:

            # Move any remaining properties to the unsuccessful list.
            for physical_property in server_request.queued_properties:

                substance_id = physical_property.substance.identifier

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

                server_request.unsuccessful_properties[substance_id].append(physical_property)

                server_request.queued_properties = []

            self._queued_calculations.pop(server_request.id)
            self._finished_calculations[server_request.id] = server_request

            logging.info(f'Finished server request {server_request.id}')
            return

        current_layer_type = server_request.options.allowed_calculation_layers.pop(0)

        if current_layer_type not in available_layers:

            # Kill all remaining properties if we reach an unsupported calculation layer.
            error_object = PropertyEstimatorException(message=f'The {current_layer_type} layer is not '
                                                              f'supported by / available on the server.')

            server_request.exceptions.append(error_object)

            server_request.options.allowed_calculation_layers.append(current_layer_type)
            server_request.queued_properties = []

            self._schedule_server_request(server_request)
            return

        logging.info(f'Launching server request {server_request.id} using the {current_layer_type} layer')

        layer_directory = path.join(self._working_directory, current_layer_type, server_request.id)

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

        current_layer = available_layers[current_layer_type]

        current_layer.schedule_calculation(self._calculation_backend,
                                           self._storage_backend,
                                           layer_directory,
                                           server_request,
                                           self._schedule_server_request)

    def start_listening_loop(self):
        """Starts the main (blocking) server IOLoop which will run until
        the user kills the process.
        """
        logging.info('Server listening at port {}'.format(self._port))

        try:
            IOLoop.current().start()
        except KeyboardInterrupt:
            self.stop()

    def stop(self):
        """Stops the property calculation server and it's
        provided backend.
        """
        self._calculation_backend.stop()
        IOLoop.current().stop()
