import time
import asyncio
import multiprocessing as mp

from functools import wraps
from concurrent.futures import ProcessPoolExecutor
from typing import Any, Coroutine, Callable, Optional

from .errors import MLServerError
from .settings import ModelSettings
from .model import MLModel
from .types import InferenceRequest, InferenceResponse
from .utils import get_wrapped_method
from .logging import logger

_InferencePoolAttr = "__inference_pool__"

# NOTE: Workaround for mypy
_mp_model: MLModel


class InvalidParallelMethod(MLServerError):
    def __init__(self, method_name: str, reason: Optional[str] = None):
        msg = f"Method {method_name} can't be parallelised"
        if reason:
            msg += f": {reason}"

        super().__init__(msg)


def _mp_load(model_settings: ModelSettings):
    """
    This method is meant to run internally in the multiprocessing workers.
    The loading needs to run synchronously, since the initializer argument
    doesn't support coroutines.
    """
    # NOTE: The global `_mp_model` variable is shared with the `_mp_predict`
    # method.
    # This global variable should only be used within the inference
    # multiprocessing workers.
    global _mp_model

    model_class = model_settings.implementation
    _mp_model = model_class(model_settings)  # type: ignore
    return asyncio.run(_mp_model.load())


def _mp_predict(payload: InferenceRequest) -> InferenceResponse:
    """
    This method is meant to run internally in the multiprocessing workers.
    The prediction needs to run synchronously, since multiprocessing
    doesn't know how to serialise coroutines.
    """
    # NOTE: `_mp_model` is a global variable initialised in the `_mp_load`
    # method.
    # This global variable is only to be used within the inference worker
    # context.
    global _mp_model
    return asyncio.run(_mp_model.predict(payload))


def _mp_noop():
    """
    This method is called to "warm workers."
    """
    # This sleep is important. Without it, workers in
    # the pool can finish the noop too quickly and run multiple operations.
    # Thus not "warming" all the workers in the pool
    time.sleep(0.1)
    return None


class InferencePool:
    """
    The InferencePool class represents a pool of workers where we can run
    inference on.

    Under the hood, it's responsible for managing a pool of multiprocessing
    workers, where the model is loaded.
    This approach lets MLServer work around the GIL to make sure that inference
    can occur in parallel across multiple models or instances of a model.
    """

    def __init__(self, model: MLModel):
        parallel_workers = model.settings.parallel_workers

        # Use 'spawn' instead of 'fork' to ensure that models are loaded in a
        # clean environment (e.g. to avoid issues like
        # https://github.com/tensorflow/tensorflow/issues/8220)
        ctx = mp.get_context("spawn")
        self._executor = ProcessPoolExecutor(
            max_workers=parallel_workers,
            mp_context=ctx,
            initializer=_mp_load,
            initargs=(model.settings,),
        )

    async def predict(self, payload: InferenceRequest) -> InferenceResponse:
        # What if we serialise payload?
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._executor, _mp_predict, payload)

    async def warm_workers(self, model: MLModel):
        logger.info(
            f"Warm workers enabled - loading {model.name} in inference pool workers"
        )
        loop = asyncio.get_event_loop()
        # Submit tasks to trigger initializer
        loading_tasks = [
            loop.run_in_executor(self._executor, _mp_noop)
            for _ in range(model.settings.parallel_workers)
        ]
        logger.info("Waiting for models to be loaded")
        setattr(self, "_warmed", True)
        return await asyncio.wait(loading_tasks)

    def __del__(self):
        self._executor.shutdown(wait=True)


def parallel(f: Callable[[InferenceRequest], Coroutine[Any, Any, InferenceResponse]]):
    """
    Decorator to attach to model's methods so that they run in parallel.
    By default, this will get attached to every model's "inference" method.

    NOTE: At the moment, this method only works with `predict()`.
    """
    # TODO: Extend to multiple methods
    @wraps(f)
    async def _inner(payload: InferenceRequest) -> InferenceResponse:
        wrapped_f = get_wrapped_method(f)
        if not hasattr(wrapped_f, "__self__"):
            raise InvalidParallelMethod(
                wrapped_f.__name__, reason="method is not bound"
            )

        model = getattr(wrapped_f, "__self__")
        if not hasattr(model, _InferencePoolAttr):
            raise InvalidParallelMethod(
                wrapped_f.__name__, reason="inference pool has not been loaded"
            )

        pool = getattr(model, _InferencePoolAttr)
        return await pool.predict(payload)

    return _inner


async def load_inference_pool(model: MLModel):
    if model.settings.parallel_workers == 0:
        # When parallel workers is set to 0, disable parallel inference
        return model

    pool = InferencePool(model)
    setattr(model, _InferencePoolAttr, pool)

    # Decorate predict method
    setattr(model, "predict", parallel(model.predict))

    # Conditionally load models to all workers in thread pool executor
    # This will cut down on initial inference response times but
    # will increase the amount of RAM utilized
    if model.settings.warm_workers:
        await pool.warm_workers(model)

    return model


async def unload_inference_pool(model: MLModel):
    pool = getattr(model, _InferencePoolAttr)
    if not pool:
        return

    pool.__del__()
    delattr(model, _InferencePoolAttr)
