#!python
import argparse

from time import sleep
from datetime import datetime
import os
import logging
import sys
import signal
from typing import List
from pathlib import Path

# Exit if psutil is not installed
try:
    import psutil
except ImportError:
    print("Error: Install psutil to get statistics!", file=sys.stderr)
    exit(0)


class GracefulKiller:
    """
    Kills the process graefully when it gets a signal from the OS.
    https://stackoverflow.com/questions/18499497/how-to-process-sigterm-signal-gracefully
    """
    kill_now = False

    def __init__(self, old_prefix=None, new_prefix=None, filename=None):
        signal.signal(signal.SIGINT, self.exit_gracefully)
        signal.signal(signal.SIGTERM, self.exit_gracefully)
        self.old_prefix = old_prefix
        self.new_prefix = new_prefix
        self.filename = filename

    def exit_gracefully(self, *args):
        if self.filename != None:
            logging.debug(
                f"Moving {self.old_prefix}/{self.filename} to {self.new_prefix}/{self.filename}\n")

            Path(f"{self.old_prefix}/{self.filename}")\
                .rename(f"{self.new_prefix}/{self.filename}")

        logging.info("Killing Process\n\n")
        self.kill_now = True


def get_all_user_procs(username=None) -> List[int]:
    # Get all processes running on the system
    total_procs = psutil.pids()
    # Make a list to return the processes to return
    user_procs = []

    # If username is not given get it from the os
    if username is None:
        try:
            username = os.getlogin()
        except OSError:
            user = "runner"

    total_procs = [int(p) for p in total_procs]
    for p in total_procs:
        # In case the process stoped between
        try:
            user = psutil.Process(p).username()
            name = psutil.Process(p).name()
        except psutil.NoSuchProcess:
            continue

        # Add user processes to list to return
        if username == user and p != os.getpid() and name != 'slurm_script':
            user_procs.append(p)

    return user_procs


def get_iocounters(pData):
    if 'io_counters' in pData and pData['io_counters'] is not None:
        read_count = pData['io_counters'].read_count
        write_count = pData['io_counters'].write_count
        read_chars = pData['io_counters'].read_chars
        write_chars = pData['io_counters'].write_chars
    else:
        read_count = "nan"
        write_count = "nan"
        read_chars = "nan"
        write_chars = "nan"
    return read_count, write_count, read_chars, write_chars


def get_meminfo(pData):
    rss = pData['memory_info'].rss
    vms = pData['memory_info'].vms

    return rss, vms


def get_cputimes(pData):
    # pcputimes(user=0.05, system=0.02, children_user=0.0, children_system=0.0, iowait=0.0)
    if 'cpu_times' in pData and pData['cpu_times'] is not None:
        user = pData['cpu_times'].user
        system = pData['cpu_times'].system
        try:
            children_user = pData['cpu_times'].children_user
            children_system = pData['cpu_times'].children_system
        except:
            children_system = "nan"
            children_user = "nan"

        try:
            iowait = pData['cpu_times'].iowait
            cpu_num = pData['cpu_num']
        except:
            iowait = "nan"
            cpu_num = "nan"

        try:
            idle = pData['cpu_times'].idle
        except:
            idle = "nan"

    else:
        user = "nan"
        system = "nan"
        iowait = "nan"
        cpu_num = "nan"
        idle = "nan"
        children_system = "nan"
        children_user = "nan"

    return cpu_num, user, system, iowait, children_system, children_user, idle


def cmd_data(pData):
    """
    Gets data from the command line arguments and serilaizes it for csv
    """
    cmd = pData['cmdline']
    if len(cmd) == 0:
        return "nan"
    cmd = [c.replace(",", "|") for c in cmd]

    return "|".join(cmd)


def runner(
        path: str = ".", filename: str = "stats.csv",
        pole_rate: float = 0.1, username: str = "",
        write_header: bool = True,
        move: bool = False):
    """
    Runs while your executable is still running and logs info
    about running process to a csv file.

    Args:
        outfile (str, optional): output filename. Defaults to "stats.csv".
        poleRate (float, optional): Time to sleep before getting new data. Defaults to 0.1.
        username (str, optional): User name to look for when getting statistics.
    """

    sleep(1)

    if move:
        killer = GracefulKiller(old_prefix=f"{path}/running",
                                new_prefix=f"{path}/done",
                                filename=filename)

        Path(f"{path}/running").mkdir(exist_ok=True)
        Path(f"{path}/done").mkdir(exist_ok=True)
        outfile = f"{path}/running/{filename}"
    else:
        killer = GracefulKiller()
        Path(f"{path}").mkdir(exist_ok=True)
        outfile = f"{path}/{filename}"

    stats_file = open(outfile, "w")

    header = ["@timestamp", "pid", "name", "num_threads", "cpu_num",
              "cpu_user", "cpu_system", "cpu_iowait",
              "cpu_children_system", "cpu_children_user", "idle",
              "mem_rss", "mem_vms", 'memory_percent',
              "num_fds", "read_count", "write_count", "read_chars",
              "write_chars", "cmdline", "current_dir"]
    # Make formater based on number of metrics in header
    ftm_writer = ",".join(["{}" for _ in range(len(header))])
    ftm_writer = ftm_writer + "\n"

    if write_header:
        stats_file.write(ftm_writer.format(*header))
        stats_file.flush()

    # Keep pulling data from the process while it's running
    while not killer.kill_now:
        user_procs = get_all_user_procs(username=username)
        for proc_num in user_procs:
            try:
                proc = psutil.Process(proc_num)
                pData = proc.as_dict()

                # Add new line to the file with relevant data
                stats_file.write(ftm_writer.format(
                    datetime.now().strftime("%m-%d-%Y %H:%M:%S.%f"),
                    proc_num,
                    pData['name'],
                    pData['num_threads'],
                    *get_cputimes(pData),
                    *get_meminfo(pData),
                    pData['memory_percent'],
                    pData['num_fds'],
                    *get_iocounters(pData),
                    cmd_data(pData),
                    pData['cwd']
                ))

            except psutil.NoSuchProcess as e:
                logging.error(f'Error: {e}')
                continue
            except Exception as e:
                logging.error(f'Error: {e}')
                # Breaks out of just the loop and not the function

        # Write and Sleep for a number of seconds before going to the next loop
        stats_file.flush()
        sleep(pole_rate)

    # Finally we close the file.
    stats_file.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--tag", type=str,
                        help="Tags the process and gives to the statistcs csv file.",
                        default="stats")
    parser.add_argument("-o", "--outfile", type=str,
                        help="File name for csv.",
                        default="stats.csv")
    parser.add_argument("-p", "--path", type=str,
                        help="Path to put csv file.",
                        default=".")
    parser.add_argument("-d", "--debug", action="store_true",
                        help="Run with debugging info.",
                        default=False)
    parser.add_argument("-r", "--rate", type=float,
                        help="Polling rate for process.", default=0.1)
    parser.add_argument("-u", "--user", type=str,
                        help="Username to get stats for.", default=None)
    parser.add_argument("-noh", "--no-header",
                        help="Turn off writting the header.", default=True, action='store_false')
    parser.add_argument("-mv", "--move",
                        help="Moves file from 'running' to 'complete'", default=False, action='store_true')

    args = parser.parse_args()

    # Turn on logging if in debug mode
    if args.debug:
        logging.basicConfig(
            format='%(asctime)s %(levelname)s ==> %(message)s', level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.FATAL)

    logging.debug("Running in debug mode")

    # Start the recorder
    runner(path=args.path, filename=args.outfile, pole_rate=args.rate,
           username=args.user, write_header=args.no_header, move=args.move)
