#!/usr/libexec/platform-python

import sys
import os
import logging

PROG = os.path.basename(sys.argv[0])
try:
    tool_output_dir = sys.argv[1]
except IndexError:
    logging.error("[%s]: missing required tool output directory argument",
            PROG)
    sys.exit(1)
else:
    if not os.path.isdir(tool_output_dir):
        logging.error("[%s]: the provided tool output directory, '%s', does"
                " not appear to be a valid directory", PROG, tool_output_dir)
        sys.exit(1)

try:
    interval_param = sys.argv[2]
except IndexError:
    logging.error("[%s]: missing required interval argument", PROG)
    sys.exit(1)
else:
    try:
        interval = int(interval_param)
    except Exception:
        logging.error("[%s]: invalid interval argument", PROG)
        sys.exit(1)

# This is the location to which the inventory file parameter will be copied.
inv_file = os.path.join(tool_output_dir, "inv_hosts")

try:
    inv_file_param = sys.argv[3]
except IndexError:
    logging.error("[%s]: missing required host inventory file", PROG)
    sys.exit(1)
else:
    # First, we copy the inventory file provided to our tool output directory
    # for safe keeping.
    import shutil
    try:
        shutil.copyfile(inv_file_param, inv_file, follow_symlinks=True)
    except Exception as exc:
        logging.error("[%s]: unable to copy provided inventory file, '%s': %s",
                PROG, inv_file_param, exc)
        sys.exit(1)

import warnings
with warnings.catch_warnings():
    # Unfortunately the `ansible` module imported below will itself import the
    # `cryptography` module.  With Python 3.6 marked End-of-Life, the
    # `cryptography` module will emit a warning about deprecated support for
    # Python 3.6.  That output pollutes the test case causing it to fail.
    # Ignoring all warnings for this datalog is fine since the v0.69 branch is
    # also nearing End-of-Life.
    warnings.filterwarnings(action="ignore", category=Warning)
    try:
        from ansible.parsing.dataloader import DataLoader
        from ansible.inventory.manager import InventoryManager
    except ImportError as exc:
        logging.error("[%s]: required Ansible python3 modules not available, '%s'",
                PROG, exc)
        sys.exit(1)

try:
    data_loader = DataLoader()
    inventory = InventoryManager(loader = data_loader, sources = [inv_file])
except Exception as exc:
    logging.error("[%s]: invalid ansible inventory file, '%s': %s", PROG,
            inv_file, exc)
    sys.exit(1)

try:
    groups = inventory.get_groups_dict()
except Exception as exc:
    logging.error("[%s]: error processing groups in ansible inventory file,"
            " '%s': %s", PROG, inv_file, exc)
    sys.exit(1)

try:
    masters = groups['masters']
    nodes = groups['pprof']
except KeyError as exc:
    logging.error("[%s]: missing expected group in ansible inventory file,"
            " '%s': %s", PROG, inv_file, exc)
    sys.exit(1)

goroot = os.environ.get('GOROOT')
ut = os.environ.get('_PBENCH_UNIT_TESTS')
if goroot is not None and ut is None:
    go_cmd = os.path.join(goroot, "bin", "go")
else:
    # If GOROOT is not set, or we are running unit tests, assume we can find
    # go in the PATH.
    go_cmd = "go"


# The Operation class using subprocess internally.
import subprocess
import math
import errno
from time import time, sleep
from datetime import datetime


class Operation(object):
    """A simple class to keep track of a given operation to perform, and how to
    perform it.

    work_dir   - Working directory in which data will be collected & processed
    interval   - Interval at which to perform the operation
    server     - Server host name (FQDN)
    op_type    - One of "profile" or "heap"
    interface  - One of "normal", "api", or "controller"
    """
    _port_map = {
        "api": 8443,
        "controller": 8444,
        "normal": 10250
        }
    _suffix_map = {
        "api": "_api",
        "controller": "_controller",
        "normal": ""
        }
    op_types = ( 'profile', 'heap' )
    master_interfaces = ( 'api', 'controller' )

    def __init__(self, tool_output_dir, interval, server, op_type,
                interface='normal'):
        self.interval = interval
        self.server = server
        self.work_dir = os.path.join(tool_output_dir, server)
        if op_type not in self.op_types:
            raise Exception(f"Invalid operation type: '{op_type}'")
        self.op_type = op_type
        if interface not in self.master_interfaces and interface != 'normal':
            raise Exception(f"Invalid interface type: '{interface}'")
        self.interface = interface
        self.op_id = f"{server}:{interface}:{op_type}"
        try:
            os.mkdir(self.work_dir)
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise
        stdout_name = os.path.join(self.work_dir,
                f"{interface}_{op_type}.out")
        self.stdout_fp = open(stdout_name, "w")
        stderr_name = os.path.join(self.work_dir,
                f"{interface}_{op_type}.err")
        self.stderr_fp = open(stderr_name, "w")
        _server = f"https://{server}:{self._port_map[interface]}"
        if op_type == 'heap':
            self.oc_cmd = f"oc get --raw /debug/pprof/heap --server {_server}"
        else:
            self.oc_cmd = f"oc get --raw /debug/pprof/profile" \
                    f"?seconds={interval} --server {_server}"
        self.raw_out = os.path.join(self.work_dir, f"{op_type}_{interface}.raw")
        self.suffix = f"{self._suffix_map[self.interface]}"

    def run(self, timestamp):
        # Handles the sequence of "oc get" followed by the go tool invocation
        # required.

        ret_code = 0
        with open(self.raw_out, "w") as raw_ofp:
            oc_proc = subprocess.run(self.oc_cmd, shell=True,
                    stdout=raw_ofp,
                    stderr=subprocess.STDOUT)
        if oc_proc.returncode != 0:
            logging.warning("[%s] (%s) '%s' failed (%d)", PROG, self.op_id,
                    self.oc_cmd, oc_proc.returncode)
            ret_code = 1
        else:
            final_out = os.path.join(self.work_dir,
                    f"{timestamp}{self.suffix}_{self.op_type}.txt")
            global go_cmd
            pp_cmd = f"{go_cmd} tool pprof -text -seconds={self.interval}" \
                    f" -output {final_out} {self.raw_out}"
            pp_proc = subprocess.run(pp_cmd, shell=True,
                    stdout=self.stdout_fp,
                    stderr=self.stderr_fp)
            if pp_proc.returncode != 0:
                logging.warning("[%s] (%s) '%s' failed (%d)", PROG, self.op_id,
                        pp_cmd, pp_proc.returncode)
                ret_code = 1
        try:
            os.unlink(self.raw_out)
        except OSError as exc:
            if exc.errno != errno.ENOENT:
                logging.warning(
                        "[%s] (%s) unable to remove raw output file, '%s': %s",
                        PROG, self.op_id, self.raw_out, exc)
        return ret_code


def gen_operations(masters, nodes):
    # First yield all the combination of master, interfaces, and operation
    # types.
    for master in masters:
        for op_type in Operation.op_types:
            for interface in Operation.master_interfaces:
                try:
                    operation = Operation(tool_output_dir, interval, master,
                            op_type, interface)
                except Exception:
                    logging.exception(f"{master}")
                else:
                    yield operation
    # Yield all the combinations of node operations
    for node in nodes:
        for op_type in Operation.op_types:
            try:
                operation = Operation(tool_output_dir, interval, node, op_type)
            except Exception:
                logging.exception(f"{node}")
            else:
                yield operation

# We keep track of pids by operation "ID", as that should be unique.
parent = True
pids = {}
for operation in gen_operations(masters, nodes):
    if operation.op_id in pids:
        logging.warning("[%s]: encountered duplicate operation, %s, line %d",
                PROG, operation.op_id, linenum)
        continue
    try:
        pid = os.fork()
    except OSError:
        logging.exception("[%s]: os.fork() failed [operation=%s, line=%d]",
                PROG, operation.op_id, linenum)
        sys.exit(1)
    else:
        if pid == 0:
            # Child continue below
            parent = False
            break
        else:
            pids[operation.op_id] = (pid, operation)

if parent:
    # If we exited the loop above and parent is still True, we are the parent;
    # just wait for the terminate signal, and pass it along to our children.
    import signal

    try:
        terminate = False

        if os.environ.get('_PBENCH_UNIT_TESTS'):
            def mock_pause():
                for op_id, val in pids.items():
                    pid, _ = val
                    try:
                        os.waitpid(pid, 0)
                    except OSError:
                        pass
                global terminate
                terminate = True
            signal.pause = mock_pause
        else:
            def handler(signum, frame):
                global terminate
                terminate = True

            signal.signal(signal.SIGTERM, handler)
            signal.signal(signal.SIGQUIT, handler)
            signal.signal(signal.SIGINT, handler)

        while not terminate:
            signal.pause()
    finally:
        for op_id, val in pids.items():
            pid, operation = val
            try:
                os.kill(pid, signal.SIGTERM)
            except OSError as exc:
                if exc.errno != errno.ESRCH:
                    logging.warning(
                            "[%s]: os.kill(TERM, %d) [operation=%s] failed: %s",
                            PROG, pid, operation.op_id, exc)
    sys.exit(0)

# We are now in the "child" ONLY at this point.

_unit_test_iterations = os.environ.get('_PBENCH_UNIT_TESTS')
if _unit_test_iterations is not None:
    # Mock "time.time" and "time.sleep" functions so that we ensure we only
    # execute as many iterations ("sleep" cycles) as specified by the unit
    # test environment variable.
    _unit_test_iterations = int(_unit_test_iterations)
    _sleep_iterations = 0
    _ts = 1234567890
    def _time(*args, **kwargs):
        """Each call to mock'd _time() increments the internal clock by
           one second.
        """
        global _ts
        _ts += 1
        return _ts
    time = _time
    _orig_sleep = sleep
    def _sleep(length):
        """Each call to mock'd _sleep() increments the internal clock by
           the amount of time requested for the sleep, and checks to see
           if we have exceeded the number of times time.sleep() should be
           called before exiting the datalog. It calls the actual time.sleep
           for one second no matter how much time the actual sleep would
           have been.
        """
        global _ts
        global _sleep_iterations
        global _unit_test_iterations
        _sleep_iterations += 1
        if _sleep_iterations > _unit_test_iterations:
            sys.exit(0)
        global _orig_sleep
        _orig_sleep(1)
        _ts += int(length)
    sleep = _sleep

# Loop infinitely executing the given "operation" once every
# interval.

# Initial timestamp and next interval calcultion
time_stamp = time()
next_start = time_stamp + interval
intervals = 0
total_duration = 0
drifted = 0

while True:
    # Run the operation
    ts = datetime.utcfromtimestamp(time_stamp)
    operation.run(ts.strftime("%Y-%m-%dT%H:%M:%S.%f"))

    # Record how long the operation took for use when we run beyond the given
    # interval.
    endtime = time()

    intervals += 1
    duration = endtime - time_stamp
    total_duration += duration
    avg_duration = total_duration / intervals

    if endtime < next_start:
        # Sleep for the remainder of the interval
        sleep(next_start - endtime)
    elif endtime > next_start:
        drifted += 1
        logging.warning("[%s] interval exceeded now %d times: [interval: %d,"
                " start: %d, duration: %d]; please consider changing the"
                " interval to at least %d seconds",
                PROG, drifted, interval, time_stamp, duration,
                int(math.ceil(avg_duration * 1.5)))
        # The current time, endtime, is now the end of this interval.
        next_start = endtime

    # The next interval is calculated from the end of the previous
    # interval.
    time_stamp = next_start
    next_start = time_stamp + interval
