#!/usr/bin/env python3
"""
This script runs Supervisorctl to issue commands to
remote FELIX nodes on which the Supervisor HTTP interface
is enabled. Authentication information is read from the
config files (that need to be accessible).

- The taget hosts are passed by option -t that supports
wildcarding e.g. -t pc-tdq-flx-nsw-*

- The action is passed with option -a
e.g. -a shutdown or -a status

- If the action is applied to a group/process use -p
e.g. multivisor.py -t pc-tdq-flx-nsw-* -a stop -p 50_daq:*


Author: carlo.gottardo@cern.ch
"""

import os
import logging as log
import argparse
import subprocess
from threading import Thread

log.basicConfig(format='[%(asctime)s] %(levelname)s: %(message)s', datefmt='%Y/%m/%d %H:%M:%S', level=log.INFO)


class multivisor():

    def __init__(self, sv_path, hosts, action, process, serial):
        self.conf_path = sv_path
        self.conf_list = []
        self.selected = {}
        if process != '':
            self.action = " ".join([action, process])
        else:
            self.action = action
        self.hostnames = hosts
        self.serial = serial
        self.parallel_out = []

    def get_host_list(self):
        """Get list of FELIX hosts from supervisor conf files"""
        for f in os.scandir(self.conf_path):
            if f.is_file() and f.name.startswith("supervisord"):
                if not ('flx' in f.name or 'felix' in f.name):
                    log.warning("Discarding config file: %s", f.name)
                    continue
                self.conf_list.append(f.name)
        if len(self.conf_list) == 0:
            log.error("No FELIX hosts found. Terminating.")
            exit(1)

    def select_hosts(self):
        """Select hosts to act on"""
        for idx, host in enumerate(self.hostnames):
            self.hostnames[idx] = host.replace('*', '')

        for c in self.conf_list:
            for n in self.hostnames:
                if n in c:
                    h = c.removeprefix('supervisord_').removesuffix('.conf')
                    if h not in self.selected.keys():
                        self.selected[h] = c

    def confirmation(self):
        """Get confirmation before proceeding"""
        log.info("You are requesting \"supervisorctl " + self.action + "\" on:")
        counter = 1
        for h in self.selected.keys():
            print(counter, ":", h)
            counter += 1
        log.info("Do you want to continue? [Yes/No]")
        check = input("")
        if check == "No":
            log.info("Aborting")
        elif check == "Yes":
            self.do_action()
        else:
            log.error("Invalid answer. Terminating.")

    def do_action(self):
        """Run supervisorctl command"""
        if self.serial is True:
            self.do_serial_action()
        else:
            self.do_parallel_action()

    def do_serial_action(self):
        for h in self.selected.keys():
            conf = os.path.join(self.conf_path, self.selected[h])
            iface = "http://" + h + ":9001"
            cmd = " ".join(["supervisorctl -c", conf, "-s", iface, self.action])
            try:
                stdoutdata = subprocess.check_output(cmd, shell=True, universal_newlines=True)
            except Exception as e:
                stdoutdata = str(e.output)
            log.info("Processing %s", h)
            log.info("Result:\n%s", stdoutdata)

    def worker(self, host, cmd):
        output_msg = "HOST: " + host + "\n-------------------------------\n"
        try:
            output_msg += subprocess.check_output(cmd, shell=True, universal_newlines=True)
        except Exception as e:
            output_msg += str(e.output)
        self.parallel_out.append(output_msg)

    def do_parallel_action(self):
        threads = []
        for h in self.selected.keys():
            conf = os.path.join(self.conf_path, self.selected[h])
            iface = "http://" + h + ":9001"
            cmd = " ".join(["supervisorctl -c", conf, "-s", iface, self.action])
            log.debug("Spawning thread for \"%s\"", cmd)
            threads.append(Thread(target=self.worker, args=(h, cmd)))
            threads[-1].start()
        for t in threads:
            t.join()
        log.info("All threads returned. Output below.")
        for e in self.parallel_out:
            print("==========================================")
            print(e)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(dest='action', type=str, default='', help='Action, e.g. start / stop / status...')
    parser.add_argument(dest='host', nargs='+', default=[], help='Target hostname(s)')
    parser.add_argument('-p', '--process', dest='proc', type=str, default='', help='Target application or group')
    parser.add_argument('-s', '--sdir', dest='svdir', type=str, default='/det/tdaq/felix-operation/supervisord/', help='Location of SV config files')
    parser.add_argument('-r', '--serial', dest='serial', action='store_true', default=False, help='Process hosts serially')
    args = parser.parse_args()

    mvisor = multivisor(args.svdir, args.host, args.action, args.proc, args.serial)
    mvisor.get_host_list()
    mvisor.select_hosts()
    mvisor.confirmation()
