#!/bin/python3
import subprocess
import datetime
import time
import sys, os
import traceback
import json
import copy
import yaml
import re

lockfile = "/tmp/autoscaling_lock"
queues_conf_file = "/opt/oci-hpc/conf/queues.conf"

# seconds for a cluster to stay alive
idle_time=600 #seconds

script_path='/opt/oci-hpc/bin'

def israckaware():
    rackware=False
    if os.path.isfile("/opt/oci-hpc/conf/variables.tf"):
        variablefile=open("/opt/oci-hpc/conf/variables.tf",'r')
        for line in variablefile:
            if "\"rack_aware\"" in line and ("true" in line or "True" in line or "yes" in line or "Yes" in line) :
                rackware=True
                break
    return rackware

def getTopology(clusterName):
    out = subprocess.Popen(['scontrol','show','topology',clusterName], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    stdout,stderr = out.communicate()
    for item in stdout.strip().split():
        if item.startswith("Nodes="):
            nodes_condensed=item.split("Nodes=")[1]
            out2 = subprocess.Popen(['scontrol','show','hostname',nodes_condensed], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
            stdout2,stderr2 = out2.communicate()    
            print(stdout2.strip().split())
            return stdout2.strip().split()
    return []
# Get the list of Jobs in all states
def getJobs():
    # changing the position of Dependency as it is giving blank instead of null. to handle that, putting it at the end.
    out = subprocess.Popen(['squeue','-r','-O','STATE,JOBID,FEATURE:100,NUMNODES,Partition,UserName,Dependency'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    stdout,stderr = out.communicate()
    return stdout.split("\n")[1:]

def getClusters():
    out = subprocess.Popen(['sinfo','-hN','-o','\"%T %E %D %N\"'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    stdout,stderr = out.communicate()
    return stdout.split("\n")

def getNodeDetails(node):
    out = subprocess.Popen(['sinfo','-h','-n',node,'-o','"%f %R"'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    stdout,stderr = out.communicate()
    for pot_output in stdout.split("\n"):
        if not "(null)" in pot_output and pot_output.strip() != '':
            output=pot_output
            if output[0] == '"':
                output=output[1:]
            if output[-1] == '"':
                output=output[:-1]
        else:
            continue
    return output

def getIdleTime(node):
    out = subprocess.Popen(["sacct -X -n -S 01/01/01 -N "+node+" -o End | tail -n 1"],stdout=subprocess.PIPE,stderr=subprocess.STDOUT,shell=True,universal_newlines=True)
    stdout,stderr = out.communicate()
    last_end_time = None
    try:
        last_end_time = datetime.datetime.strptime(stdout.strip(),"%Y-%m-%dT%H:%M:%S")
    except:
        pass
    out = subprocess.Popen(["scontrol show node "+node+" | grep SlurmdStartTime | awk '{print $2}'"],stdout=subprocess.PIPE, stderr=subprocess.STDOUT,shell=True,universal_newlines=True)
    stdout,stderr = out.communicate()
    try: 
        cluster_start_time=datetime.datetime.strptime(stdout.split("\n")[0].split("=")[1],"%Y-%m-%dT%H:%M:%S")
    except:
        print ("The cluster start time of node "+node+" could not be found")
        print ("There seems to be an issue with the command")
        print ("scontrol show node "+node+" | grep SlurmdStartTime | awk '{print $2}'")
        print ("Here is the output it generated")
        print (stdout)
        print ("The node will be deleted")
        cluster_start_time=datetime.datetime.now()-datetime.timedelta(hours=24)
    if last_end_time is None:
        right_time=cluster_start_time
    else:
        right_time=max([cluster_start_time,last_end_time])
    return ( datetime.datetime.now() - right_time ).total_seconds()

# Get the last time a node state was changed. This is used to get how long a cluster has been idle for
def getQueueConf(queue_file):
    with open(queue_file) as file:
        try:
            data = yaml.load(file,Loader=yaml.FullLoader)
        except:
            data = yaml.load(file)
        return data["queues"]

def getQueue(config,queue_name):
    for queue in config:
        if queue["name"]==queue_name:
            return queue
    return None

def getDefaultsConfig(config,queue_name):
    for partition in config:
        if queue_name == partition["name"]:
            for instance_type in partition["instance_types"]:
                if "default" in instance_type.keys():
                    if instance_type["default"]:
                        return {"queue":partition["name"], "instance_type":instance_type["name"], "shape":instance_type["shape"], "cluster_network":instance_type["cluster_network"], "hostname_convention":instance_type["hostname_convention"]}
            if len(partition["instance_types"])>0:
                instance_type=partition["instance_types"][0]
                print ("No default configuration was found, there may be a problem in your queues.conf file")
                print ("Selecting "+instance_type["name"]+" as default")
                return {"queue":partition["name"], "instance_type":instance_type["name"], "shape":instance_type["shape"], "cluster_network":instance_type["cluster_network"], "hostname_convention":instance_type["hostname_convention"]}
    print ("The queue "+queue_name+" was not found in the queues.conf file")
    return None

def getJobConfig(config,queue_name,instance_type_name):
    for partition in config:
        if queue_name == partition["name"]:
            for instance_type in partition["instance_types"]:
                if instance_type_name == instance_type["name"]:
                    return {"queue":partition["name"], "instance_type":instance_type["name"], "shape":instance_type["shape"], "cluster_network":instance_type["cluster_network"], "hostname_convention":instance_type["hostname_convention"]}
    return None

def getQueueLimits(config,queue_name,instance_type_name):
    for partition in config:
        if queue_name == partition["name"]:
            for instance_type in partition["instance_types"]:
                if instance_type_name == instance_type["name"]:
                    return {"max_number_nodes": int(instance_type["max_number_nodes"]), "max_cluster_size": int(instance_type["max_cluster_size"]),"max_cluster_count": int(instance_type["max_cluster_count"])}
    return {"max_number_nodes": 0, "max_cluster_size": 0,"max_cluster_count": 0}

def getInstanceType(config,queue_name,hostname_convention):
    for partition in config:
        if queue_name == partition["name"]:
            for instance_type in partition["instance_types"]:
                if hostname_convention == instance_type["hostname_convention"]:
                    return instance_type["name"]
    return None

def isPermanent(config,queue_name,instance_type_name):
    for partition in config:
        if queue_name == partition["name"]:
            for instance_type in partition["instance_types"]:
                if instance_type_name == instance_type["name"]:
                    return instance_type["permanent"]
    return None

def getAllClusterNames(config):
    availableNames={}
    for partition in config:
        availableNames[partition["name"]]={}
        for instance_type in partition["instance_types"]:
            availableNames[partition["name"]][instance_type["name"]]=range(1,int(instance_type["max_cluster_count"])+1)
    return availableNames

def getClusterName(node):
    details=getNodeDetails(node)
    clusterName="NOCLUSTERFOUND"
    for feature in details[0].split(","):
        if feature.startswith('CN__'):
            clusterName=feature[4:]
    if clusterName == "NOCLUSTERFOUND":
        out = subprocess.Popen(['scontrol','show','topology',node], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
        stdout,stderr = out.communicate()
        clusterName = None
        try:
            if len(stdout.split('\n')) > 2:
                for output in stdout.split('\n')[:-1]:
                    if "Switches=" in output:
                        clusterName=output.split()[0].split('SwitchName=')[1]
                        break
                    elif "SwitchName=inactive-" in output:
                        continue
                    else:
                        clusterName=output.split()[0].split('SwitchName=')[1]
            elif len(stdout.split('\n')) == 2:
                clusterName=stdout.split('\n')[0].split()[0].split('SwitchName=')[1]
            if clusterName.startswith("inactive-"):
                return "NOCLUSTERFOUND"
        except: 
            print('No ClusterName could be found for '+node)
            print('There seems to be some issues in the slurm topology file')
            return "NOCLUSTERFOUND"
    return clusterName

def getstatus_slurm():
    cluster_to_build=[]

    # Get cluster to build
    # squeue -O STATE,JOBID,FEATURE:100,NUMNODES,Partition,UserName,Dependency
    for line in getJobs():
        if len(line.split())>3:
            new_line=re.split(r"\s{1,}", line)
            if new_line[0] == 'PENDING' and ('null' in new_line[6] or len(new_line[6])==0):
                queue = new_line[4]
                user = new_line[5]
                features=new_line[2].split('&')
                instanceType= None
                possible_types=[inst_type["name"] for inst_type in getQueue(config,queue)["instance_types"]]
                default_config=getDefaultsConfig(config,queue)
                if instanceType is None:
                    instanceType = default_config["instance_type"]
                    for feature in features:
                        if feature in possible_types:
                            instanceType=feature
                            break

                nodes=int(new_line[3])
                jobID=int(new_line[1])
                if isPermanent(config,queue,instanceType) is None :
                    continue
                elif isPermanent(config,queue,instanceType):
                    continue
                else:
                    cluster_to_build.append([nodes,instanceType,queue,jobID,user])

    cluster_to_destroy=[]
    current_nodes={}
    building_nodes={}
    running_cluster=[]
    nodes_to_destroy_temp={}
    nodes_to_destroy={}
    # Get Cluster to destroy, or nodes to destroy
    for line in getClusters():
        if len(line.split()) == 0:
            break
        old_nodes=line.split()[-1].split(',')
        brokenListOfNodes=False
        nodes=[]
        for node in old_nodes:
            if brokenListOfNodes:
                if ']' in node:
                    brokenListOfNodes=False
                    nodes.append(currentNode+','+node)
                else:
                    currentNode=currentNode+','+node
            elif '[' in node and not ']' in node:
                brokenListOfNodes=True
                currentNode=node
            else:
                nodes.append(node)
        clusters = []
        for node in nodes:
            if node[-1]=="\"":
                node=node[:-1]
            if node[0]=="\"":
                node=node[1:]
            details=getNodeDetails(node).split(' ')
            features=details[0].split(',')
            queue=details[-1]
            clustername=getClusterName(node)
            if clustername is None:
                continue
            instanceType=features[0]
            if queue in current_nodes.keys():
                if instanceType in current_nodes[queue].keys():
                    current_nodes[queue][instanceType]+=1
                else:
                    current_nodes[queue][instanceType]=1
            else:
                current_nodes[queue]={instanceType:1}
            if line.split()[0] == '\"idle' or line.split()[0] == '\"down' or line.split()[0].endswith('*'):
                if not line.split()[0].endswith('*') and clustername == "NOCLUSTERFOUND": 
                    continue
                if not os.path.isdir(os.path.join(clusters_path,clustername)) and clustername != "NOCLUSTERFOUND":
                    continue
                node_idle_time=getIdleTime(node)
                if node_idle_time<idle_time:
                    print (clustername + " is too young to die : "+str(node_idle_time) + " : "+node)
                    continue
                if isPermanent(config,queue,instanceType) is None :
                    continue
                elif isPermanent(config,queue,instanceType):
                    continue
                if not clustername in nodes_to_destroy_temp.keys():
                    nodes_to_destroy_temp[clustername]=[]
                nodes_to_destroy_temp[clustername].append(node)
            elif line.split()[0] == '\"allocated' or line.split()[0] == '\"mixed':
                if not clustername in running_cluster:
                    running_cluster.append(clustername)
    cluster_to_destroy=[]
    for clustername in nodes_to_destroy_temp.keys():
        destroyEntireCluster=True
        if clustername == "NOCLUSTERFOUND":
            destroyEntireCluster=False
        elif clustername in running_cluster:
            nodes_to_destroy[clustername]=nodes_to_destroy_temp[clustername]
            destroyEntireCluster=False
        else:
            for node in getTopology(clustername):
                if not node in nodes_to_destroy_temp[clustername]:
                    nodes_to_destroy[clustername]=nodes_to_destroy_temp[clustername]
                    destroyEntireCluster=False
                    break
        if destroyEntireCluster:
            cluster_to_destroy.append([clustername])

    cluster_building=[]
    cluster_destroying=[]

    used_index={}
    for clusterName in os.listdir(clusters_path):
        if len(clusterName.split('-')) < 3:
            continue
        hostname_convention='-'.join(clusterName.split('-')[2:])
        clusterNumber=int(clusterName.split('-')[1])
        queue=clusterName.split('-')[0]
        instanceType=getInstanceType(config,queue,hostname_convention)
        if not queue in used_index.keys():
            used_index[queue]={}
        if not instanceType in used_index[queue].keys():
            used_index[queue][instanceType]=[]
        used_index[queue][instanceType].append(clusterNumber)

        if os.path.isfile(os.path.join(clusters_path,clusterName,'currently_building')):
            with open(os.path.join(clusters_path,clusterName,'currently_building'),'r') as f:
                line = f.read()
                nodes = line.split()[0]
                instance_type = line.split()[1]
                queue = line.split()[2]
                try:
                    cluster_building.append([int(nodes),instance_type,queue])
                    if queue in building_nodes.keys():
                        if instance_type in building_nodes[queue].keys():
                            building_nodes[queue][instance_type]+=int(nodes)
                        else:
                            building_nodes[queue][instance_type]=int(nodes)
                    else:
                        building_nodes[queue]={instance_type:int(nodes)}
                except ValueError:
                    print ('The cluster '+ clusterName + ' does not have a valid entry for \"currently_building\"')
                    print ('Ignoring')
                    continue
        if os.path.isfile(os.path.join(clusters_path,clusterName,'currently_destroying')):
            cluster_destroying.append(clusterName)
    return cluster_to_build,cluster_to_destroy,nodes_to_destroy,cluster_building,cluster_destroying,used_index,current_nodes,building_nodes

def getAutoscaling():
    out = subprocess.Popen(["cat /etc/ansible/hosts | grep 'autoscaling =' | awk -F  '= ' '{print $2}'"],stdout=subprocess.PIPE, stderr=subprocess.STDOUT,shell=True,universal_newlines=True)
    stdout,stderr = out.communicate()
    output = stdout.split("\n")
    autoscaling_value=False
    for i in range(0,len(output)-1):
        autoscaling_value=output[i]
    return autoscaling_value

autoscaling = getAutoscaling()

if autoscaling == "true":

    if os.path.isfile(lockfile):
        print( "Lockfile "+lockfile + " is present, exiting" )
        exit()
    open(lockfile,'w').close()
    try:
        path = os.path.dirname(os.path.dirname(os.path.realpath(sys.argv[0])))
        clusters_path = os.path.join(path,'clusters')
        config = getQueueConf(queues_conf_file)

        cluster_to_build,cluster_to_destroy,nodes_to_destroy,cluster_building,cluster_destroying,used_index,current_nodes,building_nodes=getstatus_slurm()

        print (time.strftime("%Y-%m-%d %H:%M:%S"))
        print (cluster_to_build,'cluster_to_build')
        print (cluster_to_destroy,'cluster_to_destroy')
        print (nodes_to_destroy,'nodes_to_destroy')
        print (cluster_building,'cluster_building')
        print (cluster_destroying,'cluster_destroying')
        print (current_nodes,'current_nodes')
        print (building_nodes,'building_nodes')

        for i in cluster_building:
            for j in cluster_to_build:
                if i[0]==j[0] and i[1]==j[1] and i[2]==j[2]:
                    cluster_to_build.remove(j)
                    break
        for cluster in cluster_to_destroy:
            cluster_name=cluster[0]
            print ("Deleting cluster "+cluster_name)
            subprocess.Popen([script_path+'/delete_cluster.sh',cluster_name])
            time.sleep(5)

        for cluster_name in nodes_to_destroy.keys():
            print ("Resizing cluster "+cluster_name)
            initial_nodes=[]
            unreachable_nodes=[]
            if cluster_name == "NOCLUSTERFOUND":
                subprocess.Popen([script_path+'/resize.sh','remove_unreachable','--nodes']+nodes_to_destroy[cluster_name],'--quiet')
                continue
            for node in nodes_to_destroy[cluster_name]:
                try:
                    alt_names=subprocess.check_output(["cat /etc/hosts | grep "+node],shell=True,universal_newlines=True)
                    for alt_name in alt_names.split("\n")[0].split():
                        if alt_name.startswith('inst-'):
                            initial_nodes.append(alt_name)
                            break
                except:
                    unreachable_nodes.append(node)    
            if len(initial_nodes) > 0:
                subprocess.Popen([script_path+'/resize.sh','--force','--cluster_name',cluster_name,'remove','--remove_unreachable','--nodes']+initial_nodes,'--quiet')
            if len(unreachable_nodes) > 0:
                subprocess.Popen([script_path+'/resize.sh','--cluster_name',cluster_name,'remove_unreachable','--nodes']+unreachable_nodes,'--quiet')
            time.sleep(1)

        for index,cluster in enumerate(cluster_to_build):
            nodes=cluster[0]
            instance_type = cluster[1]
            queue=cluster[2]
            jobID=str(cluster[3])
            user=str(cluster[4])
            jobconfig=getJobConfig(config,queue,instance_type)
            limits=getQueueLimits(config,queue,instance_type)
            try:
                clusterCount=len(used_index[queue][instance_type])
            except:
                clusterCount=0
            if clusterCount>=limits["max_cluster_count"]:
                print ("This would go over the number of running clusters, you have reached the max number of clusters")
                continue
            nextIndex=None
            if clusterCount==0:
                if queue in used_index.keys():
                    used_index[queue][instance_type]=[1]
                else:
                    used_index[queue]={instance_type:[1]}
                nextIndex=1
            else:
                for i in range(1,10000):
                    if not i in used_index[queue][instance_type]:
                        nextIndex=i
                        used_index[queue][instance_type].append(i)
                        break
            clusterName=queue+'-'+str(nextIndex)+'-'+jobconfig["hostname_convention"]
            if not queue in current_nodes.keys():
                current_nodes[queue]={instance_type:0}
            else:
                if not instance_type in current_nodes[queue].keys():
                    current_nodes[queue][instance_type]=0
            if not queue in building_nodes.keys():
                building_nodes[queue]={instance_type:0}
            else:
                if not instance_type in building_nodes[queue].keys():
                    building_nodes[queue][instance_type]=0
            if nodes > limits["max_cluster_size"]:
                print ("Cluster "+clusterName+" won't be created, it would go over the total number of nodes per cluster limit")
            elif current_nodes[queue][instance_type] + building_nodes[queue][instance_type] + nodes > limits["max_number_nodes"]:
                print ("Cluster "+clusterName+" won't be created, it would go over the total number of nodes limit")
            else:
                current_nodes[queue][instance_type]+=nodes
                clusterCount+=1
                print ("Creating cluster "+clusterName+" with "+str(nodes)+" nodes")
                subprocess.Popen([script_path+'/create_cluster.sh',str(nodes),clusterName,instance_type,queue,jobID,user])
                time.sleep(5)

    except Exception:
        traceback.print_exc()
    os.remove(lockfile)
else:
    print("Autoscaling is false (set in /etc/ansible/hosts)")
    exit()
