#!/usr/bin/ruby

require 'rbvmomi'
require 'redis'
require 'yaml'

$:.unshift(File.dirname(__FILE__))
require 'lib/logger'
require 'lib/require_relative'
require 'lib/vsphere_helper'

$logger = Logger.new
$vsphere_helper = VsphereHelper.new

Dir.chdir(File.dirname(__FILE__))

# Load the configuration file
config_file = File.expand_path('vmware-host-pooler.yaml')
config = YAML.load_file(config_file)

pools   = config[:pools]
vsphere = config[:vsphere]

# Connect to Redis
$redis = Redis.new



# Clone a VM
def clone_vm template, pool, folder, datastore
  Thread.new {
    vm = {}

    if template =~ /\//
      templatefolders = template.split('/')
      vm['template'] = templatefolders.pop
    end

    if templatefolders
      vm[vm['template']] = $vsphere_helper.find_folder(templatefolders.join('/')).find(vm['template'])
    else
      raise "Please provide a full path to the template"
    end

    if vm['template'].length == 0
      raise "Unable to find template '#{vm['template']}'!"
    end

    # Generate a randomized hostname
    o = [('a'..'z'),('0'..'9')].map{|r| r.to_a}.flatten
    vm['hostname'] = o[rand(25)]+(0...14).map{o[rand(o.length)]}.join

    # Annotate with creation time, origin template, etc.
    configSpec = RbVmomi::VIM.VirtualMachineConfigSpec(
      :annotation =>
        'Base template: ' + vm['template'] + "\n" +
        'Creation time: ' + Time.now.strftime("%Y-%m-%d %H:%M")
    )

    # Put the VM in the specified folder and resource pool
    relocateSpec = RbVmomi::VIM.VirtualMachineRelocateSpec(
      :datastore    => $vsphere_helper.find_datastore(datastore),
      :pool         => $vsphere_helper.find_pool(pool),
      :diskMoveType => :moveChildMostDiskBacking
    )

    # Create a clone spec
    spec = RbVmomi::VIM.VirtualMachineCloneSpec(
      :location      => relocateSpec,
      :config        => configSpec,
      :powerOn       => true,
      :template      => false
    )

    # Clone the VM
    $logger.log('d', "[ ] '#{vm['hostname']}' is being cloned from '#{vm['template']}'")

    start = Time.now
    vm[vm['template']].CloneVM_Task(
      :folder => $vsphere_helper.find_folder(folder),
      :name => vm['hostname'],
      :spec => spec
    ).wait_for_completion
    finish = '%.2f' % (Time.now-start)

    # Add VM to Redis inventory ('pending' pool)
    $redis.sadd('vmware_host_pool__pending__'+vm['template'], vm['hostname'])

    $logger.log('s', "[+] '#{vm['hostname']}' cloned from '#{vm['template']}' in #{finish} seconds")

    # Metrics
    $redis.lpush('vmware_host_pool_metrics__deploy', finish)
    $redis.ltrim('vmware_host_pool_metrics__deploy', 0, 100)
  }
end

# Check the state of a VM
def check_vm vm, pool
  Thread.new {
    if (
      ($vsphere_helper.find_vms(vm)[vm]) and
      ($vsphere_helper.find_vms(vm)[vm].summary.guest.toolsRunningStatus == 'guestToolsRunning') and
      ($vsphere_helper.find_vms(vm)[vm].summary.guest.hostName) and
      ($vsphere_helper.find_vms(vm)[vm].summary.guest.hostName == vm) and
      ($vsphere_helper.find_vms(vm)[vm].summary.guest.ipAddress != nil)
    )
      begin
        Socket.getaddrinfo(vm, nil)
      rescue
        if (
          ($vsphere_helper.find_vms(vm)[vm].runtime) and
          ($vsphere_helper.find_vms(vm)[vm].runtime.bootTime) and
          (((( Time.now - $vsphere_helper.find_mvs(vm)[vm].runtime.bootTime ) / 60 ) / 60 ) >= 1)
        )
          $redis.srem('vmware_host_pool__pending__'+pool, vm)
          $redis.sadd('vmware_host_pool__failed__'+pool, vm)

          # Metrics
          $redis.lpush('vmware_host_pool_metrics__deploy_fail', '1')

          $logger.log('s', "[<] '#{vm}' moved to 'failed' queue")
        end

        next
      end

      $redis.sadd('vmware_host_pool__ready__'+pool, vm)
      $redis.srem('vmware_host_pool__pending__'+pool, vm)

      $logger.log('s', "[>] '#{vm}' moved to 'ready' queue")

      # Metrics
      $redis.lpush('vmware_host_pool_metrics__deploy_fail', '0')
    end
  }
end

# Destroy a VM
def destroy_vm vm
  Thread.new {
    host = $vsphere_helper.find_vms(vm)[vm]

    if (
      (host) and
      (host.runtime)
    )
      start = Time.now

      if host.runtime.powerState == 'poweredOn'
        $logger.log('d', "[ ] '#{vm}' is being shut down")
        host.PowerOffVM_Task.wait_for_completion
      end

      host.Destroy_Task.wait_for_completion
      finish = '%.2f' % (Time.now-start)

      $logger.log('s', "[-] '#{vm}' destroyed in #{finish} seconds")
    end
  }
end



# Promotion loop
Thread.new {
loop do
  pools.each do |pool|
    inventory = {}

    # Locate the resource pool
    base = $vsphere_helper.find_pool(pool['pool'])

    # Make sure all VMs in resource pool are accounted-for
    base.vm.each do |vm|
      if (
        (! $redis.sismember('vmware_host_pool__pending__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__ready__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__running__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__completed__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__failed__'+pool['name'], vm['name']))
      )
        $redis.sadd('vmware_host_pool__pending__'+pool['name'], vm['name'])
      end

      inventory[vm['name']] = 1
    end

    Thread.new {
      $redis.smembers('vmware_host_pool__pending__'+pool['name']).each do |vm|
        if ! inventory[vm]
          $redis.srem('vmware_host_pool__pending__'+pool['name'], vm)
        end

        Thread.new {
          check_vm(
            vm,
            pool['name']
          )
        }
      end
    }
  end
end
}

# Cleanup loop
Thread.new {
loop do
  pools.each do |pool|
    inventory = {}

    # Locate the resource pool
    base = $vsphere_helper.find_pool(pool['pool'])

    # Make sure all VMs in resource pool are accounted-for
    base.vm.each do |vm|
      if (
        (! $redis.sismember('vmware_host_pool__pending__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__ready__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__running__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__completed__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__failed__'+pool['name'], vm['name']))
      )
        $redis.sadd('vmware_host_pool__pending__'+pool['name'], vm['name'])
      end

      inventory[vm['name']] = 1
    end

    Thread.new {
      $redis.smembers('vmware_host_pool__completed__'+pool['name']).each do |vm|
        if ! inventory[vm]
          $redis.srem('vmware_host_pool__completed__'+pool['name'], vm)
        end

        Thread.new {
          destroy_vm(vm)
        }
      end
    }

  end
end
}

# Update loop
loop do
  pools.each do |pool|
    inventory = {}

    # Locate the resource pool
    base = $vsphere_helper.find_pool(pool['pool'])

    # Make sure all VMs in resource pool are accounted-for
    base.vm.each do |vm|
      if (
        (! $redis.sismember('vmware_host_pool__pending__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__ready__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__running__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__completed__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__failed__'+pool['name'], vm['name']))
      )
        $redis.sadd('vmware_host_pool__pending__'+pool['name'], vm['name'])
      end

      inventory[vm['name']] = 1
    end

    total = $redis.scard('vmware_host_pool__ready__'+pool['name']) +
            $redis.scard('vmware_host_pool__pending__'+pool['name'])

    # Check 'ready' pool
    Thread.new {
      $redis.smembers('vmware_host_pool__ready__'+pool['name']).each do |vm|
        if ! inventory[vm]
          $redis.srem('vmware_host_pool__ready__'+pool['name'], vm)
        end

        Thread.new {
          if (
            (! $vsphere_helper.find_vms(vm)[vm]) or
            ($vsphere_helper.find_vms(vm)[vm].summary.guest.toolsRunningStatus != 'guestToolsRunning') or
            ($vsphere_helper.find_vms(vm)[vm].summary.guest.hostName != vm) or
            ($vsphere_helper.find_vms(vm)[vm].summary.guest.ipAddress == nil)
          )
            $redis.srem('vmware_host_pool__ready__'+pool['name'], vm)
            $redis.sadd('vmware_host_pool__failed__'+pool['name'], vm)

            $logger.log('s', "[<] '#{vm}' moved to 'failed' queue")

            next
          end

          begin
            Socket.getaddrinfo(vm, nil)
          rescue
            $redis.srem('vmware_host_pool__ready__'+pool['name'], vm)
            $redis.sadd('vmware_host_pool__failed__'+pool['name'], vm)

            $logger.log('s', "[<] '#{vm}' moved to 'failed' queue")

            # Metrics
            $redis.lpush('vmware_host_pool_metrics__deploy_fail', '1')
          end
        }
      end
    }

    # Check 'failed' pool
    Thread.new {
      $redis.smembers('vmware_host_pool__failed__'+pool['name']).each do |vm|
        if ! inventory[vm]
          $redis.srem('vmware_host_pool__failed__'+pool['name'], vm)
        end

        if (
          ($vsphere_helper.find_vms(vm)[vm]) and
          ($vsphere_helper.find_vms(vm)[vm].summary.guest.toolsRunningStatus == 'guestToolsRunning') and
          ($vsphere_helper.find_vms(vm)[vm].summary.guest.hostName == vm) and
          ($vsphere_helper.find_vms(vm)[vm].summary.guest.ipAddress != nil)
        )
          begin
            Socket.getaddrinfo(vm, nil)
          rescue
            next
          end

          $redis.sadd('vmware_host_pool__ready__'+pool['name'], vm)
          $redis.srem('vmware_host_pool__failed__'+pool['name'], vm)

          $logger.log('s', "[>] '#{vm}' moved to 'ready' queue")
        else
          Thread.new {
            destroy_vm(vm)
          }
        end
      end
    }

    # Bring the pool up to the desired size
    Thread.new {
      if total < pool['size']
        # Provision VMs
        (1..(pool['size']-total)).each { |i|
          Thread.new {
            clone_vm(
              pool['template'],
              pool['pool'],
              pool['folder'],
              pool['datastore']
            )
          }
        }
      end
    }

    # Metrics
    $redis.ltrim('vmware_host_pool_metrics__deploy_fail', 0, 100)
  end
end