#!python

# Monitoring system for submitted jobs via mcmsubmit
# The code is a recycled version of bucoffea monitoring system, bumon
# which is stored in this link:
# https://github.com/bu-cms/bucoffea/blob/master/bucoffea/execute/bumon

import curses
import logging
import os
import pickle
import sys
import time

from tabulate import tabulate
from mcmtest.lib.logging import LogManager

pjoin = os.path.join

logger = logging.getLogger('mcmlog')
format = '%(levelname)s (%(name)s) [%(asctime)s]: %(message)s'
date = '%F %H:%M:%S'
logging.basicConfig( level='DEBUG', format=format, datefmt=date, filename='monitor_log.txt', filemode='w')

settings = {
    'update_interval' : 20,
    'hide_finished' : False
}

def string_padding(string, length):
    return string + ' '  * (length - len(string))

def write_job_pad(pad, lm, hide=lambda x: False):
    colors = [3,3]
    table = []
    counter = 1
    for job in lm.jobs:
        if hide(job):
            continue
        if job.code == 0:
            colors.append(1)
        elif job.code != '-':
            colors.append(2)
        else:
            colors.append(3)

        table.append([counter, job.name, job.cluster, job.status, job.code, job.resubcount, time.strftime('%H:%M:%S', time.gmtime(job.runtime))])
        counter += 1

    tab = tabulate(sorted(table), headers=["","Name", "Cluster", "Status", "Return", "Resubmission count", "Runtime"])
    
    for i,l in enumerate(tab.split('\n')):
        pad.addstr(i,0,string_padding(l, curses.COLS), curses.color_pair(colors[i]))

def get_log_lines(maxlines):
    with open('monitor_log.txt', 'r') as f:
        lines = f.readlines()
        n = min([len(lines),maxlines-2])
        return lines[-n:]

def write_log_pad(logpad, logpadlen):
    logpad.addstr(0,0,'-'*curses.COLS)
    for i,l in enumerate(reversed(get_log_lines(logpadlen))):
        logpad.addstr(logpadlen-i-2,0,l)

def main(stdscr):
    curses.start_color()
    curses.use_default_colors()
    curses.init_pair(1, curses.COLOR_GREEN, -1)
    curses.init_pair(2, curses.COLOR_RED, -1)
    curses.init_pair(3, curses.COLOR_WHITE, -1)
    
    directory = sys.argv[1]
    if not os.path.exists(directory):
        raise RuntimeError(f'Directory not found: {directory}')

    # Initiate new pad
    lm = LogManager(directory)
    logging.info(f'Operating on directory {directory}.')
    logging.info(f'Found {len(lm.jobs)} jobs.')
    padlen = len(lm.jobs) + 5
    logpadlen = 10
    pad = curses.newpad(padlen,curses.COLS)
    logpad = curses.newpad(logpadlen,curses.COLS)
    stdscr.nodelay(True)
    stdscr.timeout(0)
    pad.nodelay(True)
    pad_pos = 0

    while True:
        lm.update()
        try:
            if settings['hide_finished']:
                hide = lambda x: x.code == 0
            else:
                hide = lambda x: False
            write_job_pad(pad, lm, hide=hide)
        
            # Scrolling
            start = time.time()
            stop = start
        
            while stop-start < settings['update_interval']:
                stop = time.time()

                pad.refresh(
                    pad_pos, 0,
                    1, 1,
                    curses.LINES-1 - logpadlen, curses.COLS-1
                    )
                logpad.refresh(
                    0, 0,
                    curses.LINES - logpadlen, 1,
                    curses.LINES-1, curses.COLS-1
                )
                try:
                    cmd = stdscr.getkey()
                    if cmd == 'q':
                        lm.save()
                        sys.exit(0)
                    elif cmd == 'r':
                        logger.info('Resubmit failed jobs.')
                        njobs = lm.resubmit_failed()
                        logger.info(f'Resubmitted {njobs} jobs.')
                    elif cmd == 'a':
                        lm.autoresub = not lm.autoresub
                        logger.info(f'Automatic resubmission set to {lm.autoresub}.')
                    elif cmd == 'h':
                        settings['hide_finished'] = not settings['hide_finished']
                        logger.info(f'Will {"" if settings["hide_finished"] else "not"} hide finished jobs.')
                        break
                    elif  cmd == 'KEY_HOME':
                        pad_pos = 0
                    elif  cmd == 'KEY_END':
                        pad_pos = padlen - curses.LINES + 1
                    elif  cmd == 'KEY_NPAGE':
                        pad_pos += 25
                    elif  cmd == 'KEY_PPAGE':
                        pad_pos -= 25
                    elif  cmd == 'KEY_DOWN':
                        pad_pos += 1
                    elif cmd == 'KEY_UP':
                        pad_pos -= 1
                except curses.error:
                    pass
                
                # Post process
                if pad_pos < 0:
                    pad_pos = 0

                if pad_pos >= padlen:
                    pad_pos = padlen-1

                write_log_pad(logpad,logpadlen)
                time.sleep(0.01)

        except KeyboardInterrupt:
            lm.save()
            break

if __name__ == "__main__":
    curses.wrapper(main)


