#!/usr/bin/python3

import logging
import os
import sys

PROG = os.path.basename(sys.argv[0])
PROG_dir = os.path.dirname(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 subprocess  # noqa E402

pmn_bin = os.path.join(PROG_dir, "pprof-master-n-nodes")
pmn_cmd = f"{pmn_bin} {inv_file}"
pmn_proc = subprocess.run(
    pmn_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
if pmn_proc.returncode != 0:
    logging.error("%r", pmn_proc)
    sys.exit(1)

json_str = pmn_proc.stdout.decode("UTF-8")
import json  # noqa E402

doc = json.loads(json_str)
masters = doc["masters"]
nodes = doc["nodes"]

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"


from datetime import datetime  # noqa E402
import errno  # noqa E402

# The Operation class using subprocess internally.
import math  # noqa E402
from time import sleep, time  # noqa E402


class Operation:
    """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",
            PROG,
            operation.op_id,
        )
        continue
    try:
        pid = os.fork()
    except OSError:
        logging.exception(
            "[%s]: os.fork() failed [operation=%s]",
            PROG,
            operation.op_id,
        )
        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)  # lgtm [py/unreachable-statement]

# 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

    _orig_time = time
    time = _time

    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)

    _orig_sleep = sleep
    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
