#!/usr/bin/env python3

import logging
import os
import sys

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)

try:
    clear_counters_param = sys.argv[3]
except IndexError:
    logging.error("[%s]: missing required clear counters argument", PROG)
    sys.exit(1)
else:
    if clear_counters_param in ("true", "1", "T", "Y", "Yes", "yes", "True"):
        clear_counters = True
    else:
        clear_counters = False

import errno  # noqa E402
import subprocess  # noqa E402
from time import sleep, time  # noqa E402

# FIXME: haproxy router namespace, "default"
_router_namespace = "default"
haproxy_config = "/var/lib/haproxy/conf/haproxy.config"
haproxy_socket_path = "/var/lib/haproxy/run/haproxy.sock"
clear_counters_cmd = (
    f"echo 'clear counters all' | socat - UNIX-CONNECT:{haproxy_socket_path}"
)
show_stats_cmd = f"echo 'show stat' | socat - UNIX-CONNECT:{haproxy_socket_path}"


def fetch_routers():
    """fetch_routers - Generator to fetch the current list of routers."""
    oc_proc = subprocess.run(
        f"oc get pod --no-headers -n {_router_namespace}",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    if oc_proc.returncode != 0:
        logging.error(
            "[%s] failed to fetch list of router pods, stderr: '%r'",
            PROG,
            oc_proc.stderr,
        )
        sys.exit(1)
    try:
        lines = oc_proc.stdout.split("\n")
        for line in lines:
            if line.startswith("router-"):
                yield line.split(" ", 1)[0]
    except Exception:
        logging.error("[%s] failed to parse list of routers")
        sys.exit(1)


class Router:
    """An encapsulation of the data and methods needed to work with collecting
    metrics, logs, version, and configuration data from OCP haproxy router
    pods.

    A Router() object keeps track of four background processes:

        a. A process collecting the version information ("haproxy -vv" from
           within a pod)

        b. A process collecting the configuration information of an pod

        c. A process collecting the pod's logs

        d. A process collecting the pod's statistics ("show stat" or "clear
           counters all" issued over the haproxy UNIX socket)

    Data collected for  metrics, logs, configuration, and version information
    are stored in a directory per router:

        {tool_output_dir}/{router}/
            metrics.csv[.err]
        {tool_output_dir}/{router}/( "start" | "stop" )/
            haproxy-vv.txt[.err]
            haproxy.config[.err]
            pod.log[.err]

    In addition, we also maintain a "stderr" file for each operation (.err
    suffix added to the files of the same name) to capture any output emitted
    via stderr from various operations performed on a router.
    """

    def __init__(self, pod_namespace, pod_name, clear_counters):
        self._pod_namespace = pod_namespace
        self._pod_name = pod_name
        self._clear_counters = clear_counters
        self._dir = os.path.join(tool_output_dir, pod_name)
        try:
            os.mkdir(self._dir)
        except OSError as exc:
            if exc.errno == errno.EEXIST:
                logging.error("")
            else:
                raise Exception("")
        # The files used to capture stdout and stderr for metrics collection
        # are kept open so that we don't have to re-open the files each time
        # run a sub-process to capture data.
        cvs_name = os.path.join(self._dir, "metrics.csv")
        self.csv_fp = open(cvs_name, "w+")
        err_name = os.path.join(self._dir, "metrics.csv.err")
        self.err_fp = open(err_name, "w+")
        # The oc exec command prefix is always the same for this router pod.
        self._oc_exec_prefix = (
            f"oc exec {pod_name} -n {self._pod_namespace} -- /bin/sh -c"
        )
        # The following dict() maintains the active background processes being
        # run to collect data.
        self.proc = dict()

    def collect_metrics(self, timestamp):
        """Starts a background sub-process to clearing the counters."""
        print(timestamp, file=self.csv_fp)
        if self._clear_counters:
            # We only clear the counters once.
            self._clear_counters = False
            oc_exec_cmd = f"{self._oc_exec_prefix} {clear_counters_cmd}"
        else:
            oc_exec_cmd = f"{self._oc_exec_prefix} {show_stats_cmd}"
        proc = subprocess.Popen(
            oc_exec_cmd, shell=True, stdout=self.csv_fp, stderr=self.err_fp
        )
        self.proc["metrics"] = (proc, None, None)

    def _mktagdir(self, tag):
        tag_dir = os.path.join(self._dir, tag)
        try:
            os.mkdir(tag_dir)
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise
        return tag_dir

    def collect_version(self, timestamp, tag):
        tag_dir = self._mktagdir(tag)
        stdout_name = os.path.join(tag_dir, "haproxy-vv.txt")
        stdout_fp = open(stdout_name, "w")
        stderr_name = f"{stdout_name}.err"
        stderr_fp = open(stderr_name, "w")
        oc_exec_cmd = f"{self._oc_exec_prefix} haproxy -vv"
        proc = subprocess.Popen(
            oc_exec_cmd, shell=True, stdout=stdout_fp, stderr=stderr_fp
        )
        self.proc["ver"] = (proc, stdout_fp, stderr_fp)

    def collect_config(self, timestamp, tag):
        tag_dir = self._mktagdir(tag)
        stdout_name = os.path.join(tag_dir, "haproxy.config")
        stdout_fp = open(stdout_name, "w")
        stderr_name = f"{stdout_name}.err"
        stderr_fp = open(stderr_name, "w")
        oc_exec_cmd = f"{self._oc_exec_prefix} cat {haproxy_config}"
        proc = subprocess.Popen(
            oc_exec_cmd, shell=True, stdout=stdout_fp, stderr=stderr_fp
        )
        self.proc["config"] = (proc, stdout_fp, stderr_fp)

    def collect_logs(self, timestamp, tag):
        tag_dir = self._mktagdir(tag)
        stdout_name = os.path.join(tag_dir, "pod.log")
        stdout_fp = open(stdout_name, "w")
        stderr_name = f"{stdout_name}.err"
        stderr_fp = open(stderr_name, "w")
        oc_exec_cmd = f"oc logs {self._pod_name} -n {self._pod_namespace}"
        proc = subprocess.Popen(
            oc_exec_cmd, shell=True, stdout=stdout_fp, stderr=stderr_fp
        )
        self.proc["logs"] = (proc, stdout_fp, stderr_fp)

    def wait(self, all=True):
        # We wait on any "metrics" background process first. Once that
        # completes, we "reap" (poll()) any other background process so that
        # we don't skew the interval.
        errs = 0
        for name, proc in self.proc.items():
            proc, stdout_fp, stderr_fp = proc
            if all or name == "metrics":
                # We only actively wait() for metrics processes ...
                returncode = proc.wait()
                if returncode != 0:
                    errs += 1
                    logging.warning(
                        "[%s] reaped %s process for %s with error return code:" " %d",
                        PROG,
                        name,
                        self._pod_name,
                        returncode,
                    )
            else:
                # ... all other processes we just poll()
                returncode = proc.poll()
                if returncode is not None:
                    if returncode != 0:
                        errs += 1
                        logging.warning(
                            "[%s] reaped %s process for %s with error return code:"
                            " %d",
                            PROG,
                            name,
                            self._pod_name,
                            returncode,
                        )
                    # Close the provided stdout/err files regardless of
                    # success or failure.
                    stdout_fp.close()
                    stderr_fp.close()
        return True if errs == 0 else False

    def shutdown(self):
        for name, proc in self.proc.items():
            proc, stdout_fp, stderr_fp = proc
            proc.terminate()
            if stdout_fp is not None:
                stdout_fp.close()
            if stderr_fp is not None:
                stderr_fp.close()
            proc.wait()


if clear_counters:
    logging.info("[%s] clearing counters on all routers before collecting metrics.")

# The current list of active routers indexed by router pod name.
_routers = {}


def operation(routers, ts_str):
    # The set of current router pod names.
    curr_routers = set(pod_name for pod_name in fetch_routers())
    prev_routers = set(pod_name for pod_name in routers.keys())
    # These are all the new routers that have showed up since the last time.
    new_routers = curr_routers - prev_routers
    # These are all the routers that have died since the last time.
    dead_routers = prev_routers - curr_routers

    for pod_name in dead_routers:
        # Clear out any dead routers
        del routers[pod_name]

    for pod_name, router in routers.items():
        # Initiate metrics gathering for existing routers
        router.collect_metrics(ts_str)

    for pod_name in new_routers:
        # For new routers we first create a Router() object from an initial
        # list of router pods, optionally clearing the counters for each pod,
        # and initiating the version, config, and logs collection.
        routers[pod_name] = r = Router(_router_namespace, pod_name, clear_counters)
        r.collect_metrics(ts_str)
        r.collect_version(ts_str, "start")
        r.collect_config(ts_str, "start")
        r.collect_logs(ts_str, "start")

    # At this point, all the background processes are started.

    # Wait on all metrics collection (which also reaps any other background
    # processes as they finish).
    for router in routers.values():
        router.wait()


# Flag for terminating the operations loop below.
terminate = False

_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:
            global terminate
            terminate = True
        global _orig_sleep
        _orig_sleep(1)
        _ts += int(length)

    _orig_sleep = sleep
    sleep = _sleep
else:
    import signal  # noqa E402

    # We only setup a signal handler for the non-unit test case.
    def handler(signum, frame):
        global terminate
        terminate = True

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

from datetime import datetime  # noqa E402
import math  # noqa E402

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

while not terminate:
    operation(_routers, datetime.utcfromtimestamp(time_stamp).strftime("%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

# First wait on any existing sub-processes to complete.
for router in _routers.values():  # lgtm [py/unreachable-statement]
    router.wait(all=True)

ts_str = datetime.utcfromtimestamp(time()).strftime("%s.%f")
# Collect the final version, config, and logs for each existing pod.
for router in _routers.values():
    router.collect_version(ts_str, "stop")
    router.collect_config(ts_str, "stop")
    router.collect_logs(ts_str, "stop")

# Wait for all routers to finish collecting their required information.
for router in _routers.values():
    router.wait(all=True)
