# Copyright 2019 VMware, Inc.  All rights reserved. -- VMware Confidential
#
''' Module responsible to invoke shutdown hooks in preliminary defined order on
shutdown signal.
'''
import logging
import time
import multiprocessing
import collections
import errno
import os
import signal
import transport

logger = logging.getLogger(__name__)

class ShutdownhookPriority(object):
    '''Shutdown hook order definition.
    '''
    # Component hooks are executed first
    COMPONENT_HOOK = 0
    # Then system shutdown hooks(a.k.a SDK) get executed
    SYSTEM_HOOK = 100
    # Then installer shutdown hooks get executed
    INSTALLER_HOOK = 200

class InterruptedCallError(IOError):
    '''Called when the program need to be interrupted.
    '''

class _ShutdownHook(object):
    '''Shutdownhook defintion.
    '''
    def __init__(self, callback, context, timeout, priority):
        '''Creates shutdown hook.

        @param callback: Shutdown callback handler
        @type callback: Function having the follow definition:
          def _myShutdownHook(context):
              pass

        @param context: Shutdown context passed to the shutdown hook
        @type context: Any object

        @param timeout: Timeout in seconds is the window in which the shutdown
          hook must complete
        @type timeout: int

        @param priority: The parameter specifies the priority with which the hook
          has to be executed in case of shutdown. 0es priority hooks will be
          executed first, then 1es priority and so on. If two hooks have same
          priority the order of their execution is undefined.
        @type priority: int
        '''
        self.callback = callback
        self.context = context
        self.timeout = timeout
        self.priority = priority

class _ExecutedProcess(object):
    '''Definition of started process.
    '''
    def __init__(self, p, startedAt, timeout):
        '''Creates an already started process.

        @param p: Already started process
        @type p: multiprocessing.Process

        @param startedAt: The time when the process has been started.
        @type startedAt: long

        @param timeout: The time in seconds for which the process has to complete
        @type timeout: long
        '''
        self.p = p
        self.callback = p._target
        self.startedAt = startedAt
        self.timeout = timeout

    def _terminate(self):
        '''Terminate the started process.
        '''
        self.p.terminate()
        # Wait until it fully completes
        self.p.join()

    def isCompleted(self):
        '''Check if the process has completed or its timeout elapsed.

        @return: True if process completed or in case it timeout elapsed and has
          been terminated, otherwise False
        @rtype: bool
        '''
        elapsedTime = time.time() - self.startedAt
        remainingTime = self.timeout - elapsedTime
        if not self.p.is_alive():
            logger.info('Shutdown hook %s has completed', self.callback)
            return True

        if remainingTime >= 0:
            return False

        # Process took too long. Kill it.
        logger.warning('%s shutdown hook did not complete in %s secs. '
                            'Terminating it...', self.callback, self.timeout)
        # Terminate the process
        self._terminate()

        logger.info('Shutdown hook %s was terminated', self.callback)
        return True

class _ShutdownHookExecutor(object):
    '''Executor of shutdown hooks. The executor executes the shutdown hooks based
    on their priority, as the least priority hooks will be executed first.
    '''
    def  __init__(self, maxParallelCount=10, pollingTime=0.2):
        '''Creates shutdown hook executor.

        @param maxParallelCount: Maximum count of hooks which could be executed
          in parallel
        @type maxParallelCount: int

        @param pollingTime: Time when the executor would sleep until check if
          a process has completed
        @type pollingTime: int
        '''
        # List of shutdown hooks sorted by their priority
        self.hooks = []
        # Maximum shutdown hooks executed in parallel
        self.maxParallelCount = maxParallelCount
        self.pollingTime = pollingTime

    def _findShutdownHook(self, callback):
        '''Finds shutdown hook by given callback.

        @param callback: Shutdown callback handler
        @type callback: Function

        @return: Found shutdown hook in case of match and None otherwise
        @rtype: _ShutdownHook
        '''
        for hook in self.hooks:
            if hook.callback == callback:
                return hook
        return None

    def addShutdownHook(self, callback, context=None, timeout=60,
                        priority=ShutdownhookPriority.COMPONENT_HOOK):
        '''Adds custom shutdown hook. The callback will be invoked in case the
        shutdown signal from external source is received and the context will
        be passed to the callback. This callback will not be invoked if the
        program exits normally, when the last non-daemon thread exits or when
        the exit (equivalently, system.exit) method is invoked. If the callback
        has been already registered its properties would be updated.

        @param callback: Shutdown callback handler
        @type callback: Function having the follow definition:
          def _myShutdownHook(context):
              pass

        @param context: Shutdown context passed to the shutdown hook
        @type context: Any object

        @param timeout: Timeout in seconds is the window in which the shutdown
          hook must complete
        @type timeout: int

        @param priority: The parameter specifies the priority with which the hook
          has to be executed in case of shutdown. 0es priority hooks will be
          executed first, then 1es priority and so on. If two hooks have same
          priority the order of their execution is undefined.
        @type priority: int

        @raise ValueError: In case the callback is None
        @raise KeyError: If that callback has been already registered
        '''
        if callback is None:
            error = 'Shutdown callback is None'
            logger.error(error)
            raise ValueError(error)

        if self._findShutdownHook(callback):
            error = 'Shutdown hook %s has been already registered' % callback
            logger.error(error)
            raise KeyError(error)
        else:
            newHook = _ShutdownHook(callback, context, timeout, priority)
            idx = -1
            # Keep the hooks sorted by priority, and then by timeout
            for idx, hook in enumerate(self.hooks):
                hasLessPriority = newHook.priority < hook.priority
                needMoreTime = newHook.priority == hook.priority and \
                                        newHook.timeout > hook.timeout
                if hasLessPriority or needMoreTime:
                    self.hooks.insert(idx, newHook)
                    break
            else:
                self.hooks.insert(idx + 1, newHook)

            logger.debug('Shutdown hook %s is registered', callback)

    def removeShutdownHook(self, callback):
        '''Removes the already registered shutdown hook.

        @param callback: Already registered shutdown callback handler
        @type callback: Function

        @raise KeyError: In case the callback has not been registered or
          already unregistered.
        @raise ValueError: In case the callback is None
        '''
        if callback is None:
            error = 'Shutdown callback is None'
            logger.error(error)
            raise ValueError(error)

        oldHook = self._findShutdownHook(callback)
        if oldHook:
            self.hooks.remove(oldHook)
            logger.debug('Shutdown hook %s is unregistered', oldHook.callback)
        else:
            error = 'Shutdown hook %s is not registered' % callback
            logger.error(error)
            raise KeyError(error)

    def shutdown(self):
        '''Executes all shutdown hooks based on their priority. The function must
        be invoked when a shutdown signal is received.
        '''
        if not self.hooks:
            logger.info('No shutdown hooks are defined')
            return

        logger.info('Start executing custom shutdown hooks')
        samePriorityHooks = []
        priority = None
        for hook in self.hooks:
            if samePriorityHooks and priority != hook.priority:
                self._executeSamePriorityHooks(samePriorityHooks)
                samePriorityHooks = []

            priority = hook.priority
            samePriorityHooks.append(hook)

        # Execute last batch of shutdown hooks
        if samePriorityHooks:
            self._executeSamePriorityHooks(samePriorityHooks)

    def _executeSamePriorityHooks(self, hooks):
        '''Executes given hooks in parallel, but do not start more than maxParallelCount
        in parallel.

        @param hooks: List of shutdown hooks
        @type hooks: list

        @param maxParallelCount: Maximum count of hooks which could be executed
          in parallel
        @type maxParallelCount: int

        @param pollingTime: Time when the executor would sleep until check if
          a process has completed
        @type pollingTime: int
        '''

        ####################################################################XXXXXXXXXX
        # Prevent serializing of Logger in Win (due Lock object)
        # XXX: This is the safe, if all child processes will create own logger.
        # XXX: It's mandatory, because logging handlers usually are not serializable
        # XXX: (opened streams/sockets/locks/etc.)
        if os.name == "nt":
            logging.Logger.__getstate__ = lambda self: {}

        processes = []
        def _removeCompletedProcesses(processes):
            '''The method blocks until at least one process complete.
            '''
            hasCompleted = False
            while not hasCompleted:
                for p in processes:
                    if p.isCompleted():
                        hasCompleted = True
                        processes.remove(p)

                # Sleep for some time if noone has completed
                if not hasCompleted:
                    time.sleep(self.pollingTime)

        # Start a different processes for every shutdown hook
        for hook in hooks:
            p = multiprocessing.Process(target=hook.callback, args=[hook.context])
            logger.debug('Invoke shutdown hook %s', hook.callback)
            p.start()

            # Keep the processes sorted by hook timeout
            startedProcess = _ExecutedProcess(p, time.time(), hook.timeout)
            processes.append(startedProcess)

            if len(processes) == self.maxParallelCount:
                _removeCompletedProcesses(processes)

        # Wait until all started processes complete
        while processes:
            _removeCompletedProcesses(processes)

def hardShutdown(_signal, _frame):
    '''Called on shutdown request. The function execute every shutdown hook and
    then terminate the process, as throwing exception from shutdown(signal)
    thread context - which will force main thread to exit.

    @param _signal: Sent signal
    @type _signal: int

    @param _frame: Main thread stack frame
    '''
    logger.info('Shutdown request has been received')
    shutdownHookExecutor.shutdown()

    logger.info('Terminating the process..')
    raise InterruptedCallError(errno.EINTR)

def softShutdown(_signal, _frame):
    '''Called on shutdown request. The function execute every shutdown hook, but
    does not force main process to exit, thus the main process could exit normally.

    @param _signal: Sent signal
    @type _signal: int

    @param _frame: Main thread stack frame
    '''
    logger.info('Shutdown request has been received')
    shutdownHookExecutor.shutdown()

def listenForShutdownSignal(onShutdownCallback=hardShutdown):
    '''Registers shutdown callback which will be called on shutdown signal
      request.

    @param onShutdownCallback: Function which is called when the program is
      shut down. If omitted then default shutdown callback will be registered;
      for more information take a look at #hardShutdown.
    @type onShutdownCallback: Function with follow structure
      def onShutdownCallback(signal, frame):
    '''
    signal.signal(transport.LISTEN_SHUTDOWN_SIGNAL, onShutdownCallback)

def emitShutdownSignal(frame=None):
    ''' Emits shutdown signal. Hides cross-platform adaption
    (send signal to itself does not work on Win with proper handlers catching)

    @param frame: frame like in usual signal handler
    @type frame: FrameType
    '''
    from transport.local import killProcess
    if os.name == "nt":
        handler = signal.getsignal(transport.LISTEN_SHUTDOWN_SIGNAL)
        if isinstance(handler, collections.Callable):
            handler(transport.LISTEN_SHUTDOWN_SIGNAL, frame)
            return
        logger.info('No shutdown handler, so kill itself')
    killProcess(os.getpid(), transport.LISTEN_SHUTDOWN_SIGNAL)

shutdownHookExecutor = _ShutdownHookExecutor()