import hudson.model.*
import jenkins.model.*
import groovy.json.JsonSlurper
import java.io.BufferedReader
import java.io.InputStreamReader
import java.io.OutputStreamWriter
import java.net.URL
import java.net.URLConnection

def sendPostRequest(urlString, paramString, username, password) {
    def url = new URL(urlString)
    def conn = url.openConnection()
    String userpass = username + ":" + password;
    String basicAuth = "Basic " + javax.xml.bind.DatatypeConverter.printBase64Binary(userpass.getBytes());
    conn.setRequestProperty ("Authorization", basicAuth);
    conn.setDoOutput(true)
    def writer = new OutputStreamWriter(conn.getOutputStream())

    writer.write(paramString)
    writer.flush()
    String line
    def reader = new BufferedReader(new     InputStreamReader(conn.getInputStream()))
    while ((line = reader.readLine()) != null) {
      println line
    }
    writer.close()
    reader.close()
}

def execute_command(cmd, args, verbose=false, return_stdout=false) {
    if (!cmd.toString()?.trim()) {
        cmd = "python "+System.getenv("JENKINS_HOME")+"/jenkins-scripts/common/aws_ec2.py"
    }

    println "${cmd} ${args}"

    Process process = "${cmd} ${args}".execute()
    def out = new StringBuffer()
    def err = new StringBuffer()
    process.consumeProcessOutput( out, err )
    process.waitFor()
    if (verbose == true) {
        println "\n<<<< START CMD: ${args.split()[0]} >>>>"
        println "OUT:\n" + out.toString()
        println "ERR:\n" + err.toString()
        println "<<<< END CMD >>>>\n"
    }
    if (return_stdout == true) {
        def rs = "\n" + out.toString()
        return rs
    }

    def HashMap ret_items
    def ret_err = err.toString()
    def ret_out = out.toString()

    ret_err = ret_err.replaceAll(
                      "(?m).*unable to resolve host.*", "").replaceAll(
                      "(?m).*redirecting to systemctl stop obsworker.*", "").replaceAll(
                      "(?m).*Permanently added .* to the list of known hosts.*", "")

    if (ret_err?.trim()) {
        println "You got error message:"
        //assert false
        return 'FATAL'
    } else {
        out_str = ret_out.replace("None", "''")
        if (out_str.contains("-----BEGIN RC OUTPUT-----") && \
            out_str.contains("-----END RC OUTPUT-----")) {
            ret_data = out_str.substring( \
                           out_str.indexOf("-----BEGIN RC OUTPUT-----") + 26, \
                           out_str.indexOf("-----END RC OUTPUT-----"))
            def jsonSlurper = new JsonSlurper()
            ret_items = jsonSlurper.parseText(ret_data.replaceAll("'","\""))
        }
    }
    return ret_items
}

public class WorkerConf {
    protected Properties conf
    protected slot_number = ""
    protected boolean need_slave_attach = false

    WorkerConf() {
        conf = new Properties()
        conf.load(new FileInputStream(System.getenv('JENKINS_HOME') \
                                      + '/init.groovy.d/setup.properties'))
    }

    String make_init_script() {
        return """#!bin/bash
echo "ONDEMAND SCAILING"
"""
    }

    def set_slot_number( f ) { this.slot_number = f }
    String get_slot_number() { this.slot_number }

    boolean slave_attach_needed() { this.need_slave_attach }

    Integer max_slaves() { conf.EC2_INSTANCE_CAP_STR.toInteger() }
    Integer executors_per_slave() { conf.EC2_NUMBER_OF_EXECUTORS.toInteger() }
    String name_prefix() { conf.EC2_DESCRIPTION }
    String keyname() { conf.AWS_CLOUD_KEYNAME }

    String ami_id() { conf.EC2_AMI_ID }
    String availability_zone() { conf.EC2_AV_ZONE }
    String placement() {
        return "Tenancy:default," \
             + "GroupName:," \
             + "AvailabilityZone:${this.availability_zone()}"
    }
    String security_groups() { conf.EC2_SECURITY_GROUPS.replaceAll(" ","") }
    String remote_fs() { conf.EC2_REMOTE_FS }
    String instance_type() { conf.EC2_INSTANCE_TYPE }
    String labels() { conf.EC2_LABEL_STRING }
    String remote_user() { conf.EC2_REMOTE_ADMIN }
    String ssh_port() { conf.EC2_SSH_PORT }
    String subnet_id() { conf.EC2_SUBNET_ID }
    String launch_timeout() { conf.EC2_LAUNCH_TIMEOUT }
    long idle_termination_minutes() { conf.EC2_IDLE_TERMINATION_MINUTES.toInteger() }
    String credential_id() { conf.EC2_CREDENTIAL_ID }
    String tag_name() { conf.EC2_TAG_NAME }
    String tag_hostname() { conf.EC2_TAG_HOSTNAME + this.get_slot_number() }
    String tag_env() { conf.EC2_TAG_ENV }
    String tag_source() { conf.EC2_TAG_SOURCE }
    String tags() {
        return "Name:${this.tag_name()}," \
             + "hostname:${this.tag_hostname()}," \
             + "env:${this.tag_env()}," \
             + "source:${this.tag_source()}," \
             + "slot:${this.get_slot_number()}," \
             + "jenkins_slave_type:demand_${this.name_prefix()}"
    }
    String userdata() { return this.make_init_script() }
    Integer instance_base() { conf.EC2_WORKER_INSTANCE_BASE.toInteger() }
    boolean health_check_enabled() { return false }
    String get_docker_compose_path() { return "/root/docker-compose.yml" }
}

public class WorkerConf_JENKINS_IMAGER extends WorkerConf {
    public WorkerConf_JENKINS_IMAGER() {
        need_slave_attach = true
    }
    @Override
    String name_prefix() { conf.EC2_WORKER_IMAGER_DESCRIPTION }
    String tag_source() { conf.EC2_WORKER_TAG_SOURCE }
    List check_queue_list() { conf.EC2_WORKER_IMAGER_QUEUE_CHECK_LIST.split(",") }
    String get_remote_ssh_priv_key() { conf.EC2_WORKER_IMAGER_REMOTE_SSH_PRIV_KEY }
    String get_remote_ssh_known_hosts() { conf.EC2_WORKER_IMAGER_REMOTE_SSH_KNOWN_HOSTS }
    String make_init_script() {
        def String fileContents = new File(this.get_remote_ssh_priv_key()).text
        def String fileContentsKnownHosts = new File(this.get_remote_ssh_known_hosts()).text
        return """#!bin/bash
sed -i "s/[0-9]\\+/0/g" /etc/apt/apt.conf.d/20auto-upgrades
sed -i "s/[0-9]\\+/0/g" /etc/apt/apt.conf.d/10periodic
sed -i "s/^\t/\\/\\/\t/g" /etc/apt/apt.conf.d/50unattended-upgrades
echo "${fileContents}" >> /home/${this.remote_user()}/.ssh/id_rsa
echo "User ${this.remote_user()}" > /home/${this.remote_user()}/.ssh/config
echo "IdentityFile /home/${this.remote_user()}/.ssh/id_rsa" >> /home/${this.remote_user()}/.ssh/config
chown -R ${this.remote_user()}:${this.remote_user()} ${this.remote_fs()}
chmod 0600 /home/${this.remote_user()}/.ssh/id_rsa
"""
    }
}

public class WorkerConf_OBS_WORKER_NORMAL extends WorkerConf {
    public WorkerConf_OBS_WORKER_NORMAL() {
        need_slave_attach = false
    }
    @Override
    String name_prefix() { conf.EC2_WORKER_OBS_NORMAL_DESCRIPTION }
    String labels() { conf.EC2_WORKER_OBS_NORMAL_LABEL_STRING }
    String ami_id() { conf.EC2_WORKER_OBS_NORMAL_AMI_ID }
    String security_groups() { conf.EC2_WORKER_OBS_SECURITY_GROUPS.replaceAll(" ","") }
    String instance_type() { conf.EC2_WORKER_OBS_NORMAL_INSTANCE_TYPE }
    String tag_name() { conf.EC2_WORKER_OBS_NORMAL_TAG_NAME }
    String tag_hostname() { conf.EC2_WORKER_OBS_NORMAL_TAG_HOSTNAME + this.get_slot_number() }
    String get_remote_ssh_pub_key() { conf.EC2_WORKER_OBS_REMOTE_SSH_PUB_KEY }
    String get_hostname_prefix() { return "obsnw"; }
    Integer get_obs_backend_02_num() { conf.EC2_WORKER_OBS_NORMAL_BACKEND02_NUM }
    String make_init_script() {
        def hostname_prefix = this.get_hostname_prefix();
        def docker_compose_file = this.get_docker_compose_path()
        def String fileContents = new File(this.get_remote_ssh_pub_key()).text
        def backend02_num = this.get_obs_backend_02_num();
        def instance_base = this.instance_base()
        return """#!bin/bash
docker exec obs_worker rcobsworker stop
echo "ONDEMAND SCAILING: Configuring hostname for $hostname_prefix to $hostname_prefix$slot_number"
echo $hostname_prefix$slot_number > /etc/hostname
sed -i "s/#restart/restart/g" $docker_compose_file
sed -i "s/$hostname_prefix[0-9]\\+/$hostname_prefix$slot_number/g" /etc/hosts
sed -i "s/$hostname_prefix[0-9]\\+/$hostname_prefix$slot_number/g" $docker_compose_file
#Configure multiple backends
slot_number=$slot_number
instance_base=$instance_base
backend02_num=$backend02_num
backend02_max=\$((instance_base + backend02_num))
if [ \$slot_number -lt \$backend02_max ]; then
  sed -i "s/obsbackend01:5252/obsbackend01:5252 obsbackend02:5252/g" /root/obs_worker/obs-server.j2
fi
docker-compose -f $docker_compose_file up -d
useradd -m ${this.remote_user()} -s /bin/bash
mkdir -p /home/${this.remote_user()}/.ssh/
echo "${fileContents}" > /home/${this.remote_user()}/.ssh/authorized_keys
adduser ${this.remote_user()} sudo
echo "${this.remote_user()} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers
chown -R ${this.remote_user()}:${this.remote_user()} ${this.remote_fs()}
"""
    }
    Integer executors_per_slave() { conf.EC2_WORKER_OBS_NORMAL_NUMBER_OF_EXECUTORS.toInteger() }
    Integer max_slaves() { conf.EC2_WORKER_OBS_NORMAL_INSTANCE_CAP_STR.toInteger() }
    String tag_source() { conf.EC2_WORKER_TAG_SOURCE }
    boolean health_check_enabled() { return true }
}

class WorkerConf_OBS_WORKER_POWER extends WorkerConf_OBS_WORKER_NORMAL {
    @Override
    String name_prefix() { conf.EC2_WORKER_OBS_POWER_DESCRIPTION }
    String labels() { conf.EC2_WORKER_OBS_POWER_LABEL_STRING }
    String ami_id() { conf.EC2_WORKER_OBS_POWER_AMI_ID }
    String instance_type() { conf.EC2_WORKER_OBS_POWER_INSTANCE_TYPE }
    String tag_name() { conf.EC2_WORKER_OBS_POWER_TAG_NAME }
    String tag_hostname() { conf.EC2_WORKER_OBS_POWER_TAG_HOSTNAME + this.get_slot_number() }

    String get_hostname_prefix() { return "obspw"; }
    Integer executors_per_slave() { conf.EC2_WORKER_OBS_POWER_NUMBER_OF_EXECUTORS.toInteger() }
    Integer max_slaves() { conf.EC2_WORKER_OBS_POWER_INSTANCE_CAP_STR.toInteger() }
    Integer get_obs_backend_02_num() { conf.EC2_WORKER_OBS_POWER_BACKEND02_NUM }
}

class SlaveStatus {
    private Map slave_stat = ['CURR_NUMBER_OF_NODES': 0,
                              'CURR_TOTAL_EXECUTORS': 0,
                              'CURR_BUSY_EXECUTORS': 0,
                              'CURR_IDLE_EXECUTORS': 0]
    private Map slave_info = [:]

    SlaveStatus(String name_prefix) {
        for (slave in Hudson.instance.slaves) {
            def slaveComputer = slave.getComputer()
            if ("${slaveComputer.getName()}".startsWith(name_prefix)) {
                slave_stat['CURR_NUMBER_OF_NODES']++
                slave_stat['CURR_TOTAL_EXECUTORS'] += slaveComputer.countExecutors()
                slave_stat['CURR_BUSY_EXECUTORS'] += slaveComputer.countBusy()
                if (!slaveComputer.offline && slaveComputer.isAlive()) {
                    slave_stat['CURR_IDLE_EXECUTORS'] += slaveComputer.countIdle()
                }
                def this_slave = ["name":slaveComputer.getName(),
                                  "idle_since":System.currentTimeMillis() \
                                               - slaveComputer.getIdleStartMilliseconds(),
                                  "host":slave.getLauncher().getHost(),
                                  "object":slaveComputer]
                slave_info[slaveComputer.getName()] = this_slave
            }
        }
    }

    Integer number_of_nodes() {
        slave_stat['CURR_NUMBER_OF_NODES']
    }
    Integer total_executors() {
        slave_stat['CURR_TOTAL_EXECUTORS']
    }
    Integer idle_executors() {
        slave_stat['CURR_IDLE_EXECUTORS']
    }
    Integer busy_executors() {
        slave_stat['CURR_BUSY_EXECUTORS']
    }
    Map get_slave_info() {
        return slave_info
    }
    void disconnect_node(slaveName) {
        if (slaveName in slave_info) {
            slave_info[slaveName]["object"].setTemporarilyOffline(true, null)
            slave_info[slaveName]["object"].disconnect(null)
        }
    }
    void delete_node(slaveName) {
        if (slaveName in slave_info) {
            slave_info[slaveName]["object"].doDoDelete()
        }
    }
}

def get_aws_status(worker_conf) {
    args = "ping -t ${worker_conf.tags()}"
    println args;
    return execute_command("", args, verbose=true)
}

def terminate_aws_ec2_instances(instance_ids) {
    args = "terminate -i ${instance_ids.join(',')}"
    return execute_command("", args, verbose=true)
}

def create_aws_ec2_instances(ami_id,
                             min_count,
                             max_count,
                             keyname,
                             security_group_ids,
                             instance_type,
                             placement,
                             subnet_id,
                             ebsoptimized,
                             tags,
                             userdata) {

    String encoded = userdata.bytes.encodeBase64().toString()

    args = "create --amiid ${ami_id} " \
           + " --keyname ${keyname} " \
           + " --mincount ${min_count} --maxcount ${max_count} " \
           + " --securitygroupids ${security_group_ids} " \
           + " --instancetype ${instance_type} " \
           + " --placement ${placement} " \
           + " --subnetid ${subnet_id} " \
           + " --ebsoptimized ${ebsoptimized} " \
           + " --tags ${tags}" \
           + " --userdata ${encoded}"
    return execute_command("", args, verbose=true)
}

def get_worker_conf(purpose) {

    def worker_conf = Class.forName("WorkerConf_${purpose}", true, \
                                    this.class.classLoader).newInstance()
    assert worker_conf
    println "\nWORKER CONFIGURATION (${purpose})"
    println "  MAX_SLAVES: ${worker_conf.max_slaves()}"
    println "  EXECUTORS_PER_SLAVE: ${worker_conf.executors_per_slave()}"
    println "  NAME_PREFIX: ${worker_conf.name_prefix()}\n"
    return worker_conf
}

def check_healthy_status(worker_conf, vm_list) {

    if (worker_conf.health_check_enabled() != true) {
        return
    }
    def inst_for_check = []

    vm_list.each { k, v ->
        if (v["state"] != "terminated" && v["state"] != "shutting-down") {
            v["tags"].each { tv ->
                if (tv["Key"] == "slot") {
                    inst_for_check.add(vm_list[k])
                }
            }
        }
    }

    def sendStatus = inst_for_check.each { inst ->
        println "Healthy Checking " + inst["private_ip_address"] + " " + inst["instance_id"]
        def ssh_username = worker_conf.remote_user()
        def ssh_hostname = "${inst["private_ip_address"]}"
        def ssh_known_hosts_file = System.getProperty("user.home") + "/.ssh/known_hosts"
        def ssh_known_hosts_file_my = ssh_known_hosts_file + ".obs_scailing"
        def ssh_options = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=" + ssh_known_hosts_file_my
        def docker_command = "sudo docker exec obs_worker rcobsworker status"
        // Remove out-of-date known hosts
        execute_command('cp', ssh_known_hosts_file + ' ' + ssh_known_hosts_file_my)
        execute_command('chmod', '0777 ' + ssh_known_hosts_file_my)
        execute_command('ssh-keygen', '-f ' + ssh_known_hosts_file_my + ' -R ' + ssh_hostname, verbose=true)
        def ret_val = execute_command("ssh", ssh_options + " -X " + ssh_username + "@" + ssh_hostname + " " + docker_command,
                        verbose=false, return_stdout=true)
        if (!ret_val.contains("   Active: active (running)") || !ret_val.contains("Checking for obsworker: ..running")) {
            println 'HEALTHY:FAIL for ' + ssh_hostname + '\nShutdown the VM right now!'
            terminate_aws_ec2_instances([inst["instance_id"]])
            //println 'HEALTHY:FAIL for ' + ssh_hostname + '\nRestarting docker...'
            //docker_command = "sudo docker-compose -f /root/docker-compose.yml up -d"
            //execute_command("ssh", ssh_options + " -X " + ssh_username + "@" + ssh_hostname + " " + docker_command,
            //                verbose=false, return_stdout=true)
        }
    }

}


//TODO: FIXME:
e = { filepath ->
    evaluate(new File(System.getenv("JENKINS_HOME") + "/init.groovy.d/" + filepath))
}
create_slave_node = e("Module_Node")

def create_slaves(worker_conf, num_nodes_to_create) {

    def jsonSlurper = new JsonSlurper()
    ec2s = create_aws_ec2_instances(ami_id=worker_conf.ami_id(),
                                    min_count=1,
                                    max_count=num_nodes_to_create,
                                    keyname=worker_conf.keyname(),
                                    security_group_ids=worker_conf.security_groups(),
                                    instance_type=worker_conf.instance_type(),
                                    placement=worker_conf.placement(),
                                    subnet_id=worker_conf.subnet_id(),
                                    ebsoptimized='False',
                                    tags=worker_conf.tags(),
                                    userdata=worker_conf.userdata())

    Date date = new Date()
    String curr_date = date.format("yyyyMMdd.HHmmss")

    println ec2s

    def buildEnv = build.getEnvironment(listener)

    // Jenkins slave nodes only for imager...
//    if (worker_conf.getClass().getName() == "WorkerConf_JENKINS_IMAGER") {
    println "Slave attach needed: " + worker_conf.slave_attach_needed()
    if (worker_conf.slave_attach_needed() == true) {
        print "Creating slaves for " + worker_conf.getClass().getName()
        ec2s.each{ k, v ->
            println "Creating jenkins node ${worker_conf.name_prefix()} (${v.instance_id}) at " + curr_date
            println "  Instance ID: ${v.instance_id}"
            println "  IP Address: ${v.private_ip_address}"
            println "  Launch Time: ${v.launch_time}"
            myslot = ""
            v.tags.each{ nn ->
                if( nn['Key'] == 'slot' ) {
                    myslot = String.format("%02d", nn['Value'].toInteger())
                }
            }
            create_slave_node(
                instance = Jenkins.getInstance(),
                name = "${worker_conf.name_prefix()} (${v.instance_id})",
                remoteFS = "${worker_conf.remote_fs()}",
                numExecutors = "${worker_conf.executors_per_slave()}",
                //labelString = "${worker_conf.labels()} ${worker_conf.labels()}-slot-${myslot}",
                labelString = "${worker_conf.labels()}-slot-${myslot}",
                sshHost = "${v.private_ip_address}",
                sshPort = "${worker_conf.ssh_port()}",
                sshCredentials = "${worker_conf.credential_id()}",
                userId = "${worker_conf.remote_user()}",
                description = "${v.launch_time}"
            )
            instance.save()
            // TODO: VM PRE CHECK
            println "  VM TESTER for ${myslot}"
            if ( false ) {
                File file = new File(buildEnv["WORKSPACE"] + "/IMAGER_VM_PRE_CHECK_SLOT_${myslot}.env")
                file.write "BACKEND_SELECTION=${worker_conf.labels()}-slot-${myslot}"
            } else {
                def jenkins_url = buildEnv["JENKINS_URL_INTERNAL"]
                def jenkins_job = buildEnv["VM_CHECKER_JOB_NAME"]
                sendPostRequest("${jenkins_url}/job/${jenkins_job}/buildWithParameters", "BACKEND_SELECTION=${worker_conf.labels()}-slot-${myslot}",
                                buildEnv["JENKINS_USER"], buildEnv["JENKINS_PW"])

            }

        }
    } else {
        print "Skip creating slaves for " + worker_conf.getClass().getName()
    }
    println ""
    return ec2s
}

Integer how_many_slaves_to_create(worker_conf, num_requested_executors) {

    println worker_conf.getClass(); 
    println (worker_conf instanceof WorkerConf_JENKINS_IMAGER);
    println (worker_conf instanceof WorkerConf_OBS_WORKER_NORMAL);
    println (worker_conf instanceof WorkerConf_OBS_WORKER_POWER);

    // Check slave nodes are free enough to allocate image jobs
    def slave_stat = new SlaveStatus(worker_conf.name_prefix())
    println "\nSLAVE STATUS (${worker_conf.name_prefix()}"
    println "  NUMBER_OF_NODES: ${slave_stat.number_of_nodes()}"
    println "  TOTAL EXECUTORS: ${slave_stat.total_executors()}"
    println "  IDLE EXECUTORS: ${slave_stat.idle_executors()}"
    println ""
    slave_stat.get_slave_info().each{ k, v ->
        println "    ${k}:${v}"
    }
    println ""
    def Integer final_nodes_to_create = 0
    if (slave_stat.number_of_nodes() >= worker_conf.max_slaves() ) {
        println "REACHED MAXIMUM SLAVES"
        return 0
    }
    if (slave_stat.idle_executors() >= num_requested_executors) {
        println "HAVE ENOUGH EXECUTORS"
        return 0
    }

    println "CALCULATING..."
    def Integer min_to_create = \
        Math.min(num_requested_executors - slave_stat.idle_executors(), \
                  (worker_conf.max_slaves() - slave_stat.number_of_nodes()) \
                  * worker_conf.executors_per_slave())
    def Double nodes_to_create = min_to_create.div(worker_conf.executors_per_slave())
    final_nodes_to_create = Math.min(Math.ceil(nodes_to_create).intValue(), \
                             worker_conf.max_slaves() - slave_stat.number_of_nodes())
    println "MIN(" + num_requested_executors + "-" + slave_stat.idle_executors() \
            + ", (" + worker_conf.max_slaves() + "-" + slave_stat.number_of_nodes() \
            + ")*" + worker_conf.executors_per_slave() + ")" \
            + " = " + min_to_create
    println "FINAL NUMBER OF SLAVES TO CREATE: ${final_nodes_to_create} \n"
    return final_nodes_to_create
}

Integer how_many_slaves_to_create_wo_nodes(worker_conf, num_requested_executors, current_running_slaves_num) {

    println "CALCULATING..."
    def Integer min_to_create = \
        Math.min(num_requested_executors, \
                  (worker_conf.max_slaves() - current_running_slaves_num) \
                  * worker_conf.executors_per_slave())
    def Double nodes_to_create = min_to_create.div(worker_conf.executors_per_slave())
    final_nodes_to_create = Math.ceil(nodes_to_create).intValue()
    println "MIN(" + num_requested_executors \
            + ", (" + worker_conf.max_slaves() + "-" + current_running_slaves_num \
            + ")*" + worker_conf.executors_per_slave() + ")" \
            + " = " + min_to_create
    println "FINAL NUMBER OF SLAVES TO CREATE: ${final_nodes_to_create} \n"
    return final_nodes_to_create
}

def worker_ondemand_create_request(worker_conf, Integer num_requested_executors) {

    println "YOU REQUESTED ${num_requested_executors} executors!"

    // Find empty slot from 01 ~ 99
    def free_slots = []
    def allocated_slots = []
    for (i = worker_conf.instance_base(); i <= worker_conf.instance_base()+worker_conf.max_slaves(); i++) { free_slots.add(String.format("%02d", i)) }
    def current_aws_status_list = get_aws_status(worker_conf)
    current_aws_status_list.each { k, v ->
        println k
        println v['tags']
        if (v['state'] == 'terminated' || v['state'] == 'shutting-down') {
            return
        }
        v['tags'].each { n -> 
            if (n['Key'] == 'slot') {
                used_slot_number = n['Value']
                println "Removing unused $used_slot_number from allocated pool."
                free_slots.remove(used_slot_number)
                allocated_slots.add(used_slot_number)
            }
        }
    }
    free_slots = free_slots.unique { a, b -> a <=> b }
    println "Free slots(" + free_slots.size() + "): " + free_slots
    println "Allocated slots(" + allocated_slots.size() + "): " + allocated_slots

    // Check slave nodes are free enough to allocate image jobs
    def num_nodes_to_create = 0
    if (worker_conf.slave_attach_needed() == true) {
        num_nodes_to_create = how_many_slaves_to_create(worker_conf, num_requested_executors)
    } else {
        num_nodes_to_create = how_many_slaves_to_create_wo_nodes(worker_conf, num_requested_executors, allocated_slots.size())
    }

    assert (num_nodes_to_create < 50)

    if (num_nodes_to_create > 0) {
        def created_inst = []
        for (f = 0; f < num_nodes_to_create; f++) {
            println "Processing ${free_slots[f]}"
            worker_conf.set_slot_number(free_slots[f])
            //println worker_conf.userdata()
            inst_info = create_slaves(worker_conf, 1)
            created_inst.add(inst_info)
        }
        println "\"TitleDisplay\": \"Create(${created_inst.size()}/${num_requested_executors})\""
    }

    // Make VMs online if its not ready
    check_healthy_status(worker_conf, current_aws_status_list)
}

def worker_ondemand_revoke_request_wo_nodes(worker_conf) {

    def requested_slot_numbers = build.getEnvironment(listener)["SLOT_NUMBERS"].split(',')
    println "requested_slot_numbers: " + requested_slot_numbers
    def existing_instance_ids = []
    def inst_for_terminate = []

    def current_aws_status = get_aws_status(worker_conf)
    current_aws_status.each { k, v ->
        if (v["state"] != "terminated" && v["state"] != "shutting-down") {
            existing_instance_ids.add(k)
            println "+ Existing " + k
            v["tags"].each { tv ->
                println " * Check " + tv["Key"] + " " + tv["Value"]
                if (tv["Key"] == "slot" && tv["Value"] in requested_slot_numbers) {
                    println " *** Term " + tv["Key"] + " " + tv["Value"]
                    inst_for_terminate.add(current_aws_status[k])
                }
            }
        }
    }

    println "\nEXISTING INSTANCE INFO:\n${existing_instance_ids}\n"
    println "\nINSTANCES TO TERMINATE:\n${inst_for_terminate}\n"

    def threads = []
    def inst_ids_to_terminate = []
    def sendStop = inst_for_terminate.each { inst ->
        inst_ids_to_terminate.add(inst["instance_id"])
        def th = new Thread({
            println "Terminating " + inst["private_ip_address"] + " " + inst["instance_id"]
            def ssh_username = worker_conf.remote_user()
            def ssh_hostname = "${inst["private_ip_address"]}"
            def ssh_known_hosts_file = System.getProperty("user.home") + "/.ssh/known_hosts"
            def ssh_known_hosts_file_my = ssh_known_hosts_file + ".obs_scailing"
            def ssh_options = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=" + ssh_known_hosts_file_my
            def docker_command = "sudo docker exec obs_worker rcobsworker stop"
            // Remove out-of-date known hosts
            execute_command('cp', ssh_known_hosts_file + ' ' + ssh_known_hosts_file_my)
            execute_command('chmod', '0777 ' + ssh_known_hosts_file_my)
            execute_command('ssh-keygen', '-f ' + ssh_known_hosts_file_my + ' -R ' + ssh_hostname, verbose=true)
            def ret_val = execute_command("ssh", ssh_options + " -X " + ssh_username + "@" + ssh_hostname + " " + docker_command,
                            verbose=true)
            if (ret_val == "FATAL") {
                inst_ids_to_terminate.remove(inst["instance_id"])
                println "Fail to stop service " + inst["instance_id"]
            }
        })
        println "Adding thread " + inst["instance_id"]
        threads << th
    }
    threads.each { it -> it.start() }
    threads.each { it -> it.join() }
    // Wait 20 seconds to complete its execution
    //sleep(20000)

    println inst_ids_to_terminate
    if (inst_ids_to_terminate) {
        terminate_aws_ec2_instances(inst_ids_to_terminate)
        println "\"TitleDisplay\": \"Delete(${inst_ids_to_terminate.size()})\""
    }
}

def worker_ondemand_revoke_request(worker_conf) {

    if (worker_conf.slave_attach_needed() != true) {
        return worker_ondemand_revoke_request_wo_nodes(worker_conf)
    }

    def bRevoke = true

    // Check slave nodes for running builds
    def slave_stat = new SlaveStatus(worker_conf.name_prefix())
    def running_build_num = slave_stat.busy_executors()
    if (running_build_num > 0) {
        println "\nYOU HAVE BUSY EXECUTORS."
        bRevoke = false
    }

    // Check build queue
    def pending_build_num = 0
    Jenkins.instance.queue.items.find { it ->
        if (it.task.name in worker_conf.check_queue_list()) {
            pending_build_num += 1
            //println "\nPENDING BUILD IN THE QUEUE (${it.task.name}:${it.id}). STOP REVOKE OPERATION."
            bRevoke = false
        }
    }

    if (running_build_num > 0 || pending_build_num > 0) {
        Date date = new Date()
        String curr_date = date.format("yyyyMMdd.HHmmss")
        println "\n[" + curr_date + "] You have " + running_build_num + " running builds and " + pending_build_num + " pending builds"
        def fname = System.getenv("JENKINS_HOME") + '/.imager_history.log'
        //f = new File(System.getenv("JENKINS_HOME"), '/.imager_history.log')
        f = new File(fname)
        f.append(curr_date + "," + running_build_num + "," + pending_build_num + "\n")
        def sync_base = build.getEnvironment(listener)['IMG_SYNC_DEST_BASE']
        def cmd = "rsync -av --delay-updates " + fname + " " + sync_base + "/snapshots/.dashboard/imager_trend/imager_history.csv"
        execute_command(cmd, " ")
    }

    // Check build
    worker_conf.check_queue_list().find { it ->
        def item = hudson.model.Hudson.instance.getItem(it)
        if (item.isInQueue() || item.isBuilding()) {
            println "YOU HAVE QUEUED OR BUILDING ITEMS FOR ${it}"
            bRevoke = false
            return true
        }
    }

    if (bRevoke != true) { return }

    existing_instance_ids = []
    get_aws_status(worker_conf).each { k, v ->
        if (v["state"] != "terminated" && v["state"] != "shutting-down") {
            existing_instance_ids.add(k)
        }
    }
    println "\nEXISTING INSTANCE INFO:\n${existing_instance_ids}\n"

    Date date = new Date()
    String curr_date = date.format("yyyyMMdd.HHmmss")

    def existing_slave_ids = []
    def inst_for_terminate = []

    slave_stat.get_slave_info().each { k, v ->
        def matcher = v["name"] =~ /^${worker_conf.name_prefix()} \((i-[0-9a-f]+)\)$/
        if (matcher?.matches()) {
            def inst_id = matcher.group(1)
            existing_slave_ids.add(inst_id)
            if ( v["idle_since"].toLong().div(1000*60) >= worker_conf.idle_termination_minutes()) {
                println ".... Delete candidate ${k} : " + v["idle_since"].toLong().div(1000*60)
                // Mark the node as offline and disconnect
                slave_stat.disconnect_node(v["name"])
                if (inst_id in existing_instance_ids) {
                    inst_for_terminate.add(inst_id)
                }
            } else {
                println ".... TIMER NOT EXPIRED FOR ${k} : " + v["idle_since"].toLong().div(1000*60)
            }
        } else {
            println "SLAVE \"${k}\" IS NOT MY CHILD."
        }
    }
    println "\nEXISTING NODE INFO:\n${existing_slave_ids}\n"

    println inst_for_terminate
    if (inst_for_terminate) {
        terminate_aws_ec2_instances(inst_for_terminate)
        println "\"TitleDisplay\": \"Delete(${inst_for_terminate.size()})\""
    }

    slave_stat.get_slave_info().each { k, v ->
        def matcher = v["name"] =~ /^${worker_conf.name_prefix()} \((i-[0-9a-f]+)\)$/
        if (matcher?.matches()) {
            def inst_id = matcher.group(1)
            // Delete the node
            if (inst_id in inst_for_terminate) {
                println "Deleting jenkins node ${k} at " + curr_date \
                    + " Cause: " + v["idle_since"].toLong().div(1000*60)
                slave_stat.delete_node(v["name"])
            }
        }
    }

    dangled_jenkins_nodes = existing_slave_ids.toSet() - existing_instance_ids.toSet()
    if (dangled_jenkins_nodes) {
        println "\nDangled jenkins nodes:\n${dangled_jenkins_nodes}"
        for (djn in dangled_jenkins_nodes) {
            slave_stat.get_slave_info().each { k, v ->
                def matcher = v["name"] =~ /^${worker_conf.name_prefix()} \(${djn}\)$/
                if (matcher?.matches()) {
                    println 'Delete dangled node:' + v["name"]
                    slave_stat.delete_node(v["name"])
                }
            }
        }
    }

    dangled_ec2_instances = existing_instance_ids.toSet() - existing_slave_ids.toSet()
    if (dangled_ec2_instances) {
        println "\nDangled EC2 instances:\n${dangled_ec2_instances}"
        println "\nFORCE TERMINATE!"
        terminate_aws_ec2_instances(dangled_ec2_instances)
    }
}

def __main__() {

    def buildEnv = build.getEnvironment(listener)

    if (buildEnv["ONDEMAND_SLAVE_CONFIGURATION_ENABLED"] != "1") {
        println "FEATURE NOT ENABLED"
        return
    }

    action = buildEnv['ACTION']
    purpose = buildEnv['PURPOSE']
    requested_num_executors = buildEnv['REQUESTED_NUM_EXECUTORS']

    assert (buildEnv['PURPOSE'] == 'JENKINS_IMAGER') \
        || (buildEnv['PURPOSE'] == 'OBS_WORKER_NORMAL') \
        || (buildEnv['PURPOSE'] == 'OBS_WORKER_POWER')

    def worker_conf = get_worker_conf(purpose)

    if (action == "REQUEST_WORKER") {
        worker_ondemand_create_request(worker_conf, \
                                     requested_num_executors.toInteger())
    } else if(action == "REVOKE_WORKER") {
        worker_ondemand_revoke_request(worker_conf)
    } else {
        println "Invalid action: " + action
    }

}

__main__()

return 0

