#!/usr/bin/env python

import subprocess
import argparse
import shutil
import yaml
import sys
import os


def parse_kubectl_describe_nodes():
    """
    Kubectl describe nodes gives a lot of information that for the most part we don't really care about at all.
    This method will run that command and give slurm-esque output (number of nodes, node names, and current loads
    on each node)
    """
    kubectl_output = subprocess.check_output('kubectl describe nodes', shell=True)
    node_names = list()
    node_cpu_loads = dict()
    node_memory_loads = dict()
    node_cpu_capacity = dict()
    node_memory_capacity = dict()
    # This is not a great way to do this, but I can't seem to get the info I need in any other command.
    # Need to do some more kubectl doc reading to see if I've missed something.

    # Iterate through the output created and populate our data dictionaries.
    # This way of doing things is definitely potentially very brittle - hopefully kubectl keeps output
    # looking consistent between versions.
    kubectl_lines = kubectl_output.decode('utf-8').split('\n')
    current_node_name = None
    for i in range(len(kubectl_lines)):
        if 'Name:' in kubectl_lines[i]:
            current_node_name = kubectl_lines[i].split()[1].rstrip()
            node_names.append(current_node_name)
        if current_node_name is not None:
            if 'cpu:' in kubectl_lines[i] and 'Capacity:' in kubectl_lines[i - 1]:
                node_cpu_capacity[current_node_name] = kubectl_lines[i].split()[1].rstrip()
            elif 'memory:' in kubectl_lines[i] and 'Capacity:' in kubectl_lines[i - 2]:
                node_memory_capacity[current_node_name] = kubectl_lines[i].split()[1].rstrip()
            elif 'cpu' in kubectl_lines[i] and 'Allocated' in kubectl_lines[i - 4]:
                node_cpu_loads[current_node_name] = kubectl_lines[i].split()[2].rstrip()
            elif 'memory' in kubectl_lines[i] and 'Allocated' in kubectl_lines[i - 5]:
                node_memory_loads[current_node_name] = kubectl_lines[i].split()[2].rstrip()

    # Now write the output of data gathered in a user-friendly(ish) format.
    print('Number of nodes in cluster: {}'.format(len(node_names)))
    print('NodeName\tCPU_Capacity\tCPU_Usage\tMemory_Capacity\tMemory_Usage')
    for node_name in node_names:
        print('{name}\t{cpucapacity}\t{cpuusage}\t{memorycapacity}\t{memoryusage}'.format(name=node_name,
                                                                                          cpucapacity=node_cpu_capacity[node_name],
                                                                                          cpuusage=node_cpu_loads[node_name],
                                                                                          memorycapacity=node_memory_capacity[node_name],
                                                                                          memoryusage=node_memory_loads[node_name]))


def get_node_names(node_dict):
    node_names = list()
    for item in node_dict['items']:
        node_names.append(item['metadata']['name'])
    return node_names


def get_node_cpu_capacity(node_name, node_dict):
    for item in node_dict['items']:
        if item['metadata']['name'] == node_name:
            return item['status']['capacity']['cpu']


def get_node_memory_capacity(node_name, node_dict):
    for item in node_dict['items']:
        if item['metadata']['name'] == node_name:
            return int(item['status']['capacity']['memory'].replace('Ki', ''))


def check_cpu_request(node_names, node_dict, requested_cpus):
    """
    Checks that the requested number of CPUs can actually be hosted on one of the nodes.
    :param node_names: List of node names, generated by get_node_names.
    :param node_dict: Dictionary of info about the cluster, generated by parsing the yaml output from kubectl get nodes
    -o=yaml
    :param requested_cpus: Number of CPUs the user has requested.
    :return: True if at least one node has sufficient CPU available to satisfy the request. False if requested
    number of CPUs is too high.
    """
    for node_name in node_names:
        node_cpu_capacity = get_node_cpu_capacity(node_name=node_name,
                                                  node_dict=node_dict)
        if requested_cpus <= int(node_cpu_capacity):
            return True
    return False


def check_memory_request(node_names, node_dict, requested_memory):
    # This won't actually work yet - need to get conversions done for mem.
    """
    Checks that the requested number of CPUs can actually be hosted on one of the nodes.
    :param node_names: List of node names, generated by get_node_names.
    :param node_dict: Dictionary of info about the cluster, generated by parsing the yaml output from kubectl get nodes
    -o=yaml
    :param requested_memory: Number of CPUs the user has requested.
    :return: True if at least one node has sufficient memory available to satisfy the request. False if requested
    amount of memory is too high.
    """
    for node_name in node_names:
        node_memory_capacity = get_node_memory_capacity(node_name=node_name,
                                                        node_dict=node_dict)
        if requested_memory <= node_memory_capacity/1000000:  # We request memory in GB, but node memory capacity is in KB, so divide by a million
            return True
    return False


def check_for_job_name(job_name):
    """
    Kubernetes does not seem to like it if a job is submitted that has the same name as something else.
    It seems to do reconfiguration, which could probably mess things up. I'm going to assume people aren't
    going to be trying to reconfigure their jobs on the fly. This will check if a job with specified name
    exists so we can give the user an error.
    :param job_name: Name of job user wants to submit to kubernetes
    :return: True if job with same name already exists, False if no job with that name found.
    """
    job_exists = False
    kubectl_job_output = subprocess.check_output('kubectl get jobs', shell=True)
    job_lines = kubectl_job_output.decode('utf-8').split('\n')
    for i in range(1, len(job_lines) - 1):  # An extra line gets added when decoding output, don't need to go through it
        if job_lines[i].split()[0] == job_name:
            job_exists = True
    return job_exists


def check_python_version():
    """
    This script is python3 only - double check that a py3 interpreter is being called and give a nice(ish) error
    if it isn't.
    :return: True if python version if 3.x, False if python version is 2.x
    """
    if sys.version_info[0] < 3:
        return False
    else:
        return True


if __name__ == '__main__':
    # Argparse - should probably have subcommands here - main one for submitting jobs, others
    # for getting cluster info and whatnot.
    parser = argparse.ArgumentParser(description='KubeJobSub')
    subparsers = parser.add_subparsers(help='SubCommand Help', dest='subparsers')
    job_sub_parser = subparsers.add_parser('submit', help='Submits a job to your kubernetes cluster. Configured'
                                                          ' to assume you\'re using azure with a file mount.')
    job_sub_parser.add_argument('-j', '--job_name',
                                type=str,
                                required=True,
                                help='Name of job.')
    job_sub_parser.add_argument('-c', '--command',
                                type=str,
                                required=True,
                                help='The command you want to run. Put it in double quotes. (")')
    job_sub_parser.add_argument('-i', '--image',
                                type=str,
                                required=True,
                                help='Docker image to create container from.')
    job_sub_parser.add_argument('-n', '--num_cpu',
                                type=int,
                                default=1,
                                help='Number of CPUs to request for your job. Must be an integer greater than 0.'
                                     ' Defaults to 1.')
    job_sub_parser.add_argument('-m', '--memory',
                                type=float,
                                default=2,
                                help='Amount of memory to request, in GB. Defaults to 2.')
    job_sub_parser.add_argument('-v', '--volume',
                                type=str,
                                default='/mnt/azure',
                                help='The mountpath for your azure file share. Defaults to /mnt/azure')
    job_sub_parser.add_argument('-share', '--share',
                                type=str,
                                required=True,
                                help='Name of Azure file share that you want mounted to the point specified by -v')
    job_sub_parser.add_argument('-secret', '--secret',
                                type=str,
                                default='azure-secret',
                                help='The name of the secret created by kubectl for azure file mounting. Defaults to'
                                     ' azure-secret. See https://docs.microsoft.com/en-us/azure/aks/azure-files-volume'
                                     ' for more information on creating your own.')
    job_sub_parser.add_argument('-k', '--keep',
                                default=False,
                                action='store_true',
                                help='A YAML file will be created to submit your job. Deleted by default once'
                                     ' job is submitted, but if this flag is active it will be kept.')
    info_parser = subparsers.add_parser('info', help='Tells you about your kubernetes cluster - number of nodes,'
                                                     'and specs/usage for each node.')
    args = parser.parse_args()

    # Check that we're actually using python 3.
    if check_python_version() is False:
        print('ERROR: It looks like you\'re using python2, but this script needs python3. Exiting...')
        quit(code=1)

    if shutil.which('kubectl') is None:
        print('ERROR: kubectl not found. Please verify that kubernetes is installed.')
        quit(code=1)

    # We need to know all about the cluster. Best way to do this seems to be to make kubectl output
    # some stuff in YAML format, then parse it to figure things out.
    # Calls to kubectl don't seem to be quick, so just do this once at the start of the program,
    # and then pass the dictionary around as necessary
    yaml_info = subprocess.check_output('kubectl get nodes -o=yaml', shell=True).decode('utf-8')
    node_dict = yaml.load(yaml_info)
    node_names = get_node_names(node_dict)

    if 'submit' == args.subparsers:
        # Need to:
        # 1) Validate that the parameters put in actually make sense (i.e. request doesn't require more CPU/mem
        # than is available on nodes)
        if check_cpu_request(node_names=node_names, node_dict=node_dict, requested_cpus=args.num_cpu) is False:
            print('You requested {} CPUs for your job, but no nodes have that many CPUs available.'
                  ' Run KubeJobSub.py info to see information about your cluster.'.format(args.num_cpu))
            quit(code=1)
        if check_memory_request(node_names=node_names, node_dict=node_dict, requested_memory=args.memory) is False:
            print('You requested {} GB of RAM for your job, but no nodes have that much RAM available.'
                  ' Run KubeJobSub.py info to see information about your cluster.'.format(args.memory))
            quit(code=1)
        if check_for_job_name(args.job_name):
            print('A job with that name already exists. Kubernetes does not like that, so please select a new name '
                  'for your job.')
            quit(code=1)
        # 2) Create a YAML file using the supplied info and submit job.
        info_dict = {'apiVersion': 'batch/v1', 'kind': 'Job',
                     'metadata': {'name': args.job_name},
                     'spec': {'template': {'spec':
                                               {'containers':
                                                    [{'name': args.job_name, 'image': args.image, 'resources':
                                                        {'requests': {'cpu': args.num_cpu, 'memory': str(args.memory) + 'Gi'}},
                                                      'command': args.command.split(),
                                                      'volumeMounts': [{'mountPath': args.volume, 'name': args.job_name}]}],
                                                'restartPolicy': 'Never', 'volumes':
                                                    [{'name': args.job_name, 'azureFile':
                                                        {'secretName': args.secret,
                                                         'shareName': args.share, 'readOnly': False}}]}}}}

        with open(args.job_name, 'w') as f:
            yaml.dump(info_dict, f)
        cmd = 'kubectl apply -f {}'.format(args.job_name)
        subprocess.call(cmd, shell=True)

        # 3) Cleanup YAML file, unless specifed that it should be kept.
        if not args.keep:
            os.remove(args.job_name)

    elif 'info' == args.subparsers:
        parse_kubectl_describe_nodes()
