#!/usr/bin/env python
from __future__ import print_function
import sys
import time
import boto

#get all the instances of a layer and run some recipes on each
def update_instances(elb_name, stack, layer, min_instances, recipes):
  #connect to opsworks
  opsworks = boto.connect_opsworks()

  #grab all the instances for that layer
  instances = opsworks.describe_instances(layer_id=layer)

  #TODO: make this more robust of a heuristic
  #TODO: sort by instance name so that we always try the same one first
  #if we dont have enough machines to feel safe we give up
  instances = [ i for i in instances['Instances'] if i.get('Status') == 'online' ]
  if len(instances) < min_instances:
    raise Exception('Not enough instances (%d < %d) in layer %s to risk the update' % (len(instances), min_instances, layer))

  #connect to the elb and check which are in it
  elb = boto.connect_elb(elb_name)
  instances = [ i for i in instances if elb.describe_instance_health(elb_name, [i['Ec2InstanceId']])[0].state == 'InService' ]
  if len(instances) < min_instances:
    raise Exception('Not enough instances (%d < %d) in elb %s to risk the update' % (len(instances), min_instances, layer))

  #for each one
  for instance in instances:
    #take this instance out of the elb
    registered = [ i.id for i in elb.deregister_instances(elb_name, [instance['Ec2InstanceId']]) ]
    if instance['Ec2InstanceId'] in registered:
      raise Exception('Instance was still registered with the elb')

    #check its status
    wait = 60
    while wait > 0:
      if elb.describe_instance_health(elb_name, [instance['Ec2InstanceId']])[0].state == 'OutOfService':
        break
      wait -= 5
      time.sleep(5)
    if wait < 1:
      elb.register_instances(elb_name, [instance['Ec2InstanceId']])
      raise Exception('Deregistered instance was not OutOfService')

    #deployment task of running a recipe
    deployment_id = opsworks.create_deployment(stack, {'Name': 'execute_recipes', 'Args': {'recipes': recipes}}, instance_ids=[instance['InstanceId']])['DeploymentId']
    #wait until it ends
    while True:
      if opsworks.describe_commands(deployment_id=deployment_id)['Commands'][0]['Status'] != 'pending':
        break
      time.sleep(5)
    #check that its status is all good
    if opsworks.describe_commands(deployment_id=deployment_id)['Commands'][0]['Status'] != 'successful':
      elb.register_instances(elb_name, [instance['Ec2InstanceId']])
      print('Deployment %s to Instance %s failed' % (deployment_id, instance['InstanceId']), file=sys.stderr) 
      sys.exit(1)

    #put the instance back in the elb
    registered = [ i.id for i in elb.register_instances(elb_name, [instance['Ec2InstanceId']]) ]
    if instance['Ec2InstanceId'] not in registered:
       raise Exception('Instance was not registered with the elb')

    #check its status
    wait = 60
    while wait > 0:
      if elb.describe_instance_health(elb_name, [instance['Ec2InstanceId']])[0].state == 'InService':
        break
      wait -= 5
      time.sleep(5)
    if wait < 1:
      raise Exception('Re-registered instance was not InService')
    

#entry point for script
if __name__ == "__main__":
  #update the service instances
  elbs = '<%= node[:valhalla][:routing_service_elbs] %>'.split(',')
  layers = '<%= node[:valhalla][:routing_service_layers] %>'.split(',')
  min_layers_instances = '<%= node[:valhalla][:min_layers_instances] %>'.split(',')
  recipes = '<%= node[:valhalla][:routing_service_recipes] %>'.split(',')
  for elb, layer, min_instances in zip(elbs, layers, min_layers_instances):
    update_instances(elb, '<%= node[:valhalla][:routing_service_stack] %>', layer, int(min_instances), recipes)
