#!/usr/bin/env python3
import sys
import os
import subprocess
import argparse
import time
import signal

class term_fmt:
    cyan = '\033[96m'
    green = '\033[92m'
    yellow = '\033[93m'
    bold = '\033[1m'
    reset = '\033[0m'

def get_from_command(command):
    ret = subprocess.check_output(command).decode(sys.stdout.encoding)
    return ret

def exit_handler(signum, frame):
    os.system('clear')
    quit()


# set up argument parser
parser = argparse.ArgumentParser()
parser.add_argument('-n', help='update output every nth second [2]', default=2, type=int)

# get time interval from command line
args = parser.parse_args()
interval = args.n

# get user name to higlight in output
user_name = os.getlogin()

# connect SIGINT signal to simple handler
signal.signal(signal.SIGINT, exit_handler)

while True:
    os.system('clear')
    print("Update interval: " + str(interval) + "s" + "\n")
    queue = get_from_command(['squeue'])
    output = ""
    for line in queue.splitlines():
        if user_name in line:
            if " R " in line:
                color = term_fmt.green + term_fmt.bold
            elif " CF " in line:
                color = term_fmt.yellow + term_fmt.bold
            else:
                color = term_fmt.cyan
            line = color + line + term_fmt.reset
        output += line + "\n"
    print(output)
    time.sleep(interval)