import atexit
import getpass
import inspect
import os
import time
import datetime
import traceback
from threading import Thread, Lock

import logging

import sys

from torii import ToriiException
from .session import Session
from .process_server import start_process_server, stop_process_server
from .process import Process


class ProcessManager(object):


    def __init__(self, name=None, dir=None, logdir=None, server_host='localhost', server_port=5000, auto_cleanup=True, logger=None):
        """
        Process manager constructor

        :param name:   the name of the process manager
        :param dir:    the directory containing the traces, locks, logs...
        :param logdir: the directory containing the log files
        :param server_host:  the HTTP host for the REST server
        :param server_port:  the HTTP port for the REST server
        :param auto_cleanup: does the PM kill its children on exit
        :param logger: the logger
        """

        self.processes = []
        self.auto_cleanup = auto_cleanup
        self.name = name if name else 'process_manager'
        self.dir = dir if dir else os.getcwd()
        self.logdir = logdir if logdir else os.path.join(self.dir, 'logs')

        self.monitor_thread = None
        self.monitor_frequency = 10
        self.monitor_log_frequency = 60
        self.monitor_max_retry = 10

        self.rest_thread = None
        self.server_host = server_host
        self.server_port = server_port
        self.before_start_handler = None
        self.on_stop_handler = None
        self.on_start_handler = None
        self.lock = Lock()

        if logger:
            self.logger = logger
        else:
            # init the logger and a console handler
            self.logger = logging.getLogger(self.name)
            self.logger.setLevel(logging.DEBUG)
            ch = logging.StreamHandler(stream=sys.stdout)
            ch.setLevel(logging.DEBUG)
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            ch.setFormatter(formatter)
            self.logger.addHandler(ch)

        if not os.path.isdir(self.logdir):
            os.makedirs(self.logdir)

        # create file handler that logs in the logdir
        self.log_file = os.path.join(self.logdir, '{}_{}.log'.format(self.name, os.getpid()))
        fh = logging.FileHandler(self.log_file)
        fh.setLevel(logging.DEBUG)
        fh.setFormatter(self.logger.handlers[0].formatter)
        self.logger.addHandler(fh)

        self.active = True
        # cleanup on exit
        def _exit_handler():
            self.shutdown()

        # auto-cleanup on exit
        atexit.register(_exit_handler)

        self.logger.info('The process manager initialized')


    def new_process(self, command, name=None, start=True, check=None, check_timeout=None, **kwparam):
        """
        Instanciate a new child process

        :param command: the command line
        :param name: the logical name
        :param start: does the child starts aotomatically
        :param check: callback function to check the process has correctly started
        :param check_timeout: the timout for the check to append
        :param kwparam: extra parameters
        :return:
        """

        if not name:
            name = os.path.basename(command[0])

        self.logger.info('Create a new process \'{}\' (command = {})'.format(name, ' '.join(command)))

        # create the process
        process = Process(name=name, command=command, manager=self, **kwparam)

        # start the process
        if start:
            process.start(check=check, check_timeout=check_timeout)

        # add the process to the PM list
        self.processes.append(process)

        return process

    def set_before_start(self, handler, **kwparams):
        """
        Set the handler of the before-start event
        :param handler:
        :param kwparams:
        :return:
        """

        # add the process as argument if needed
        if 'manager' in inspect.getargspec(handler).args:
            kwparams['manager'] = self

        self.before_start_handler = (handler, kwparams)

    def set_on_start(self, handler, **kwparams):
        """
        Set the handler of the on-start event
        :param handler:
        :param kwparams:
        :return:
        """

        # check the handler is valid
        if 'process' not in inspect.getargspec(handler).args:
            raise ToriiException('Bad on_start handler, missing \'process\' argument')

        # add the manager as argument if needed
        if 'manager' in inspect.getargspec(handler).args:
            kwparams['manager'] = self

        self.on_start_handler = (handler, kwparams)

    def on_start(self, process):
        """
        Treat the on-start event for a given process
        :param process: the started process
        :return:
        """
        try:
            if self.on_start_handler:
                on_start, kwargs = self.on_start_handler
                on_start(process=process, **kwargs)
        except Exception as e:
            self.logger.exception('on_start event handling failed : {}'.format(e))
            self.shutdown()

    def set_on_stop(self, handler, **kwparams):
        """
        Set the handler of the on-stop event
        :param handler:
        :param kwparams:
        :return:
        """

        # check the handler is valid
        if 'process' not in inspect.getargspec(handler).args:
            raise ToriiException('Bad on_stop handler, missing \'process\' argument')

        # add the manager as argument if needed
        if 'manager' in inspect.getargspec(handler).args:
            kwparams['manager'] = self

        self.on_stop_handler = (handler, kwparams)

    def on_stop(self, process):
        """
        Treat the on-stop event for a given process
        :param process: the stopped process
        :return:
        """
        self.lock.acquire()
        try:
            if self.on_stop_handler:
                on_stop, kwargs = self.on_stop_handler
                on_stop(process=process, **kwargs)
            else:
                self.default_on_stop(process=process)
        except Exception as e:
            self.logger.exception('on_stop event handling failed: {}'.format(e))
            self.shutdown()
        finally:
            self.lock.release()

    def default_on_stop(manager, process):
        """
        Default behavior on process stop
        :param process:
        :return:
        """
        if process.mandatory:
            manager.shutdown(force_cleanup=True)

    def shutdown(self, force_cleanup=False):
        """
        Terminate all the processes of the manager
        :return:
        """
        # cleanup the child processes
        if self.auto_cleanup or force_cleanup:
            self.cleanup()

        stop_process_server()
        self.active = False

    def start_monitoring(self, blocking=False):
        """
        Start the monitoring of the processes stats
        :param blocking: block until the manager is shut down
        :return:
        """

        if blocking:
            self._do_monitor()
        else:
            self.monitor_thread = Thread(target=self._do_monitor)
            self.monitor_thread.daemon = True # thread dies with the program
            self.monitor_thread.start()

    def start_rest_server(self, blocking=True):
        """
        Start the REST server
        :param blocking: block until the manager is shut down
        :return:
        """
        if blocking:
            start_process_server(self, self.server_host, self.server_port)
        else:
            self.rest_thread = Thread(target=start_process_server, args=(self, self.server_host, self.server_port))
            self.rest_thread.daemon = True # thread dies with the program
            self.rest_thread.start()


    def _do_monitor(self):
        nb_errors = 0

        self.active = True
        self.logger.info('Start the process monitoring (nb_process = {})'.format(len(self.processes)))

        last = datetime.datetime.utcnow().timestamp()
        while self.active and nb_errors <= self.monitor_max_retry:
            try:
                now = datetime.datetime.utcnow().timestamp()
                total = len(self.processes)
                running = 0
                for p in self.processes:
                    if p.pid_exists():
                        running += 1

                for p in self.processes:
                    p.update_stats()

                # Log sometimes ...
                if self.monitor_log_frequency >= 0 and (now - last) > self.monitor_log_frequency:
                    last = now
                    details = ' - ' + '\n - '.join(p.get_abstract() for p in self.processes)
                    self.logger.info('Monitoring {0} running processes (total is {1})\n{2}'
                                      .format(running, total, details))

                # reset the error counter
                nb_errors = 0

            except Exception as e:
                nb_errors += 1
                trace = traceback.format_exc()
                self.logger.error('Failed to monitor processes ({}/{}) : {}\n{}'
                                  .format(nb_errors, self.monitor_max_retry, e, trace))


            try:
                # Time to sleep a bit
                time.sleep(self.monitor_frequency)
            except:
                pass

        self.logger.info('End of the process monitoring')

    @property
    def session(self):
        """
        :return: the session object corresponding to this manager
        """
        session = Session(dir=self.dir,
                          user=getpass.getuser(),
                          active=True,
                          pid=os.getpid(),
                          processe_pids={x.pid:x.name for x in self.processes if x.pid and x.active})

        return session

    def list(self):
        return [process.desc for process in self.processes]

    def cleanup(self, reverse=True):
        if self.processes:
            self.logger.info('Kill all the remaining processes')
            list_processes = self.processes
            if reverse:
                list_processes = reversed(self.processes)
            for process in list_processes:
                if process.active:
                    self.logger.info('... kill process \'{}\''.format(process.name))
                    try:
                        process.kill()
                    except:
                        pass

            self.processes = []

    def poll(self, name=None):
        return self.global_operation('poll', name=name, only_active=False)

    def send_signal(self, signal=0, name=None):
        self.global_operation('signal', name=name, signal=signal)

    def kill(self, name=None):
        self.global_operation('kill', name=name)

    def get_process(self, name):
        if name is None:
            return [x for x in self.processes]
        elif type(name) is list:
            return [x for x in self.processes if x.name in name]
        else:
            return next((x for x in self.processes if x.name == name), None)

    def global_operation(self, function, name=None, only_active=True,**kwparams):
        method = getattr(Process, function)

        if name is None:
            if only_active:
                processes = (x for x in self.processes if x.active)
            else:
                processes = self.processes
            return [method(x, **kwparams) for x in processes]
        elif type(name) is list:
            return [method(x, **kwparams) for x in self.processes if x.name in name]
        else:
            return next((method(x, **kwparams) for x in self.processes if x.name == name), None)


    def daemonize(self):
        """Detach a process from the controlling terminal and run it in the
        background as a daemon.
        """
        self.logger.info('Daemonizing the process manager (old_pid={})'.format(os.getpid()))

        try:
            pid = os.fork()
        except OSError as e:
            raise Exception("%s [%d]" % (e.strerror, e.errno))

        if (pid == 0):  # The first child.
            os.setsid()
            try:
                pid = os.fork()  # Fork a second child.
            except OSError as e:
                raise Exception("%s [%d]" % (e.strerror, e.errno))

            if (pid == 0):  # The second child.
                # os.chdir(self.logdir)
                # os.umask(0o22)

                # cleanup on exit
                def _exit_handler():
                    self.shutdown()

                # reattach the signals
                atexit.register(_exit_handler)

                self.logger.info('Process manager is now daemonized (new_pid={})'.format(os.getpid()))
            else:
                os._exit(0)  # Exit parent (the first child) of the second child.
        else:
            os._exit(0)  # Exit parent of the first child.

