#!/usr/bin/env python3

"""
Manage (create, list, modify and delete) and starting jupyter slurm kernels using srun
"""

import argparse;
import json;
import sys;
import os;
import getpass;
import re;
import tempfile;
import subprocess
from venv import create;
import slurm_jupyter_kernel;
from slurm_jupyter_kernel import script_template;
from pathlib import Path;
from shutil import copy, rmtree 
from pexpect import pxssh;
try:
    from jupyter_client import kernelspec;
except ImportError:
    from IPython.kernel import kernelspec;

class Color:
    # Foreground
    F_Default = "\x1b[39m"
    F_Black = "\x1b[30m"
    F_Red = "\x1b[31m"
    F_Green = "\x1b[32m"
    F_Yellow = "\x1b[33m"
    F_Blue = "\x1b[34m"
    F_Magenta = "\x1b[35m"
    F_Cyan = "\x1b[36m"
    F_LightGray = "\x1b[37m"
    F_DarkGray = "\x1b[90m"
    F_LightRed = "\x1b[91m"
    F_LightGreen = "\x1b[92m"
    F_LightYellow = "\x1b[93m"
    F_LightBlue = "\x1b[94m"
    F_LightMagenta = "\x1b[95m"
    F_LightCyan = "\x1b[96m"
    F_White = "\x1b[97m"

def add_slurm_kernel (kernel_displayname, kernel_language, kernel_cmd, slurm_parameter, loginnode, user, proxyjump, srun_cmd="srun", kernel_environment=None):
   
    print(f"\n\033[94m\u2687\033[0m Try to create slurm kernel '{kernel_displayname}'... \033[0m");
    args = [str(sys.executable), '-m', 'slurm_jupyter_kernel'];

    args.extend(['--slurm-parameter', slurm_parameter]);
    args.extend(['--loginnode', user+'@'+loginnode]);
    args.extend(['--proxyjump', proxyjump]);
    if srun_cmd is None:
        srun_cmd = 'srun';
    args.extend(['--srun-cmd', srun_cmd]);

    kernel_cmd = kernel_cmd.format(connection_file='{remote_connection_file}');
    args.extend(['--kernel-cmd', kernel_cmd]);
    args.extend(['--connection-file', '{connection_file}']);

    if not re.search('slurm', str(kernel_displayname), re.IGNORECASE):
        kernel_displayname = "Slurm " + str(kernel_displayname);

    kerneldir_name = kernel_displayname.replace(' ', '_');
    username = getpass.getuser();

    kernel_env = {}
    if not kernel_environment is None:
        kernel_env = dict( (key.strip(), val.strip()) for key, val in (item.split('=') for item in kernel_environment.split(',')) ); 
        args.extend(['--environment', str(kernel_env)]);

    json_kernel = {
        'display_name': kernel_displayname,
        'env': kernel_env,
        'metadata': {"remote_slurm_kernel": True},
        'argv': args
    };
    if not kernel_language is None:
        json_kernel['language'] = kernel_language;

    # create a temporary jupyter kernel directory
    tempdir = tempfile.mkdtemp();

    kernel_loc = os.path.join(str(tempdir), 'kernel.json');
    with open(os.path.join(tempdir, 'kernel.json'), 'w') as kfile:
        json.dump(json_kernel, kfile, indent=2, sort_keys=True);

    img_dir = str(Path(slurm_jupyter_kernel.__file__).parent.parent) + '/imgs/';
    if os.path.isdir(str(img_dir)):
        copy(str(img_dir) + 'logo-32x32.png', str(tempdir));
        copy(str(img_dir) + 'logo-64x64.png', str(tempdir));

    kernel_path = kernelspec.install_kernel_spec(tempdir, kerneldir_name, user=username);

    print(f'{Color.F_LightGreen}\u2714\033[0m Successfully created kernel: {kernel_path}{Color.F_Default}');

def delete_slurm_kernel ():
    
    selected_kernel = _select_slurm_kernel();

    while True:
        print(f'Following kernel selected: \033[1m\033[93m{selected_kernel}\033[0m');
        delete_yes_no = input('Are you sure you want to delete the selected kernel? [Y,n] ');
        if delete_yes_no == '' or delete_yes_no.upper() == 'Y':
            rmtree(str(selected_kernel));
            print(f'Kernel \033[92m{selected_kernel}\033[0m deleted');
            sys.exit();
        elif delete_yes_no.upper() == 'N':
            print('\033[94mI did not delete the selected kernel\033[0m');
            sys.exit();
        else:
            print('\033[91mInvalid input. Please try again!\033[0m');
            continue;

def _select_slurm_kernel ():

    available_slurm_kernel = list_slurm_kernel(returnonly=True);

    if available_slurm_kernel is None:
        sys.exit();

    if len(available_slurm_kernel) >= 1:
        while True:
            print("\033[4mFollowing kernels found:\033[0m\n");
            iterator = 0;
            old_slurm_kernels = False;
            for kernelfile, data in available_slurm_kernel.items():
            #print('\n\033[91m\u2B51\033[0m = Maybe an old version of a jupyter slurm kernel');

                try:
                    if data[4]['maybe_old_slurm_kernel'] == True:
                        old_slurm_kernels = True;
                        print(f'\033[94m[{iterator}]\033[0m \033[91m\u2B51\033[0m {kernelfile}');
                    else:
                        print(f'\033[94m[{iterator}]\033[0m {kernelfile}');
                except:
                    print(f'\033[94m[{iterator}]\033[0m {kernelfile}');

                iterator += 1;

            if old_slurm_kernels == True:
                print('\n\033[91m\u2B51\033[0m = Maybe an old version of a jupyter slurm kernel\n');

            identifier = input('Select a kernel using the \033[94midentifier\033[0m:\033[94m ');
            try:
                key = list(available_slurm_kernel.keys())[int(identifier)];
            except IndexError:
                print(f'\033[91mInvalid identifier: {identifier}\033[0m\n');
                continue;
            print('\033[0m');
            if key:
                return key;
            else:
                return False;


def modify_slurm_kernel (editor=None):

    selected_kernel = _select_slurm_kernel();
    if selected_kernel:
        kernelspec = selected_kernel + '/kernel.json';
        if editor is None:
            try:
                editor = os.environ['EDITOR'];
            except KeyError:
                print(f'\033[91m$EDITOR not set. Please explicity set an editor with --editor (e.g. --editor vim)\033[0m\n');
                sys.exit();

        while True:
            edit_process = subprocess.Popen(f'{editor} {kernelspec}', shell=True);
            (stdout, stderr) = edit_process.communicate();
            edit_process_status = edit_process.wait();
            try:
                # check modified kernel.json
                with open(kernelspec, 'r') as f:
                    json.loads(f.read());

            except json.decoder.JSONDecodeError:
                print(f'\033[91mError parsing the modified kernel.json!\033[0m\n');
                re_run = input('Would you like to re-run the edit mode? [Y,n] ');
                if re_run == '' or re_run.upper() == 'Y':
                    continue;
                elif re_run.upper() == 'N':
                    break;
                else:
                    print(f'\033[91mInvalid input {re_run}! Aborting.\033[0m\n');
                    break;
            
            break;
            
        # Modify done
        print(f'Successfully modified {kernelspec}!');
        sys.exit();
    

def list_slurm_kernel (all=False, verbose=False, returnonly=False):

    kernelspecs = kernelspec.find_kernel_specs();
    slurm_kernels = {};

    # iterate throuh all available kernels
    for kernel, file in kernelspecs.items():
        try:
            kernel_spec = kernelspec.get_kernel_spec(kernel);
        except json.decoder.JSONDecodeError:
            print(f'\033[91mError parsing {file}/kernel.json! Invalid JSON.\033[0m');
            continue;

        add = False;
        if all:
            add = True;
        else:
            try:
                is_slurm_kernel = kernel_spec.metadata['remote_slurm_kernel'];
                if is_slurm_kernel == True:
                    add = True;
            except KeyError:
                # maybe its an old jupyter slurm kernel
                kernelname = kernel.split('/')[-1];
                kernelname = kernelname.upper();
                check_name = kernelname.split('_');
                if 'SLURM' in check_name:
                    kernel_spec.metadata['maybe_old_slurm_kernel'] = True;
                    add = True;
        if add:
            slurm_kernels.update({str(file): [kernel_spec.name, kernel_spec.display_name, kernel_spec.language, kernel_spec.env, kernel_spec.metadata]});
    
    if len(slurm_kernels) >= 1:
        # return only
        if returnonly:
            return slurm_kernels;

        maybe_old_versions = False;

        print("\033[4mFollowing kernels found:\033[0m\n");
        for kernel, data in slurm_kernels.items():
            try:
                if data[4]['maybe_old_slurm_kernel'] == True:
                    maybe_old_versions = True;
                    print(f'\033[94m\u27A4\033[0m \033[95m{kernel} (\033[91m\u2B51\033[95m)\033[0m');
                else:
                    print(f'\033[94m\u27A4\033[0m \033[95m{kernel}\033[0m');
            
            except:
                print(f'\033[94m\u27A4\033[0m \033[95m{kernel}\033[0m');
            if verbose:
                if data[0]:
                    print(f'  \u2937  \033[1mKernel Name  : \033[0m{data[0]}');
                else:
                    kernelname = kernel.split('/')[-1];
                    print(f'  \u2937  \033[1mKernel Name  : \033[0m{kernelname}') ;
                if data[2]: print(f'  \u2937  \033[1mLanguage     : \033[0m{data[2]}');
                if data[1]: print(f'  \u2937  \033[1mDisplay Name : \033[0m{data[1]}');
                if data[3]: print(f'  \u2937  \033[1mEnvironment  : \033[0m{data[3]}');
                print('');

        if maybe_old_versions == True:
            print('\n\033[91m\u2B51\033[0m = Maybe an old version of a jupyter slurm kernel');

    else:
        print('Currently no installed kernels found!');

def yninput (inputstr=''):

    while True:
        question = input(inputstr);
        if question.upper() == 'Y' or question == '':
            return True;
        elif question.upper() == 'N':
            return False;
        else:
            continue;

def main (cmd_line=None):

    parser = argparse.ArgumentParser('Tool to manage (create, list, modify and delete) and starting jupyter slurm kernels using srun');
    subparser = parser.add_subparsers(dest='command');

    add_option = subparser.add_parser('create', help='create a new slurm kernel');

    add_option.add_argument('--displayname', required=True, help='Display name of the new kernel');
    add_option.add_argument('--environment', required=False, help='Jupyter kernel environment');
    add_option.add_argument('--language', help='Programming language');
    add_option.add_argument('--loginnode', required=True, help='The login node to connect to');
    add_option.add_argument('--user', required=True, help='The username to log in to the loginnode');
    add_option.add_argument('--proxyjump', help='Add a proxy jump (SSH -J)');
    add_option.add_argument('--srun-cmd', help='Path to srun command. Default: srun');
    add_option.add_argument('--kernel-cmd', required=True, help='command to run jupyter kernel');
    add_option.add_argument('--slurm-parameter', required=True, help='Slurm job parameter');

    list_option = subparser.add_parser('list', help='list available slurm kernel');
    list_option.add_argument('-v', '--verbose', action='store_true', required=False, help='Print all kernel with the kernelspec information');
    list_option.add_argument('-a', '--all', action='store_true', required=False, help='Print all available Jupyter kernels');

    modify_option = subparser.add_parser('edit', help='edit an existing slurm kernel');
    modify_option.add_argument('-e', '--editor', help='Set a specific editor to modify the kernelspec (default: $EDITOR)');

    delete_option = subparser.add_parser('delete', help='delete an existing slurm kernel');

    template_option = subparser.add_parser('template', help='manage script templates (list, use, add, edit)');
    template_subparser = template_option.add_subparsers(dest='subcommand');

    template_list = template_subparser.add_parser('list', help='List all availabe script templates');

    template_use = template_subparser.add_parser('use', help='Use a script template (Remote-Initialization)');
    template_use.add_argument('--dry-run', action='store_true', required=False, help='Activate dry-run mode. Do not execute script lines or create kernel.');
    template_use.add_argument('--loginnode', '-l', required=False, help='The login node to connect to');
    template_use.add_argument('--user', '-u', help='The username to log in to the loginnode');
    template_use.add_argument('--proxyjump', '-p', help='Add a proxy jump (SSH -J)');
    template_use.add_argument('--template', '-t', required=False, help='The template to use (ipython, ijulia, ...)');

    template_add = template_subparser.add_parser('add', help='Add a new script template for remote initialization');
    template_add.add_argument('--template', '-t', required=True, help='Path to the script template');

    template_edit = template_subparser.add_parser('edit', help='Edit an existing script template');
    template_edit.add_argument('--template', '-t', required=True, help='The template to use (ipython, ijulia, ...)');
    template_edit.add_argument('--editor', '-e', required=False, help='Editor to use (Default: $EDITOR)');

    from slurm_jupyter_kernel.__init__ import __version__;
    parser.add_argument('--version', action='version', version='%(prog)s ({version})'.format(version=__version__));

    if len(sys.argv) == 1:
        parser.print_help();
        sys.exit();

    args = parser.parse_args(cmd_line);
    if args.command == 'create':
        add_slurm_kernel(args.displayname, args.language, args.kernel_cmd, args.slurm_parameter, args.loginnode, args.user, args.proxyjump, args.srun_cmd, args.environment);
    elif args.command == 'list':
        list_slurm_kernel(all=args.all, verbose=args.verbose);
    elif args.command == 'edit':
        modify_slurm_kernel(args.editor);
    elif args.command == 'delete':
        delete_slurm_kernel();
    elif args.command == 'template':
        if args.subcommand == 'list':
            script_template.ScriptTemplate.list_templates();
        elif args.subcommand == 'use':
            #remote_initialization(args.loginnode, args.user, args.proxyjump, args.dry_run);
            template = script_template.ScriptTemplate(args.template);
            try:
                kernel_specs_info = template.use(args.loginnode, args.user, args.proxyjump, args.dry_run);
                add_slurm_kernel(**kernel_specs_info);
            except script_template.SSHConnectionError:
                sys.exit(f'{Color.F_LightRed}Error: Could not establish a connection to host\nPlease check following things:\n* Username is correct\n* Running SSH agent with loaded key file\n* proxyjump, loginnode or username is correct{Color.F_Default}');
            except script_template.SSHAgentNotRunning:
                sys.exit(f'{Color.F_LightRed}Error: Please start your SSH agent and load your SSH key: eval $(ssh-agent) && ssh-add{Color.F_Default}');
            except script_template.RenderTemplateError:
                sys.exit(f'{Color.F_LightRed}Error: The template seems to be broken!{Color.F_Default}');

        elif args.subcommand == 'add':
            new_template = script_template.ScriptTemplate(args.template);
            try:
                save_template = new_template.save();
                if save_template:
                    sys.exit(f'{Color.F_LightGreen}Added template {args.template}!{Color.F_Default}');
                else:
                    print('Error...');
            except script_template.WrongFileType:
                sys.exit(f'{Color.F_LightRed}Wrong filetype! ".sh" needed!{Color.F_Default}');
            except script_template.FileAlreadyExists:
                sys.exit(f'{Color.F_LightRed}The template {args.template} already exists! You may want to rename your template.{Color.F_LightRed}');
        elif args.subcommand == 'edit':
            template = script_template.ScriptTemplate(args.template);
            try:
                edit = template.edit(args.editor);
                if edit:
                    sys.exit(f'{Color.F_LightGreen}Template {args.template} successfully edited!{Color.F_Default}');
            except script_template.NoEditorGiven:
                sys.exit(f'{Color.F_LightRed}No editor specified! Neither $EDITOR is set nor an editor was specified (--editor)!{Color.F_Default}');

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