import asyncio
import collections
import contextlib
import datetime
import functools
import io
import multiprocessing
import multiprocessing.pool
import os
import signal
import tempfile

from aiohttp import web
import marshmallow

from oslo_config import cfg
from oslo_log import log

LOG = log.getLogger(__name__)

CONF = cfg.CONF


UploadedFile = collections.namedtuple(
    "UploadedFile", ("name", "filename", "content_type", "original_filename")
)
"""Class to hold uploaded field metadata when passed to model's methods

.. py:attribute:: name

   Name of the argument where this file is being sent.

.. py:attribute:: filename

   Complete file path to the temporary file in the filesystem,

.. py:attribute:: content_type

   Content-type of the uploaded file

.. py:attribute:: original_filename

   Filename of the original file being uploaded.
"""

ReturnedFile = collections.namedtuple(
    "ReturnedFile", ("name", "filename", "content_type", "original_filename")
)
"""Class to pass the files returned from predict in a pickable way

.. py:attribute:: name

   Name of the argument where this file is being sent.

.. py:attribute:: filename

   Complete file path to the temporary file in the filesystem,

.. py:attribute:: content_type

   Content-type of the uploaded file

.. py:attribute:: original_filename

   Filename of the original file being uploaded.
"""


# set defaults to None, mainly for compatibility (vkoz)
UploadedFile.__new__.__defaults__ = (None, None, None, None)
ReturnedFile.__new__.__defaults__ = (None, None, None, None)


class ModelWrapper(object):
    """Class that will wrap the loaded models before exposing them.

    Whenever a model is loaded it will be wrapped with this class to create a
    wrapper object that will handle the calls to the model's methods so as to
    handle non-existent method exceptions.

    :param name: Model name
    :param model: Model object
    :raises HTTPInternalServerError: in case that a model has defined
        a response schema that is not JSON schema valid (DRAFT 4)
    """

    def __init__(self, name, model_obj, app=None):
        self.name = name
        self.model_obj = model_obj
        self._app = app

        self._loop = asyncio.get_event_loop()

        self._workers = CONF.workers
        self._executor = self._init_executor()

        if self._app is not None:
            self._setup_cleanup()

        schema = getattr(self.model_obj, "schema", None)

        if isinstance(schema, dict):
            try:
                schema = marshmallow.Schema.from_dict(
                    schema, name="ModelPredictionResponse"
                )
                self.has_schema = True
            except Exception as e:
                LOG.exception(e)
                raise web.HTTPInternalServerError(
                    reason=("Model defined schema is invalid, " "check server logs.")
                )
        elif schema is not None:
            try:
                if issubclass(schema, marshmallow.Schema):
                    self.has_schema = True
            except TypeError:
                raise web.HTTPInternalServerError(
                    reason=("Model defined schema is invalid, " "check server logs.")
                )
        else:
            self.has_schema = False

        self.response_schema = schema

    def _setup_cleanup(self):
        self._app.on_cleanup.append(self._close_executors)

    async def _close_executors(self, app):
        self._executor.shutdown()

    def _init_executor(self):
        n = self._workers
        executor = CancellablePool(max_workers=n)
        return executor

    @contextlib.contextmanager
    def _catch_error(self):
        name = self.name
        try:
            yield
        except AttributeError:
            raise web.HTTPNotImplemented(
                reason=("Not implemented by underlying model (loaded '%s')" % name)
            )
        except NotImplementedError:
            raise web.HTTPNotImplemented(
                reason=("Model '%s' does not implement this functionality" % name)
            )
        except Exception as e:
            LOG.error(
                "An exception has happened when calling method on " "'%s' model." % name
            )
            LOG.exception(e)
            if isinstance(e, web.HTTPException):
                raise e
            else:
                raise web.HTTPInternalServerError(reason=e)

    def validate_response(self, response):
        """Validate a response against the model's response schema, if set.

        If the wrapped model has defined a ``response`` attribute we will
        validate the response that

        :param response: The response that will be validated.
        :raises exceptions.InternalServerError: in case the reponse cannot be
            validated.
        """
        if self.has_schema is not True:
            raise web.HTTPInternalServerError(
                reason=(
                    "Trying to validate against a schema, but I do not "
                    "have one defined"
                )
            )

        try:
            self.response_schema().load(response)
        except marshmallow.ValidationError as e:
            LOG.exception(e)
            raise web.HTTPInternalServerError(
                reason="ERROR validating model response, check server logs."
            )
        except Exception as e:
            LOG.exception(e)
            raise web.HTTPInternalServerError(
                reason="Unknown ERROR validating response, check server logs."
            )

        return True

    def get_metadata(self):
        """Obtain model's metadata.

        If the model's metadata cannot be obtained because it is not
        implemented, we will provide some generic information so that the
        call does not fail.

        :returns dict: dictionary containing model's metadata
        """
        try:
            d = self.model_obj.get_metadata()
        except (NotImplementedError, AttributeError):
            d = {
                "id": "0",
                "name": self.name,
                "description": (
                    "Could not load description from "
                    "underlying model (loaded '%s')" % self.name
                ),
            }
        return d

    def _run_in_pool(self, func, *args, **kwargs):
        fn = functools.partial(func, *args, **kwargs)
        ret = self._loop.create_task(self._executor.apply(fn))
        return ret

    async def warm(self):
        """Warm (i.e. load, initialize) the underlying model.

        This method is called automatically when the model is loaded. You
        should use this method to initialize the model so that it is ready for
        the first prediction.

        The model receives no arguments.
        """
        try:
            func = self.model_obj.warm
        except AttributeError:
            LOG.debug("Cannot warm (initialize) model '%s'" % self.name)
            return

        try:
            n = self._workers
            LOG.debug("Warming '%s' model with %s workers" % (self.name, n))
            fs = [self._run_in_pool(func) for _ in range(0, n)]
            await asyncio.gather(*fs)
            LOG.debug("Model '%s' has been warmed" % self.name)
        except NotImplementedError:
            LOG.debug("Cannot warm (initialize) model '%s'" % self.name)

    @staticmethod
    def predict_wrap(predict_func, *args, **kwargs):
        """Wrapper function to allow returning files from predict
        This wrapper exists because buffer objects are not pickable,
        thus cannot be returned from the executor.
        """
        ret = predict_func(*args, **kwargs)
        if isinstance(ret, io.BufferedReader):
            ret = ReturnedFile(filename=ret.name)

        return ret

    def predict(self, *args, **kwargs):
        """Perform a prediction on wrapped model's ``predict`` method.

        :raises HTTPNotImplemented: If the method is not
            implemented in the wrapper model.
        :raises HTTPInternalServerError: If the call produces
            an error
        :raises HTTPException: If the call produces an
            error, already wrapped as a HTTPException
        """
        for key, val in kwargs.items():
            if isinstance(val, web.FileField):
                fd, name = tempfile.mkstemp()
                fd = os.fdopen(fd, "w+b")
                fd.write(val.file.read())
                fd.close()
                aux = UploadedFile(
                    name=val.name,
                    filename=name,
                    content_type=val.content_type,
                    original_filename=val.filename,
                )
                kwargs[key] = aux
                # FIXME(aloga); cleanup of tmpfile here

        with self._catch_error():
            return self._run_in_pool(
                self.predict_wrap, self.model_obj.predict, *args, **kwargs
            )

    def train(self, *args, **kwargs):
        """Perform a training on wrapped model's ``train`` method.

        :raises HTTPNotImplemented: If the method is not
            implemented in the wrapper model.
        :raises HTTPInternalServerError: If the call produces
            an error
        :raises HTTPException: If the call produces an
            error, already wrapped as a HTTPException
        """

        with self._catch_error():
            return self._run_in_pool(self.model_obj.train, *args, **kwargs)

    def get_train_args(self):
        """Add training arguments into the training parser.

        :param parser: an argparse like object

        This method will call the wrapped model ``add_train_args``.
        """
        try:
            args = self.model_obj.get_train_args()
        except (NotImplementedError, AttributeError):
            args = {}
        return args

    def get_predict_args(self):
        """Add predict arguments into the predict parser.

        :param parser: an argparse like object

        This method will call the wrapped model ``get_predict_args``.
        """
        try:
            args = self.model_obj.get_predict_args()
        except (NotImplementedError, AttributeError):
            args = {}
        return args


class NonDaemonProcess(multiprocessing.context.SpawnProcess):
    """Processes must use 'spawn' instead of 'fork' (which is the default
    in Linux) in order to work CUDA [1] or Tensorflow [2].

    [1] https://pytorch.org/docs/stable/notes/multiprocessing.html
    #cuda-in-multiprocessing
    [2] https://github.com/tensorflow/tensorflow/issues/5448
    #issuecomment-258934405
    """

    @property
    def daemon(self):
        return False

    @daemon.setter
    def daemon(self, value):
        pass


class NonDaemonPool(multiprocessing.pool.Pool):
    # Based on https://stackoverflow.com/questions/6974695/
    def Process(self, *args, **kwds):  # noqa
        proc = super(NonDaemonPool, self).Process(*args, **kwds)
        proc.__class__ = NonDaemonProcess

        return proc


class CancellablePool(object):
    def __init__(self, max_workers=None):
        self._free = {self._new_pool() for _ in range(max_workers)}
        self._working = set()
        self._change = asyncio.Event()

    def _new_pool(self):
        return NonDaemonPool(1, context=multiprocessing.get_context("spawn"))

    async def apply(self, fn, *args):
        """
        Like multiprocessing.Pool.apply_async, but:
         * is an asyncio coroutine
         * terminates the process if cancelled
        """
        while not self._free:
            await self._change.wait()
            self._change.clear()
        pool = usable_pool = self._free.pop()
        self._working.add(pool)

        loop = asyncio.get_event_loop()
        fut = loop.create_future()

        def _on_done(obj):
            ret = {"output": obj, "finish_date": str(datetime.datetime.now())}
            loop.call_soon_threadsafe(fut.set_result, ret)

        def _on_err(err):
            loop.call_soon_threadsafe(fut.set_exception, err)

        pool.apply_async(fn, args, callback=_on_done, error_callback=_on_err)

        try:
            return await fut
        except asyncio.CancelledError:
            # This is ugly, but since our pools only have one slot we can
            # kill the process before termination
            try:
                pool._pool[0].kill()
            except AttributeError:
                os.kill(pool._pool[0].pid, signal.SIGKILL)
            pool.terminate()
            usable_pool = self._new_pool()
        finally:
            self._working.remove(pool)
            self._free.add(usable_pool)
            self._change.set()

    def shutdown(self):
        for p in self._working:
            p.terminate()
        self._free.clear()
