#!/usr/bin/env python
""" Bureaucrat - The Procfile & Deployfile process manager for Python Virtual Environments """

import argparse
import os
import re
import subprocess
import signal
import sys
import time


# Config
EXIT_FATAL = 128
EXIT_OK = 0
PROCESS_POLL_INTERVAL = 5


class ProcessLine(object):
    """ Represents a single process instance, eg a line of a Procfile """

    def __init__(self, name, cmd):
        self.name = name
        self.cmd = cmd


class Process(object):
    """ Represents a single process manager """

    def __init__(self, name, cmd):
        self.pl = ProcessLine(name, cmd)
        self.sub_process = None
        self.pid_file = None
        self.log_file = None
        self.ended = False

    @property
    def name(self):
        return self.pl.name

    @property
    def cmd(self):
        return self.pl.cmd

    @property
    def pid(self):
        if self.pid_file:
            try:
                # open pid file
                with open(self.pid_file, 'r') as f:
                    pid = f.read()
            except IOError:
                return None
            else:
                return int(pid)
        else:
            return self.sub_process.pid

    def kill(self):
        # kill process
        if hasattr(self.sub_process, 'returncode') and self.sub_process.returncode is not None:
            self.ended = True
            return
        else:
            print("Stopping %s: %s" % (self.name, self.cmd))
            try:
                os.kill(self.pid, signal.SIGTERM)
            except OSError:
                print("Proc %s (pid: %s) not found." % (self.name, self.pid))
            finally:
                if self.pid_file:
                    self.rm_pid()
                self.ended = True

    def rm_pid(self):
        # remove pid
        try:
            os.remove(self.pid_file)
        except OSError:
            print("Couldn't remove %s" % self.pid_file)
            sys.exit(EXIT_FATAL)

    def expanded_cmd(self):
        # Replace all environment vars in the cmd path and split
        return os.path.expandvars(self.pl.cmd).split()

    def execute(self, cwd, background=True):
        if self.pid_file and os.path.exists(self.pid_file):
            print("Error %s: %s exists" % (self.name, self.pid_file))
            return False
        else:
            print("%s: %s" % (self.name, self.cmd))
        cmd = self.expanded_cmd()
        try:
            if background:
                self.sub_process = subprocess.Popen(cmd, cwd=cwd, stdout=open(self.log_file, 'w'),
                                                    stderr=open(self.log_file, 'w'))
            else:
                self.sub_process = subprocess.Popen(cmd, cwd=cwd, stdout=open(self.log_file, 'w'),
                                                    stderr=subprocess.STDOUT)
                self.sub_process.wait()
        except OSError:
            print("Error: Command %s not found" % cmd[:1])
            return False
        else:
            if self.pid_file:
                # write pid
                with open(self.pid_file, 'w') as pid:
                    pid.write(str(self.sub_process.pid))
            return self.sub_process


class ProcessManager(object):
    """ Manage all processes in a given process file  """

    def __init__(self, process_file, env_file, log_path, pid_path, create_pids, named_procs=None, debug=False):
        self.debug = debug
        self.create_pids = create_pids

        # Init
        self._set_environment(env_file)
        self.processes = self._parse_process_file(process_file, log_path, pid_path)
        if named_procs is not None and len(named_procs) > 0:  # truncate processes
            self.process_manager.processes = [p for p in self.process_manager.processes if p.name in named_procs]

    def _set_environment(self, env_file):
        """ Set environment variables from .env file """
        reg = re.compile('(?P<name>\w+)(\=(?P<value>.+))')
        for line in open(env_file):
            m = reg.match(line)
            if m:
                name = m.group('name')
                value = ''
                if m.group('value'):
                    value = m.group('value')
                os.environ[name] = value
                if self.debug:
                    print("set %s %s" % (name, value))

    def _parse_process_file(self, process_file, log_path, pid_path):
        """ Parse Processfile (eg Procfile, Deployfile)
        :return: list of process objects
        """
        processes = []
        with open(process_file) as f:
            for line in f.readlines():
                match = re.search(r'([a-zA-Z0-9_-]+):(.*)', line)
                if not match:
                    raise Exception('Bad Process file line')
                name = match.group(1).strip()
                cmd = match.group(2).strip()
                p = Process(name, cmd)
                if pid_path is not None and self.create_pids:
                    p.pid_file = os.path.join(pid_path, '%s.pid' % name)  # Assign pid file
                else:
                    p.pid_file = None
                p.log_file = os.path.join(log_path, '%s.log' % name)  # assign log file
                processes.append(p)
        return processes


class Bureaucrat(object):
    """ Main Class """

    def __init__(self, env_file, virtual_env, app_path, log_path, pid_path, debug=False):
        self.debug = debug
        self.virtual_env = virtual_env
        self.app_path = app_path
        self.env_file = env_file
        self.log_path = log_path
        self.pid_path = pid_path
        # Set Path
        # Add virtualenv root and bin to path (alternatively we source bin/activate)
        os.environ['PATH'] = '%s/bin' % self.virtual_env + os.pathsep + \
            '%s' % self.app_path + os.pathsep + os.environ['PATH']
        # Listen for shutdown
        signal.signal(signal.SIGTERM, self._sigterm_handler)

    def load_procfile(self, process_file, named_procs, create_pids=True):
        self.process_manager = ProcessManager(process_file, self.env_file, self.log_path, self.pid_path, create_pids, named_procs)

    def start(self):
        """  Start processes defined in Procfile """
        for p in self.process_manager.processes:
            sys.stdout.write("Spawning ")
            p.execute(cwd=self.app_path, background=True)

    def stop(self):
        """  Terminate processes defined in Procfile """
        for p in self.process_manager.processes:
            p.kill()

    def deploy(self):
        """  Run all tasks in Deployfile """
        for p in self.process_manager.processes:
            sys.stdout.write("Running task ")
            p.execute(cwd=self.app_path, background=False)

    def _sigterm_handler(self, signal, frame):
        """ SIGTERM, eg kill handler """
        print('Got SIGTERM. Shutting down.')
        self.stop()
        sys.exit(EXIT_OK)

    @staticmethod
    def _check_running(processes):
        running = False
        for p in processes:
            if p.sub_process is False:
                print('Failed to spawn process: %s' % p.name)
                sys.exit(EXIT_FATAL)
            elif p.sub_process is not None and p.ended is False:
                p.sub_process.poll()  # update returncode
                if p.sub_process.returncode is None:
                    running = True
                elif p.sub_process.returncode is 0:
                    print("Spawned process ended: %s (pid: %s exit: %s)" %
                          (p.name, p.sub_process.pid, p.sub_process.returncode))
                    p.ended = True  # Stop monitoring
                    if p.pid_file:
                        p.rm_pid()
                elif p.sub_process.returncode is not 0:
                    print("Spawned process exited with error: %s (pid: %s exit: %s)" %
                          (p.name, p.sub_process.pid, p.sub_process.returncode))
                    sys.exit(EXIT_FATAL)  # bail out
        return running

    def monitor(self):
        """ Monitor processes defined in Procfile """
        try:
            while True:
                # Monitor all processes
                if self._check_running(self.process_manager.processes) is False:
                    print('All spawned processes have ended.')
                    return False  # Exit Ok
                time.sleep(PROCESS_POLL_INTERVAL)
        except KeyboardInterrupt:
            print('Shutting down.')
            self.stop()
            sys.exit(EXIT_OK)


def bureaucrat_args_init(args, process_file='Procfile', create_pids=True):
    # init
    virtual_env = args.venv or os.environ.get('VIRTUAL_ENV', os.getcwd())  # read env var from venv or param or cwd
    app_path = args.app or virtual_env
    env_file = args.envfile or os.path.join(app_path, '.env')

    if process_file == 'Procfile':
        process_file = args.procfile or os.path.join(app_path, 'Procfile')
    elif process_file == 'Deployfile':
        process_file = args.deployfile or os.path.join(app_path, 'Deployfile')

    if hasattr(args, 'process'):
        named_procs = set(args.process) or set(os.environ.get('PROCFILE_TASKS', '').split()) or None
    else:
        named_procs = set(os.environ.get('PROCFILE_TASKS', '').split()) or None

    if hasattr(args, 'pid_path'):
            pid_path = args.pidpath or ''
    else:
        pid_path = ''

    if hasattr(args, 'no_pid'):
        if args.no_pid:
            create_pids = False

    log_path = args.logpath or ''
    b = Bureaucrat(env_file, virtual_env, app_path, log_path, pid_path)
    b.load_procfile(process_file, named_procs, create_pids)
    return b


def stop(args):
    b = bureaucrat_args_init(args)
    return b.stop()


def start(args):
    b = bureaucrat_args_init(args)
    return b.start()


def restart(args):
    b = bureaucrat_args_init(args)
    b.stop()
    b.start()


def deploy(args):
    b = bureaucrat_args_init(args, process_file='Deployfile', create_pids=False)
    b.deploy()


def init(args):
    deploy(args)
    b = bureaucrat_args_init(args, create_pids=False)
    b.start()
    b.monitor()


if __name__ == "__main__":

    # Parse command line arguments
    parser = argparse.ArgumentParser(prog='Bureaucrat',
                         description='Bureaucrat - the Procfile & Deployfile manager for Python Virtual Environments')
    # Venv root
    venv = argparse.ArgumentParser(add_help=False)
    venv.add_argument('--venv', type=str, help='Virtualenv root')
    # Project / app root
    app = argparse.ArgumentParser(add_help=False)
    app.add_argument('--app', type=str, help='Application root')
    # Procfile
    procfile = argparse.ArgumentParser(add_help=False)
    procfile.add_argument('--procfile', type=str, help='Procfile path')
    # Deployfile
    deployfile = argparse.ArgumentParser(add_help=False)
    deployfile.add_argument('--deployfile', type=str, help='Deployfile path')
    # .env
    envfile = argparse.ArgumentParser(add_help=False)
    envfile.add_argument('--envfile', type=str, help='.env file path')
    # logpath
    logpath = argparse.ArgumentParser(add_help=False)
    logpath.add_argument('--logpath', type=str, help='log file path')
    # pidpath
    pidpath = argparse.ArgumentParser(add_help=False)
    pidpath.add_argument('--pidpath', type=str, help='pid file path')
    # don't create pid
    no_pid = argparse.ArgumentParser(add_help=False)
    no_pid.add_argument('--no-create-pid', action='store_true', help='Don\'t create PID files')
    # process name
    process_name = argparse.ArgumentParser(add_help=False)
    process_name.add_argument('process', nargs='*', type=str, help='Procfile Process Name')
    # Start
    sp = parser.add_subparsers()
    sp_start = sp.add_parser('start', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                             help='Starts Procfile processes')
    sp_start.set_defaults(func=start)
    # Stop
    sp_stop = sp.add_parser('stop', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                            help='Stops Procfile processes')
    sp_stop.set_defaults(func=stop)
    # Restart
    sp_restart = sp.add_parser('restart', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                               help='Restarts Procfile processes')
    sp_restart.set_defaults(func=restart)
    # Deploy
    sp_deploy = sp.add_parser('deploy', parents=[venv, app, deployfile, envfile, logpath],
                              help='Run tasks in Deployfile')
    sp_deploy.set_defaults(func=deploy)
    # Init
    sp_init = sp.add_parser('init', parents=[venv, app, deployfile, procfile, envfile, logpath, pidpath, no_pid, process_name],
                               help='Run Deployfile tasks and then start Procfile processes in foreground')
    sp_init.set_defaults(func=init)

    # Init parser if args passed
    if len(sys.argv) > 1:
        args = parser.parse_args()
        args.func(args)