"""
The core functionality of the 'server' side of the
openff-evaluator framework.
"""
import copy
import json
import logging
import os
import select
import socket
import threading
import traceback
import uuid

import networkx

from openff.evaluator.attributes import Attribute, AttributeClass
from openff.evaluator.client import (
    BatchMode,
    EvaluatorClient,
    RequestOptions,
    RequestResult,
)
from openff.evaluator.datasets import PhysicalProperty
from openff.evaluator.forcefield import ParameterGradientKey
from openff.evaluator.layers import registered_calculation_layers
from openff.evaluator.storage import LocalFileStorage
from openff.evaluator.utils.exceptions import EvaluatorException
from openff.evaluator.utils.serialization import TypedJSONEncoder
from openff.evaluator.utils.tcp import (
    EvaluatorMessageTypes,
    pack_int,
    recvall,
    unpack_int,
)

logger = logging.getLogger(__name__)


class Batch(AttributeClass):
    """Represents a batch of physical properties which are being estimated by
    the server for a given set of force field parameters.

    The expectation is that this object will be passed between calculation layers,
    whereby each layer will attempt to estimate each of the `queued_properties`.
    Those properties which can be estimated will be moved to the `estimated_properties`
    set, while those that couldn't will remain in the `queued_properties` set ready
    for the next layer.
    """

    id = Attribute(
        docstring="The unique id of this batch.",
        type_hint=str,
        default_value=lambda: str(uuid.uuid4()).replace("-", ""),
    )

    force_field_id = Attribute(
        docstring="The id of the force field being used to estimate"
        "this batch of properties.",
        type_hint=str,
    )
    options = Attribute(
        docstring="The options being used to estimate this batch.",
        type_hint=RequestOptions,
    )
    parameter_gradient_keys = Attribute(
        docstring="The parameters that this batch of physical properties "
        "should be differentiated with respect to.",
        type_hint=list,
    )

    queued_properties = Attribute(
        docstring="The set of properties which have yet to be estimated.",
        type_hint=list,
        default_value=[],
    )
    estimated_properties = Attribute(
        docstring="The set of properties which have been successfully estimated.",
        type_hint=list,
        default_value=[],
    )
    unsuccessful_properties = Attribute(
        docstring="The set of properties which have been could not be estimated.",
        type_hint=list,
        default_value=[],
    )
    exceptions = Attribute(
        docstring="The set of properties which have yet to be, or "
        "are currently being estimated.",
        type_hint=list,
        default_value=[],
    )

    def validate(self, attribute_type=None):
        super(Batch, self).validate(attribute_type)

        assert all(isinstance(x, PhysicalProperty) for x in self.queued_properties)
        assert all(isinstance(x, PhysicalProperty) for x in self.estimated_properties)
        assert all(
            isinstance(x, PhysicalProperty) for x in self.unsuccessful_properties
        )
        assert all(isinstance(x, EvaluatorException) for x in self.exceptions)
        assert all(
            isinstance(x, ParameterGradientKey) for x in self.parameter_gradient_keys
        )


class EvaluatorServer:
    """The object responsible for coordinating all properties estimations to
    be ran using the openff-evaluator framework.

    This server is responsible for receiving estimation requests from the client,
    determining which calculation layer to use to launch the request, and
    distributing that estimation across the available compute resources.

    Notes
    -----
    Every client request is split into logical chunk batches. This enables batches
    of related properties (e.g. all properties for CO) to be estimated in one go
    (or one task graph in the case of workflow based layers) and returned when ready,
    rather than waiting for the full data set to complete.

    Examples
    --------
    Setting up a general server instance using a dask based calculation backend,
    and a local file storage backend:

    >>> # Create the backend which will be responsible for distributing the calculations
    >>> from openff.evaluator.backends.dask import DaskLocalCluster
    >>> calculation_backend = DaskLocalCluster()
    >>> calculation_backend.start()
    >>>
    >>> # Create the server to which all estimation requests will be submitted
    >>> from openff.evaluator.server import EvaluatorServer
    >>> property_server = EvaluatorServer(calculation_backend)
    >>>
    >>> # Instruct the server to listen for incoming requests
    >>> # This command will run until killed.
    >>> property_server.start()
    """

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

        Parameters
        ----------
        calculation_backend: CalculationBackend
            The backend to use for executing calculations.
        storage_backend: StorageBackend, optional
            The backend to use for storing information from any calculations.
            If `None`, a default `LocalFileStorage` backend will be used.
        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.
        """

        # Initialize the main 'server' attributes.
        self._port = port

        self._server_thread = None
        self._socket = None

        self._started = False
        self._stopped = True

        # Initialize the internal components.
        assert calculation_backend is not None and calculation_backend.started
        self._calculation_backend = calculation_backend

        if storage_backend is None:
            storage_backend = LocalFileStorage()

        self._storage_backend = storage_backend

        self._working_directory = working_directory
        os.makedirs(self._working_directory, exist_ok=True)

        self._queued_batches = {}
        self._finished_batches = {}

        self._batch_ids_per_client_id = {}

    def _query_request_status(self, client_request_id):
        """Queries the the current state of an estimation request
        and stores it in a `RequestResult`.

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

        Returns
        -------
        RequestResult
            The state of the request.
        EvaluatorException, optional
            The exception raised while retrieving the status,
            if any.
        """

        request_results = RequestResult()

        for batch_id in self._batch_ids_per_client_id[client_request_id]:

            # Find the batch.
            if batch_id in self._queued_batches:
                batch = self._queued_batches[batch_id]

            elif batch_id in self._finished_batches:

                batch = self._finished_batches[batch_id]

                if len(batch.queued_properties) > 0:

                    return (
                        None,
                        EvaluatorException(
                            message=f"An internal error occurred - the {batch_id} "
                            f"batch was prematurely marked us finished."
                        ),
                    )

            else:

                return (
                    None,
                    EvaluatorException(
                        message=f"An internal error occurred - the {batch_id} "
                        f"request was not found on the server."
                    ),
                )

            request_results.queued_properties.add_properties(*batch.queued_properties)
            request_results.unsuccessful_properties.add_properties(
                *batch.unsuccessful_properties
            )
            request_results.estimated_properties.add_properties(
                *batch.estimated_properties
            )
            request_results.exceptions.extend(batch.exceptions)

        return request_results, None

    def _batch_by_same_component(self, submission, force_field_id):
        """Batches a set of requested properties based on which substance they were
        measured for. Properties which were measured for substances containing the
        exact same components (but not necessarily in the same amounts) will be placed
        into the same batch.

        Parameters
        ----------
        submission: EvaluatorClient._Submission
            The full request submission.
        force_field_id: str
            The unique id of the force field to use.

        Returns
        -------
        list of Batch
            The property batches.
        """

        reserved_batch_ids = {
            *self._queued_batches.keys(),
            *self._finished_batches.keys(),
        }

        batches = []

        for substance in submission.dataset.substances:

            batch = Batch()
            batch.force_field_id = force_field_id

            # Make sure we don't somehow generate the same uuid
            # twice (although this is very unlikely to ever happen).
            while batch.id in reserved_batch_ids:
                batch.id = str(uuid.uuid4()).replace("-", "")

            batch.queued_properties = [
                x for x in submission.dataset.properties_by_substance(substance)
            ]
            batch.options = RequestOptions.parse_json(submission.options.json())

            batch.parameter_gradient_keys = copy.deepcopy(
                submission.parameter_gradient_keys
            )

            reserved_batch_ids.add(batch.id)
            batches.append(batch)

        return batches

    def _batch_by_shared_component(self, submission, force_field_id):
        """Batches a set of requested properties based on which substance they were
        measured for. Properties which were measured for substances sharing at least
        one common component (defined only by its smiles pattern and not necessarily
        in the same amount) will be placed into the same batch.

        Parameters
        ----------
        submission: EvaluatorClient._Submission
            The full request submission.
        force_field_id: str
            The unique id of the force field to use.

        Returns
        -------
        list of Batch
            The property batches.
        """

        reserved_batch_ids = {
            *self._queued_batches.keys(),
            *self._finished_batches.keys(),
        }

        all_smiles = set(x.smiles for y in submission.dataset.substances for x in y)

        # Build a graph containing all of the different component
        # smiles patterns as nodes.
        substance_graph = networkx.Graph()
        substance_graph.add_nodes_from(all_smiles)

        # Add edges to the graph based on which substances contain
        # the different component nodes.
        for substance in submission.dataset.substances:

            if len(substance) < 2:
                continue

            smiles = [x.smiles for x in substance]

            for smiles_a, smiles_b in zip(smiles, smiles[1:]):
                substance_graph.add_edge(smiles_a, smiles_b)

        # Find clustered islands of those smiles which exist in
        # overlapping substances.
        islands = [
            substance_graph.subgraph(c)
            for c in networkx.connected_components(substance_graph)
        ]

        # Create one batch per island
        batches = []

        for _ in range(len(islands)):

            batch = Batch()
            batch.force_field_id = force_field_id

            # Make sure we don't somehow generate the same uuid
            # twice (although this is very unlikely to ever happen).
            while batch.id in reserved_batch_ids:
                batch.id = str(uuid.uuid4()).replace("-", "")

            batch.options = RequestOptions.parse_json(submission.options.json())

            batch.parameter_gradient_keys = copy.deepcopy(
                submission.parameter_gradient_keys
            )

            reserved_batch_ids.add(batch.id)
            batches.append(batch)

        for physical_property in submission.dataset:

            smiles = [x.smiles for x in physical_property.substance]

            island_id = 0

            for island_id, island in enumerate(islands):

                if not any(x in island for x in smiles):
                    continue

                break

            batches[island_id].queued_properties.append(physical_property)

        return batches

    def _prepare_batches(self, submission, request_id):
        """Turns an estimation request into chunked batches to
        calculate separately.

        This enables batches of related properties (e.g. all properties
        for CO) to be estimated in one go (or one task graph in the case
        of workflow based layers) and returned when ready, rather than waiting
        for the full data set to complete.

        Parameters
        ----------
        submission: EvaluatorClient._Submission
            The full request submission.
        request_id: str
            The id that was assigned to the request.

        Returns
        -------
        list of Batch
            A list of the batches to launch.
        """

        force_field_source = submission.force_field_source
        force_field_id = self._storage_backend.store_force_field(force_field_source)

        batch_mode = submission.options.batch_mode

        if batch_mode == BatchMode.SameComponents:
            batches = self._batch_by_same_component(submission, force_field_id)
        elif batch_mode == BatchMode.SharedComponents:
            batches = self._batch_by_shared_component(submission, force_field_id)
        else:
            raise NotImplementedError()

        for batch in batches:

            self._queued_batches[batch.id] = batch
            self._batch_ids_per_client_id[request_id].append(batch.id)

        return batches

    def _launch_batch(self, batch):
        """Launch a batch of properties to estimate.

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

        Parameters
        ----------
        batch : Batch
            The batch to launch.
        """

        if (
            len(batch.options.calculation_layers) == 0
            or len(batch.queued_properties) == 0
        ):

            # Move any remaining properties to the unsuccessful list.
            batch.unsuccessful_properties = [*batch.queued_properties]
            batch.queued_properties = []

            self._queued_batches.pop(batch.id)
            self._finished_batches[batch.id] = batch

            logger.info(f"Finished server request {batch.id}")
            return

        current_layer_type = batch.options.calculation_layers.pop(0)

        if current_layer_type not in registered_calculation_layers:

            # Add an exception if we reach an unsupported calculation layer.
            error_object = EvaluatorException(
                message=f"The {current_layer_type} layer is not "
                f"supported by / available on the server."
            )

            batch.exceptions.append(error_object)
            self._launch_batch(batch)
            return

        logger.info(f"Launching batch {batch.id} using the {current_layer_type} layer")

        layer_directory = os.path.join(
            self._working_directory, current_layer_type, batch.id
        )
        os.makedirs(layer_directory, exist_ok=True)

        current_layer = registered_calculation_layers[current_layer_type]

        current_layer.schedule_calculation(
            self._calculation_backend,
            self._storage_backend,
            layer_directory,
            batch,
            self._launch_batch,
        )

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

        Parameters
        ----------
        connection:
            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.
        """

        logger.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 = recvall(connection, message_length)
        json_model = encoded_json.decode()

        request_id = None
        error = None

        try:

            # noinspection PyProtectedMember
            submission = EvaluatorClient._Submission.parse_json(json_model)
            submission.validate()

        except Exception as e:

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

            error = EvaluatorException(
                message=f"An exception occured when parsing "
                f"the submission: {formatted_exception}"
            )

            submission = None

        if error is None:

            while request_id is None or request_id in self._batch_ids_per_client_id:
                request_id = str(uuid.uuid4()).replace("-", "")

            self._batch_ids_per_client_id[request_id] = []

        # Pass the id of the submitted requests back to the client
        # as well as any error which may have occurred.
        return_packet = json.dumps((request_id, error), cls=TypedJSONEncoder)

        encoded_return_packet = return_packet.encode()
        length = pack_int(len(encoded_return_packet))

        connection.sendall(length + encoded_return_packet)

        if error is not None:
            # Exit early if there is an error.
            return

        # Batch the request into more managable chunks.
        batches = self._prepare_batches(submission, request_id)

        # Launch the batches
        for batch in batches:
            self._launch_batch(batch)

    def _handle_job_query(self, connection, message_length):
        """An asynchronous routine for handling the receiving and
        processing of request status queries from a client

        Parameters
        ----------
        connection:
            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 = recvall(connection, message_length)
        client_request_id = encoded_request_id.decode()

        response = None

        if client_request_id not in self._batch_ids_per_client_id:

            error = EvaluatorException(
                message=f"The request id ({client_request_id}) was not found "
                f"on the server.",
            )

        else:
            response, error = self._query_request_status(client_request_id)

        response_json = json.dumps((response, error), cls=TypedJSONEncoder)

        encoded_response = response_json.encode()
        length = pack_int(len(encoded_response))

        connection.sendall(length + encoded_response)

    def _handle_stream(self, connection, address):
        """A routine to handle incoming requests from
        a TCP client.
        """

        # Receive an introductory message with the message type.
        packed_message_type = recvall(connection, 4)
        message_type_int = unpack_int(packed_message_type)[0]

        packed_message_length = recvall(connection, 4)
        message_length = unpack_int(packed_message_length)[0]

        try:
            message_type = EvaluatorMessageTypes(message_type_int)
        except ValueError as e:

            trace = traceback.format_exception(None, e, e.__traceback__)
            logger.info(f"Bad message type received: {trace}")

            # Discard the unrecognised message.
            if message_length > 0:
                recvall(connection, message_length)

            return

        if message_type is EvaluatorMessageTypes.Submission:
            self._handle_job_submission(connection, address, message_length)
        elif message_type is EvaluatorMessageTypes.Query:
            self._handle_job_query(connection, message_length)

    def _handle_connections(self):
        """Handles incoming client TCP connections.
        """
        to_read = [self._socket]

        try:

            while not self._stopped:

                ready, _, _ = select.select(to_read, [], [], 0.1)

                for data in ready:

                    if data == self._socket:
                        connection, address = self._socket.accept()
                        to_read.append(connection)

                    else:

                        connection = data
                        self._handle_stream(connection, connection.getpeername())
                        connection.close()

                        to_read.remove(data)

        except Exception:
            logger.exception("Fatal error in the main server loop")

    def start(self, asynchronous=False):
        """Instructs the server to begin listening for incoming
        requests from any `EvaluatorClients`.

        Parameters
        ----------
        asynchronous: bool
            If `True` the server will run on a separate thread in the background,
            returning control back to the main thread. Otherwise, this function
            will block the main thread until this server is killed.
        """

        if self._started:
            raise RuntimeError("The server has already been started.")

        logger.info("Server listening at port {}".format(self._port))
        self._started = True
        self._stopped = False

        # Create the TCP socket
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._socket.bind(("localhost", self._port))
        self._socket.listen(128)

        try:

            if asynchronous:

                self._server_thread = threading.Thread(
                    target=self._handle_connections, daemon=True
                )
                self._server_thread.start()

            else:
                self._handle_connections()

        except KeyboardInterrupt:
            self.stop()

    def stop(self):
        """Stops the property calculation server and it's
        provided backend.
        """
        if not self._started:
            raise ValueError("The server has not yet been started.")

        self._stopped = True
        self._started = False

        if self._server_thread is not None:

            self._server_thread.join()
            self._server_thread = None

        if self._socket is not None:

            self._socket.close()
            self._socket = None

    def __enter__(self):
        self.start(asynchronous=True)
        return self

    def __exit__(self, *args):
        self.stop()

    def __del__(self):

        if self._started and not self._stopped:
            self.stop()
