#!/usr/bin/env python3

import sys
import os

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)

from datetime import datetime

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

import subprocess
import logging

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.
import signal
signal.signal(signal.SIGTERM, handler)
signal.signal(signal.SIGQUIT, handler)
signal.signal(signal.SIGINT, handler)

import time

if os.environ.get("_PBENCH_UNIT_TESTS"):
    _OC_DELAY = 0
    def mock_pause():
        time.sleep(1)
        global terminate
        terminate = True
    signal.pause = mock_pause
    def now():
        return datetime.utcfromtimestamp(0).strftime("%Y-%m-%dT%H:%M:%S.%f")
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

        # 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")
