#!/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')

pools   = YAML.load_file(config_file)[:pools]
vsphere = YAML.load_file(config_file)[:vsphere]

# Connect to vSphere
$vim = RbVmomi::VIM.connect(
  :host     => vsphere['server'],
  :user     => vsphere['username'],
  :password => vsphere['password'],
  :ssl      => true,
  :insecure => true,
  :rev      => '5.1'
)

# Connect to Redis
$redis = Redis.new

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

    # Locate the resource pool
    datacenter = $vim.serviceInstance.find_datacenter
    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__ready__'+pool['name'], vm['name'])) and
        (! $redis.sismember('vmware_host_pool__pending__'+pool['name'], vm['name']))
      )
        $redis.sadd('vmware_host_pool__pending__'+pool['name'], vm['name'])
      end

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

    # Check 'pending' pool
    $redis.smembers('vmware_host_pool__pending__'+pool['name']).each do |vm|
      if ! inventory[vm]
        $redis.srem('vmware_host_pool__pending__'+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.ipAddress != nil)
      )
        begin
          Socket.getaddrinfo(vm, nil)
        rescue
          break
        end

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

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

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

      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.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")

        break
      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")
      end
    end

    # Bring the pool up to the desired size
    if total < pool['size']

      # Provision VMs
      (1..(pool['size']-total)).each { |i|
        vm = {}

        if pool['template'] =~ /\//
          templatefolders = pool['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 '#{h['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

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

        # Create a clone spec
        spec = RbVmomi::VIM.VirtualMachineCloneSpec(
          :location      => relocateSpec,
          :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(pool['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__'+pool['name'], vm['hostname'])

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

    # ZZZzzz...
    sleep(5)
  end
end