#!/usr/bin/env python

import commands
import sys
import time
import datetime
import random
from euca2ools import Euca2ool

euca_conn = None

def instances(test_host=None):
    rval = []
    for r in euca_conn.get_all_instances([]):
        for i in r.instances:
            i.host = i.key_name.split(',')[-1].split(')')[0].strip()
            if test_host and i.host != test_host:
                continue
            rval.append(i)
    return rval


max_instance_count = 12
max_launch_count = 4
error_launch_time = 60

launch_times = {}
def actually_running(instance):
    if not instance.private_dns_name:
        return False
    id_url = "latest/meta-data/instance-id"
    options = "-s --connect-timeout 0.1"
    command = "curl %s %s/%s" % (options, instance.private_dns_name, id_url)
    instance_id = commands.getoutput(command).strip()
    if not instance_id:
        return False
    if instance_id != instance.id:
        print "incorrect (%s) instance id for %s" % (instance_id, instance.id)
        return False
    return True

def mysql(query):
    print '\n'.join(commands.getoutput('mysql nova -e "%s"' % query).split("\n")[1:])

def print_stats(host, start, succeeded, failed):
    ran = succeeded + failed
    end = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S")
    print """%s STATISTICS
  Compute
  -------
  Total:                %s
  Successes:            %s
  Failures:             %s
  Started:              %s
  Finished:             %s""" % (host.upper(), ran, succeeded, failed, start, end)
    where = " where host='%s' and created_at > '%s' " % (host, start)
    mysql("select '  MySQL Successes:' as label, count(*) from instances" +
          where + "and launched_at is not null")
    diff = "AVG(TIME_TO_SEC(TIMEDIFF(launched_at, scheduled_at)))"
    mysql("select '  MySQL Avg Time:' as label, %s from instances" % diff +
          where + "and launched_at is not null")
    mysql("select '  MySQL Failures:' as label, count(*) from instances" +
          where + "and launched_at is null and deleted_at is not null")


def main(test_host, num_to_run):
    try:
        euca = Euca2ool()
        global euca_conn
        euca_conn = euca.make_connection()
    except:
        print "environment variables not configured"
        sys.exit(1)

    user_data = open('proxy.sh').read()
    start = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S")
    total_succeeded = 0
    total_failed = 0
    try:
        while total_succeeded + total_failed < num_to_run:
            insts = instances(test_host)

            # terminate running instances and check on launched instances
            for i in insts:
                if i.state == 'shutdown':
                    if launch_times.has_key(i.id):
                        print "%s,%s,%s, FAIL" % (launch_times[i.id], i.id, i.host)
                        del launch_times[i.id]
                        euca_conn.terminate_instances([i.id])
                        total_failed = total_failed + 1
                if i.state == 'running' and actually_running(i):
                    if launch_times.has_key(i.id):
                        print '%s,%s,%s, %d' % (launch_times[i.id], i.id, i.host, time.time()-launch_times[i.id])
                        total_succeeded = total_succeeded + 1
                        del launch_times[i.id]
                    euca_conn.terminate_instances([i.id])

            # launching new instances
            launch_count = min(random.randint(0, max_instance_count-len(insts)), max_launch_count)
            if launch_count > 0:
                if test_host:
                    placement = 'test:%s' % test_host
                else:
                    placement = None
                r = euca_conn.run_instances(image_id="ami-tty",
                                            min_count=launch_count,
                                            placement=placement,
                                            user_data=user_data,
                                            instance_type="m1.tiny",
                                            max_count=launch_count)
                assert(len(r.instances) == launch_count)
                for i in r.instances:
                    launch_times[i.id] = time.time()

            time.sleep(1)
    except:
        pass
    print_stats(test_host, start, total_succeeded, total_failed)

if __name__ == "__main__":
    if len(sys.argv) > 1:
        test_host = sys.argv[1]
    else:
        print "please specify a host"
        exit(1)
    num_to_run = 1000
    if len(sys.argv) > 2:
        num_to_run = sys.argv[2]
    main(test_host, int(num_to_run))
