#!/usr/bin/python3.9

from datetime import datetime
import logging
import os
import signal
import subprocess
import sys
import time

PROG = os.path.basename(sys.argv[0])

try:
    tool_output_dir = sys.argv[1]
except IndexError:
    print(f"{PROG}: missing required tool output directory argument", file=sys.stderr)
    sys.exit(1)

components = []
for idx in range(2, len(sys.argv)):
    components.append(sys.argv[idx])
if not components:
    print(f"{PROG}: missing required components argument(s)", file=sys.stderr)
    sys.exit(1)


if os.environ.get("_PBENCH_UNIT_TESTS"):

    def _now():
        return datetime.utcfromtimestamp(0)

else:

    def _now():
        return datetime.utcnow()


def now():
    return _now().strftime("%Y-%m-%dT%H:%M:%S.%f")


def gather_nodes_ev(when):
    with open(os.path.join(tool_output_dir, f"nodes-{when}.txt"), "w") as nfp:
        nfp.write(f"timestamp: {now()}\n\n")
        nfp.flush()
        proc = subprocess.run(
            "oc get nodes --show-labels -o wide",
            stdout=nfp,
            stderr=subprocess.STDOUT,
            shell=True,
        )
        nfp.write(f"\ntimestamp: {now()}\n")
        nfp.flush()

    if proc.returncode != 0:
        logging.warning('%s: "oc get nodes" failed with %d', PROG, proc.returncode)
    with open(os.path.join(tool_output_dir, f"ev-{when}.txt"), "w") as efp:
        efp.write(f"timestamp: {now()}\n\n")
        efp.flush()
        proc = subprocess.run(
            "oc get ev --all-namespaces -o wide",
            stdout=efp,
            stderr=subprocess.STDOUT,
            shell=True,
        )
        efp.write(f"\ntimestamp: {now()}\n")
        efp.flush()
    if proc.returncode != 0:
        logging.warning('%s: "oc get ev" failed with %d', PROG, proc.returncode)


class GracefulPopen(subprocess.Popen):
    """Simple sub-class of Popen which adds a graceful way to terminate a
    process.
    """

    def silent_kill(self):
        try:
            self.kill()
        except Exception:
            # Don't bother reporting any errors on the .kill().
            pass

    def gracefully_terminate(self, name):
        try:
            self.terminate()
        except Exception:
            logging.warning(
                "%s: error gracefully terminating %s (pid %d)", PROG, name, self.pid
            )
            self.silent_kill()
        else:
            try:
                self.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logging.warning(
                    "%s: timeout gracefully terminating %s (pid %d)",
                    PROG,
                    name,
                    self.pid,
                )
                self.silent_kill()


terminate = False


def handler(signum, frame):
    global terminate
    terminate = True


# Establish signal handler for TERM, QUIT, and INT to stop process creation,
# and then tear them all down.

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

if os.environ.get("_PBENCH_UNIT_TESTS"):
    _OC_DELAY = 0

    def mock_pause():
        time.sleep(1)
        global terminate
        terminate = True

    signal.pause = mock_pause
else:
    _OC_DELAY = 5

# Gather "nodes" and "events" before we start other watchers.
gather_nodes_ev("start")

cmd_fmt = "oc get {all_ns_opt}{component} -o wide -w"
opts = {}
pids = {}

try:
    for component in components:
        if terminate:
            # During the creation of sub-processes we were told to terminate.
            break  # lgtm [py/unreachable-statement]

        # Wait 5 seconds between starting watchers.
        time.sleep(_OC_DELAY)

        opts["all_ns_opt"] = "" if component in ("cs", "pv") else "--all-namespaces "
        opts["component"] = component
        cmd = cmd_fmt.format(**opts)
        cfp = open(os.path.join(tool_output_dir, f"{component}.txt"), "w")
        cfp.write(f"timestamp: {now()}\n\n")
        cfp.flush()
        oc_cmd = GracefulPopen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True
        )
        log_ts = GracefulPopen(
            "pbench-log-timestamp",
            stdin=oc_cmd.stdout,
            stdout=cfp,
            stderr=subprocess.STDOUT,
            shell=True,
        )
        pids[component] = (oc_cmd, log_ts, cfp)

    # Wait for all sub-processes to complete
    while not terminate:
        signal.pause()
finally:
    for component, procs in pids.items():
        oc_cmd, log_ts, fp = procs
        oc_cmd.gracefully_terminate(component)
        log_ts.gracefully_terminate(f"{component}-ts")
        fp.close()
    # Gather "nodes" and "events" now that all watchers are stopped.
    gather_nodes_ev("stop")
