# This file is part of Pebble.
# Copyright (c) 2013-2020, Matteo Cafasso

# Pebble is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation,
# either version 3 of the License, or (at your option) any later version.

# Pebble is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with Pebble.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import signal

from itertools import count
from functools import wraps
from multiprocessing import Pipe
from concurrent.futures import CancelledError, TimeoutError
try:
    from multiprocessing import get_start_method
except ImportError:
    def get_start_method():
        return 'spawn' if os.name == 'nt' else 'fork'

from pebble.common import ProcessExpired, ProcessFuture
from pebble.common import launch_process, stop_process, SLEEP_UNIT
from pebble.common import process_execute, launch_thread, send_result


def process(*args, **kwargs):
    """Runs the decorated function in a concurrent process,
    taking care of the result and error management.

    Decorated functions will return a concurrent.futures.Future object
    once called.

    The timeout parameter will set a maximum execution time
    for the decorated function. If the execution exceeds the timeout,
    the process will be stopped and the Future will raise TimeoutError.
    The name parameter will set the process name.
    """
    timeout = kwargs.get('timeout')
    name = kwargs.get('name')
    daemon = kwargs.get('daemon', True)

    # decorator without parameters
    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        return _process_wrapper(args[0], timeout, name, daemon)
    else:
        # decorator with parameters
        if timeout is not None and not isinstance(timeout, (int, float)):
            raise TypeError('Timeout expected to be None or integer or float')
        if name is not None and not isinstance(name, str):
            raise TypeError('Name expected to be None or string')
        if daemon is not None and not isinstance(daemon, bool):
            raise TypeError('Daemon expected to be None or bool')

        def decorating_function(function):
            return _process_wrapper(function, timeout, name, daemon)

        return decorating_function


def _process_wrapper(function, timeout, name, daemon):
    _register_function(function)

    @wraps(function)
    def wrapper(*args, **kwargs):
        future = ProcessFuture()
        reader, writer = Pipe(duplex=False)

        if get_start_method() != 'fork':
            target = _trampoline
            args = [function.__name__, function.__module__] + list(args)
        else:
            target = function

        worker = launch_process(
            name, _function_handler, daemon, target, args, kwargs, writer)

        writer.close()

        future.set_running_or_notify_cancel()

        launch_thread(name, _worker_handler, True, future, worker, reader, timeout)

        return future

    return wrapper


def _worker_handler(future, worker, pipe, timeout):
    """Worker lifecycle manager.

    Waits for the worker to be perform its task,
    collects result, runs the callback and cleans up the process.

    """
    result = _get_result(future, pipe, timeout)

    if isinstance(result, BaseException):
        if isinstance(result, ProcessExpired):
            result.exitcode = worker.exitcode

        future.set_exception(result)
    else:
        future.set_result(result)

    if worker.is_alive():
        stop_process(worker)


def _function_handler(function, args, kwargs, pipe):
    """Runs the actual function in separate process and returns its result."""
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    result = process_execute(function, *args, **kwargs)

    send_result(pipe, result)


def _get_result(future, pipe, timeout):
    """Waits for result and handles communication errors."""
    counter = count(step=SLEEP_UNIT)

    try:
        while not pipe.poll(SLEEP_UNIT):
            if timeout is not None and next(counter) >= timeout:
                return TimeoutError('Task Timeout', timeout)
            elif future.cancelled():
                return CancelledError()

        return pipe.recv()
    except (EOFError, OSError):
        return ProcessExpired('Abnormal termination')
    except Exception as error:
        return error


################################################################################
# Spawn process start method handling logic
################################################################################
_registered_functions = {}


def _register_function(function):
    global _registered_functions

    _registered_functions[function.__name__] = function


def _trampoline(name, module, *args, **kwargs):
    """Trampoline function for decorators.

    Lookups the function between the registered ones;
    if not found, forces its registering and then executes it.

    """
    function = _function_lookup(name, module)

    return function(*args, **kwargs)


def _function_lookup(name, module):
    """Searches the function between the registered ones.
    If not found, it imports the module forcing its registration.

    """
    try:
        return _registered_functions[name]
    except KeyError:  # force function registering
        __import__(module)
        mod = sys.modules[module]
        getattr(mod, name)

        return _registered_functions[name]
