#!/usr/bin/env python3

import argparse;
import sys;
try:
    import jh_hpc_interface;
except ImportError:
    sys.exit('ImportError: Did you installed the jh_hpc_interface package? Try installing it with: python3 -m pip install jh_hpc_interface');

def config ():
    import os;

    lib_dir = os.path.dirname(jh_hpc_interface.jh_interface.__file__) + '/';
    if os.path.isfile(lib_dir + 'config/jh_config.ini'):
        print('Config path: ' +str(lib_dir) + 'config/jh_config.ini');
    else:
        sys.exit(f'Could not find config file {lib_dir}config/jh_config.ini???');

def start (stdin):
    # create object
    jupyter_job = jh_hpc_interface.jh_interface.manageInstance(stdin); 
    jupyter_job.start();

def get (jobid):
    # get object information
    import subprocess;
    import configparser;
    import os;
    import re;

    lib_dir = os.path.dirname(jh_hpc_interface.jh_interface.__file__) + '/';
    config = configparser.ConfigParser();
    if os.path.isfile(lib_dir + 'config/jh_config.ini'):
        config.read(lib_dir + 'config/jh_config.ini');
    else:
        sys.exit(f'Could not find config file {lib_dir}config/jh_config.ini');

    get_execnode_cmd = str(config['workload_manager']['cmd_job_get_exec_node']);
    get_execnode_cmd = get_execnode_cmd.replace('$JOBID', str(jobid));
    execnode_regex = str(config['workload_manager']['execnode_regex']);

    get_jobstate_cmd = str(config['workload_manager']['cmd_job_state']);
    get_jobstate_cmd = get_jobstate_cmd.replace('$JOBID', str(jobid));

    job_state = subprocess.check_output(get_jobstate_cmd, shell=True);
    get_node  = subprocess.check_output(get_execnode_cmd, shell=True);

    job_state = job_state.decode('utf-8');
    get_node = get_node.decode('utf-8');

    job_state = job_state.rstrip();
    get_node = get_node.rstrip();
    get_node = ' '.join(re.findall(str(execnode_regex), get_node));

    print(f'{job_state} {get_node}');

def stop (jobid):
    # stop started job
    import subprocess;
    import configparser;
    import os;

    lib_dir = os.path.dirname(jh_hpc_interface.jh_interface.__file__) + '/';
    config = configparser.ConfigParser();
    if os.path.isfile(lib_dir + 'config/jh_config.ini'):
        config.read(lib_dir + 'config/jh_config.ini');
    else:
        sys.exit(f'Could not find config file {lib_dir}config/jh_config.ini');

    get_stop_cmd = str(config['workload_manager']['cmd_kill_job']);
    get_stop_cmd = get_stop_cmd.replace('$JOBID', str(jobid));

    killjob_cmd = get_stop_cmd.rstrip();
    killjob = subprocess.check_output(killjob_cmd, shell=True);

def showlog (live=False):

    import configparser;
    import os;
    import subprocess;

    lib_dir = os.path.dirname(jh_hpc_interface.jh_interface.__file__) + '/';
    config = configparser.ConfigParser(interpolation=configparser.ExtendedInterpolation());
    if os.path.isfile(lib_dir + 'config/jh_config.ini'):
        config.read(lib_dir + 'config/jh_config.ini');
    else:
        sys.exit(f'Could not find config file {lib_dir}config/jh_config.ini');

    logfile = config['general']['log_file'];

    if not os.path.isfile(logfile):
        configf = lib_dir + 'config/jh_config.ini'
        sys.exit(f'Error: Log file {logfile} specified in {configf} does not exists!')

    if live:
        lf = subprocess.Popen(['tail','-f',str(logfile)], stdout=subprocess.PIPE);
        while True:
            line = lf.stdout.readline();
            sys.stdout.write(str(line));
    else:
        lf = subprocess.check_output(f'tail -n 10 {logfile}', shell=True).rstrip().decode('utf-8');
        print(lf);

def main ():
    
    parser = argparse.ArgumentParser();
    subparser = parser.add_subparsers(dest='operation');

    stop_operation = subparser.add_parser('getconfig', help='Output the path of the configuration file');
    
    showlog_operation = subparser.add_parser('showlog', help='Show last log');
    showlog_operation.add_argument('--live', action='store_true', help='Output last logs in live mode');

    get_operation = subparser.add_parser('get');
    get_operation.add_argument('--jobid', required=True, type=int, help='Outputs the execution host and the job state');

    stop_operation = subparser.add_parser('stop');
    stop_operation.add_argument('--jobid', required=True, type=int, help='Kill a job using the command specified in the configuration file jh_config.ini');

    args = parser.parse_args();

    if args.operation == 'getconfig':
        config();
    elif args.operation == 'showlog':
        if args.live:
            showlog(live=True);
        else:
            showlog();
    elif args.operation == 'get':
        get(args.jobid);
    elif args.operation == 'stop':
        stop(args.jobid);
    else:
        # default is: start();
        start(str(sys.stdin.read()).rstrip());

if __name__ == '__main__':
    main();
