# Copyright (c) 2020 All Rights Reserved
# Author: William H. Guss, Brandon Houghton

import coloredlogs
import argparse
import psutil
import time
import shutil
import logging
import subprocess
import os
import sys
import signal
from daemoniker import daemonize

logger = logging.getLogger('process_watcher')
MINERL_WATCHERS_DIR = os.path.join(os.environ.get('TMPDIR', '.'), 'logs', 'minerl_watchers')

CHILD_DIR_ARG = 'child-dirs'


def parse_args():
    """Parses the arguments for the process watcher.
    """
    parser = argparse.ArgumentParser(
        description='A general process watcher utility that ensures '
                    'that a parent and child process terminate uniformly.')
    parser.add_argument('parent_pid', type=int,
                        help='The PID of the parent process.')
    parser.add_argument('child_pid', type=int,
                        help='The PID of the child process.')
    parser.add_argument('--{}'.format(CHILD_DIR_ARG),
                        type=str, help='Temporary directories which should be deleted '
                                       'were the processes to terminate.',
                        nargs='+')
    return parser.parse_args()


def launch(parent_pid, child_pid, *temp_dirs):
    """Launches the process watcher with a PID.

    Args:
        parent_pid (int): The parent PID.
        child_pid (int): The child PID.

    Returns:
        psutil.Process: The process object of the watcher.
    """
    logger.info("Launching process watcher daemonizer.")
    subprocess.check_call([
                              sys.executable, os.path.join(os.path.dirname(__file__), "process_watcher.py"),
                              str(parent_pid), str(child_pid),
                              '--{}'.format(CHILD_DIR_ARG)] + list(temp_dirs))
    logger.info("Process watcher daemonizer launched successfully.")


def reap_process_and_children(process, timeout=5):
    "Tries hard to terminate and ultimately kill all the children of this process."

    def on_process_wait_successful(proc):
        returncode = proc.returncode if hasattr(proc, 'returncode') else None
        logger.info("Process {} terminated with exit code {}".format(
            proc, returncode))

    def get_process_info(proc):
        return "{}:{}:{} i {}, owner {}".format(
            proc.pid,
            proc.name(),
            proc.exe(),
            proc.status(),
            proc.ppid()

        )

    procs = process.children(recursive=True)[::-1] + [process]
    try:
        logger.info("About to reap process tree of {}, ".format(get_process_info(process)) +
                    "printing process tree status in termination order:")
        for p in procs:
            logger.info("\t-{}".format(get_process_info(p)))
    except psutil.ZombieProcess:
        logger.info("Zombie process found in process tree.")

    for p in procs:
        try:
            logger.info("Trying to SIGTERM {}".format(
                get_process_info(p)
            ))
            p.terminate()
            try:
                p.wait(timeout=timeout)
                on_process_wait_successful(p)
            except psutil.TimeoutExpired:

                logger.info(
                    "Process {} survived SIGTERM; trying SIGKILL on {}".format(p.pid, get_process_info(p)))
                p.kill()

                try:
                    p.wait(timeout=timeout)
                    on_process_wait_successful(p)
                except psutil.TimeoutExpired:
                    # Give up
                    logger.info(
                        "Process {} survived SIGKILL; giving up (final status) {}, (owner) {}".format(p.pid, p.status(),
                                                                                                      p.ppid()))
        except (psutil.NoSuchProcess, psutil.ZombieProcess):
            logger.info("Process {} does not exist or is zombie.".format(p))


def main(args):
    """The primary process watcher program. 

    Args:
        args (--): Arguments from the process watcher.
    """

    # Wait for processes to be launched
    logger.info(
        "Process watcher started between parent {}".format(args.parent_pid) +
        " and child {} ".format(args.child_pid))
    time.sleep(1)
    try:
        child = psutil.Process(args.child_pid)
    except psutil.NoSuchProcess:
        child = None
    try:
        parent = psutil.Process(args.parent_pid)
    except psutil.NoSuchProcess:
        parent = None

    while True:
        try:
            # Sleep for a short time, and check if subprocesses needed to be killed.
            time.sleep(0.1)

            if not parent.is_running() or parent is None:
                logger.info(
                    "Parent is not running, hence we need terminate the child.")
                if child is not None:
                    reap_process_and_children(child)
                    try:
                        for temp_dir in args[CHILD_DIR_ARG]:
                            shutil.rmtree(temp_dir)
                    except OSError:
                        logger.warning(
                            "Failed to delete temporary child directory. It may have already been removed.")
                    except KeyError:
                        # No CHILD_DIR_ARG provided.
                        pass
                return
            # Kill the watcher if the child is no longer running.
            # If you want to attempt to restart the child on failure, this
            # would be the location to do so.
            if not child.is_running():
                logger.info(
                    "Child is not running anymore, launcher can terminate."
                )
                return
        except KeyboardInterrupt:
            pass


if __name__ == '__main__':
    args = parse_args()
    os.makedirs(MINERL_WATCHERS_DIR, exist_ok=True)

    os_cur_dir = os.path.abspath(os.getcwd())
    watcher_name = 'watcher_{}-{}'.format(
        args.parent_pid,
        args.child_pid
    )

    daemonize(os.path.join(MINERL_WATCHERS_DIR, watcher_name + '.pid'))

    coloredlogs.install(level=logging.DEBUG, stream=open(
        os.path.join(MINERL_WATCHERS_DIR, watcher_name + '.log'), 'w'
    ))

    main(args)
    exit()  # Correctly invoke the at_exit handlers and any other os clean up _explicitly_!
