"""
Based on :doc:`aiotools.context`, this module provides an automated lifecycle
management for multi-process servers with explicit initialization steps and
graceful shutdown steps.

.. function:: server(func)

    A decorator wrapper for :class:`AsyncServerContextManager`.

    Usage example:

    .. code:: python

       @aiotools.server
       async def myserver(loop, pidx, args):
           await do_init(args)
           stop_sig = yield
           if stop_sig == signal.SIGINT:
               await do_graceful_shutdown()
           else:
               await do_forced_shutdown()

       aiotools.start_server(myserver, ...)
"""

import asyncio
import functools
import inspect
import logging
import os
import signal
import struct
import sys
import threading
from contextlib import (
    AbstractContextManager, ContextDecorator,
)
from contextvars import ContextVar
from typing import (
    Any,
    Callable,
    Iterable,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
)

from .compat import all_tasks, current_task, get_running_loop
from .context import AbstractAsyncContextManager
from .fork import AbstractChildProcess, afork, _has_pidfd

__all__ = (
    'main',
    'start_server',
    'process_index',
    'AsyncServerContextManager',
    'ServerMainContextManager',
    'InterruptedBySignal',
)

log = logging.getLogger(__name__)

process_index: ContextVar[int] = ContextVar('process_index')


class InterruptedBySignal(BaseException):
    """
    A new :class:`BaseException` that represents interruption by an arbitrary UNIX
    signal.

    Since this is a :class:`BaseException` instead of :class:`Exception`, it behaves
    like :class:`KeyboardInterrupt` and :class:`SystemExit` exceptions (i.e.,
    bypassing except clauses catching the :class:`Exception` type only)

    The first argument of this exception is the signal number received.
    """
    pass


class AsyncServerContextManager(AbstractAsyncContextManager):
    """
    A modified version of :func:`contextlib.asynccontextmanager`.

    The implementation detail is mostly taken from the ``contextlib`` standard
    library, with a minor change to inject ``self.yield_return`` into the wrapped
    async generator.
    """

    yield_return: Optional[signal.Signals]

    def __init__(self, func: Callable[..., Any], args, kwargs):
        if not inspect.isasyncgenfunction(func):
            raise RuntimeError('Context manager function must be '
                               'an async-generator')
        self._agen = func(*args, **kwargs)
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.yield_return = None

    async def __aenter__(self):
        try:
            return (await self._agen.__anext__())
        except StopAsyncIteration:
            raise RuntimeError("async-generator didn't yield") from None

    async def __aexit__(self, exc_type, exc_value, tb):
        if exc_type is None:
            try:
                # Here is the modified part.
                await self._agen.asend(self.yield_return)
            except StopAsyncIteration:
                return
            else:
                raise RuntimeError("async-generator didn't stop") from None
        else:
            if exc_value is None:
                exc_value = exc_type()
            try:
                await self._agen.athrow(exc_type, exc_value, tb)
                raise RuntimeError("async-generator didn't stop after athrow()")
            except StopAsyncIteration as exc_new_value:
                return exc_new_value is not exc_value
            except RuntimeError as exc_new_value:
                if exc_new_value is exc_value:
                    return False
                if isinstance(exc_value, (StopIteration, StopAsyncIteration)):
                    if exc_new_value.__cause__ is exc_value:
                        return False
                raise
            except (BaseException, asyncio.CancelledError) as exc:
                if exc is not exc_value:
                    raise


class ServerMainContextManager(AbstractContextManager,
                               ContextDecorator):
    """
    A modified version of :func:`contextlib.contextmanager`.

    The implementation detail is mostly taken from the ``contextlib`` standard
    library, with a minor change to inject ``self.yield_return`` into the wrapped
    generator.
    """

    yield_return: Optional[signal.Signals]

    def __init__(self, func, args, kwargs):
        self.gen = func(*args, **kwargs)
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.yield_return = None

    def __enter__(self):
        del self.args, self.kwargs, self.func
        try:
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                self.gen.send(self.yield_return)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                value = type()
            try:
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    return False
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:  # noqa
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")


# This is a dirty hack to implement "module callable".
# NOTE: only works in Python 3.5 or higher.

def _server_ctxmgr(func):
    @functools.wraps(func)
    def helper(*args, **kwargs):
        return AsyncServerContextManager(func, args, kwargs)
    return helper


class _ServerModule(sys.modules[__name__].__class__):  # type: ignore
    def __call__(self, func):
        return _server_ctxmgr(func)


sys.modules[__name__].__class__ = _ServerModule


def _main_ctxmgr(func):
    """
    A decorator wrapper for :class:`ServerMainContextManager`

    Usage example:

    .. code:: python

       @aiotools.main
       def mymain():
           server_args = do_init()
           stop_sig = yield server_args
           if stop_sig == signal.SIGINT:
               do_graceful_shutdown()
           else:
               do_forced_shutdown()

       aiotools.start_server(..., main_ctxmgr=mymain, ...)
    """
    @functools.wraps(func)
    def helper(*args, **kwargs):
        return ServerMainContextManager(func, args, kwargs)
    return helper


main = _main_ctxmgr


def setup_child_watcher(loop: asyncio.AbstractEventLoop) -> None:
    try:
        watcher_cls = getattr(asyncio, 'PidfdChildWatcher', None)
        if _has_pidfd and watcher_cls:
            watcher = watcher_cls()
            asyncio.set_child_watcher(watcher)
        else:
            # Just get the default child watcher.
            watcher = asyncio.get_child_watcher()
        if not watcher.is_active():
            watcher.attach_loop(loop)
    except NotImplementedError:
        pass  # for uvloop


async def cancel_all_tasks() -> None:
    loop = get_running_loop()
    cancelled_tasks = []
    for task in all_tasks():
        if not task.done() and task is not current_task():
            task.cancel()
            cancelled_tasks.append(task)
    await asyncio.gather(*cancelled_tasks, return_exceptions=True)
    for task in cancelled_tasks:
        if task.cancelled():
            continue
        if task.exception() is not None:
            loop.call_exception_handler({
                'message': 'unhandled exception during loop shutdown',
                'exception': task.exception(),
                'task': task,
            })


def _worker_main(
    worker_actxmgr: Callable[
        [asyncio.AbstractEventLoop, int, Sequence[Any]],
        AsyncServerContextManager,
    ],
    stop_signals: Set[signal.Signals],
    intr_pipe_wfd: int,
    proc_idx: int,
    args: Sequence[Any],
) -> int:
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    setup_child_watcher(loop)
    interrupted = asyncio.Event()
    ctx = worker_actxmgr(loop, proc_idx, args)
    process_index.set(proc_idx)
    forever_future = loop.create_future()

    def handle_stop_signal(signum):
        if interrupted.is_set():
            pass
        else:
            interrupted.set()
            ctx.yield_return = signum
            forever_future.cancel()

    for signum in stop_signals:
        loop.add_signal_handler(
            signum,
            functools.partial(handle_stop_signal, signum),
        )
    # Allow the worker to be interrupted during initialization
    # (in case of initialization failures in other workers)
    signal.pthread_sigmask(signal.SIG_UNBLOCK, stop_signals)

    async def _wrapped_worker():
        err_ctx = 'enter'
        try:
            async with ctx:
                err_ctx = 'body'
                try:
                    await forever_future
                except asyncio.CancelledError:
                    pass
                finally:
                    err_ctx = 'exit'
        except Exception:
            if err_ctx != 'body':
                err_ctx_str = 'initialization' if err_ctx == 'enter' else 'shutdown'
                log.exception(f'Worker {proc_idx}: '
                              f'Error during context manager {err_ctx_str}')
                os.write(intr_pipe_wfd, struct.pack('i', proc_idx))
            raise

    try:
        loop.run_until_complete(_wrapped_worker())
    finally:
        try:
            loop.run_until_complete(cancel_all_tasks())
            loop.run_until_complete(loop.shutdown_asyncgens())
            try:
                loop.run_until_complete(loop.shutdown_default_executor())
            except (AttributeError, NotImplementedError):  # for uvloop
                pass
        finally:
            loop.close()
        return 0


def _extra_main(
    extra_func,
    stop_signals: Set[signal.Signals],
    proc_idx: int,
    args: Sequence[Any],
) -> int:
    interrupted = threading.Event()
    process_index.set(proc_idx)

    # Since signals only work for the main thread in Python,
    # extra processes in use_threading=True mode should check
    # the intr_event by themselves (probably in their loops).

    def raise_stop(signum, frame):
        if interrupted.is_set():
            pass
        else:
            interrupted.set()
            if signum == signal.SIGINT:
                raise KeyboardInterrupt
            elif signum == signal.SIGTERM:
                raise SystemExit
            else:
                raise InterruptedBySignal(signum)

    # restore signal handler.
    for signum in stop_signals:
        signal.signal(signum, raise_stop)
    signal.pthread_sigmask(signal.SIG_UNBLOCK, stop_signals)
    intr_event = None

    try:
        if not interrupted.is_set():
            extra_func(intr_event, proc_idx, args)
    except (SystemExit, KeyboardInterrupt, InterruptedBySignal):
        log.warning(f'extra_proc[{proc_idx}] did not handle stop signals.')
    finally:
        # same as in _worker_main()
        signal.pthread_sigmask(signal.SIG_BLOCK, stop_signals)
        return 0


def start_server(
    worker_actxmgr: Callable[
        [asyncio.AbstractEventLoop, int, Sequence[Any]],
        AsyncServerContextManager],
    main_ctxmgr: Optional[Callable[[], ServerMainContextManager]] = None,
    extra_procs: Iterable[Callable] = tuple(),
    stop_signals: Iterable[signal.Signals] = (
        signal.SIGINT,
        signal.SIGTERM
    ),
    num_workers: int = 1,
    args: Iterable[Any] = tuple(),
    wait_timeout: Optional[float] = None,
) -> None:
    """
    Starts a multi-process server where each process has their own individual
    asyncio event loop.  Their lifecycles are automantically managed -- if the
    main program receives one of the signals specified in ``stop_signals`` it
    will initiate the shutdown routines on each worker that stops the event
    loop gracefully.

    Args:
        worker_actxmgr: An asynchronous context manager that dicates the
                        initialization and shutdown steps of each worker.
                        It should accept the following three arguments:

                        * **loop**: the asyncio event loop created and set
                          by aiotools
                        * **pidx**: the 0-based index of the worker
                          (use this for per-worker logging)
                        * **args**: a concatenated tuple of values yielded by
                          **main_ctxmgr** and the user-defined arguments in
                          **args**.

                        aiotools automatically installs an interruption handler
                        that calls ``loop.stop()`` to the given event loop,
                        regardless of using either threading or
                        multiprocessing.

        main_ctxmgr: An optional context manager that performs global
                     initialization and shutdown steps of the whole program.
                     It may yield one or more values to be passed to worker
                     processes along with **args** passed to this function.
                     There is no arguments passed to those functions since
                     you can directly access ``sys.argv`` to parse command
                     line arguments and/or read user configurations.

        extra_procs: An iterable of functions that consist of extra processes
                     whose lifecycles are synchronized with other workers.
                     They should set up their own signal handlers.

                     It should accept the following three arguments:

                     * **intr_event**: Always ``None``, kept for legacy
                     * **pidx**: same to **worker_actxmgr** argument
                     * **args**: same to **worker_actxmgr** argument

        stop_signals: A list of UNIX signals that the main program to
                      recognize as termination signals.

        num_workers: The number of children workers.

        args: The user-defined arguments passed to workers and extra
              processes.  If **main_ctxmgr** yields one or more values,
              they are *prepended* to this user arguments when passed to
              workers and extra processes.

        wait_timeout: The timeout in seconds before forcibly killing all
                      remaining child processes after sending initial stop
                      signals.

    Returns:
        None

    .. versionchanged:: 0.3.2

       The name of argument **num_proc** is changed to **num_workers**.
       Even if **num_workers** is 1, a child is created instead of
       doing everything at the main thread.

    .. versionadded:: 0.3.2

       The argument ``extra_procs`` and ``main_ctxmgr``.

    .. versionadded:: 0.4.0

       Now supports use of threading instead of multiprocessing via
       **use_threading** option.

    .. versionchanged:: 0.8.0

       Now **worker_actxmgr** must be an instance of
       :class:`AsyncServerContextManager` or async generators decorated by
       ``@aiotools.server``.

       Now **main_ctxmgr** must be an instance of :class:`ServerMainContextManager`
       or plain generators decorated by ``@aiotools.main``.

       The usage is same to asynchronous context managers, but optionally you can
       distinguish the received stop signal by retrieving the return value of the
       ``yield`` statement.

       In **extra_procs** in non-threaded mode, stop signals are converted into
       either one of :class:`KeyboardInterrupt`, :class:`SystemExit`, or
       :class:`InterruptedBySignal` exception.

    .. versionadded:: 0.8.4

       **start_method** argument can be set to change the subprocess spawning
       implementation.

    .. deprecated:: 1.2.0

       The **start_method** and **use_threading** arguments, in favor of our new
       :func:`afork()` function which provides better synchronization and pid-fd
       support.

    .. versionchanged:: 1.2.0

       The **extra_procs** will be always separate processes since **use_threading**
       is deprecated and thus **intr_event** arguments are now always ``None``.

    .. versionadded:: 1.5.5

       The **wait_timeout** argument.
    """

    @_main_ctxmgr
    def noop_main_ctxmgr():
        yield

    assert stop_signals

    if hasattr(asyncio, 'get_running_loop'):
        # only for Python 3.7+
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            # there should be none.
            pass
        else:
            raise RuntimeError(
                'aiotools.start_server() cannot be called inside '
                'a running event loop.')

    if main_ctxmgr is None:
        main_ctxmgr = noop_main_ctxmgr

    children: List[AbstractChildProcess] = []
    sigblock_mask = frozenset(stop_signals)

    main_ctx = main_ctxmgr()

    # temporarily block signals and register signal handlers to main_loop
    signal.pthread_sigmask(signal.SIG_BLOCK, sigblock_mask)

    main_loop = asyncio.new_event_loop()
    asyncio.set_event_loop(main_loop)
    main_future = main_loop.create_future()

    # to make subprocess working in child threads
    setup_child_watcher(main_loop)

    # build a main-to-worker interrupt channel using signals
    def handle_stop_signal(signum: signal.Signals) -> None:
        main_ctx.yield_return = signum  # type: ignore
        for child in children:
            child.send_signal(signum)
        # instead of main_loop.stop(), we use a future here
        # NOT to interrupt the cleanup routines in
        # an arbitrary timing upon async child failures.
        main_future.cancel()

    for signum in stop_signals:
        main_loop.add_signal_handler(
            signum,
            functools.partial(handle_stop_signal, signum),
        )

    # build a reliable worker-to-main interrupt channel using a pipe
    # (workers have no idea whether the main interrupt is enabled/disabled)
    def handle_child_interrupt(fd: int) -> None:
        child_idx = struct.unpack('i', os.read(fd, 4))[0]  # noqa
        log.debug(f'Child {child_idx} has interrupted the main program.')
        # self-interrupt to initiate the main-to-worker interrupts
        signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT})
        os.kill(0, signal.SIGINT)

    child_intr_pipe: Tuple[int, int] = os.pipe()
    rfd = child_intr_pipe[0]
    main_loop.add_reader(rfd, handle_child_interrupt, rfd)

    # start
    with main_ctx as main_args:

        # retrieve args generated by the user-defined main
        if main_args is None:
            main_args = tuple()
        if not isinstance(main_args, tuple):
            main_args = (main_args, )

        # spawn managed async workers
        for i in range(num_workers):
            try:
                p = main_loop.run_until_complete(afork(functools.partial(
                    _worker_main,
                    worker_actxmgr,
                    stop_signals,
                    child_intr_pipe[1],
                    i,
                    main_args + args,
                )))
            except RuntimeError as e:
                if 'loop stopped' in e.args[0]:
                    log.warning(
                        "skipping spawning of child[worker]:%d due to "
                        "async failure(s) of other child process",
                        i,
                    )
                    continue
                raise
            children.append(p)

        # spawn extra workers
        for i, f in enumerate(extra_procs):
            try:
                p = main_loop.run_until_complete(afork(functools.partial(
                    _extra_main,
                    f,
                    stop_signals,
                    num_workers + i,
                    main_args + args,
                )))
            except RuntimeError as e:
                if 'loop stopped' in e.args[0]:
                    log.warning(
                        "skipping spawning of child[extra]:%d due to "
                        "async failure(s) of other child process",
                        num_workers + i,
                    )
                    continue
                raise
            children.append(p)
        try:
            # unblock the stop signals for user/external interrupts.
            signal.pthread_sigmask(signal.SIG_UNBLOCK, sigblock_mask)

            # run!
            try:
                main_loop.run_until_complete(main_future)
            except asyncio.CancelledError:
                pass
            finally:
                # if interrupted, wait for workers to finish.
                try:
                    main_loop.run_until_complete(
                        asyncio.wait_for(asyncio.gather(
                            *[child.wait() for child in children],
                            return_exceptions=True,
                        ), wait_timeout)
                    )
                except asyncio.TimeoutError:
                    log.warning(
                        "Timeout during waiting for child processes; "
                        "killing all",
                    )
                    for child in children:
                        child.send_signal(signal.SIGKILL)
        finally:
            try:
                main_loop.run_until_complete(cancel_all_tasks())
                main_loop.run_until_complete(main_loop.shutdown_asyncgens())
                try:
                    main_loop.run_until_complete(
                        main_loop.shutdown_default_executor()
                    )
                except (AttributeError, NotImplementedError):  # for uvloop
                    pass
            finally:
                main_loop.close()
                asyncio.set_event_loop(None)
