#!/usr/bin/env python

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

euca_conn = None

def volumes(test_host):
    rval = []
    for v in euca_conn.get_all_volumes([]):
        v.host = v.status.split(',')[1].strip()
        if test_host and v.host != test_host:
            continue
        rval.append(v)
    return rval


max_volume_count = 50
max_launch_count = 4
volume_size = 10

launch_times = {}

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
  Volume
  ------
  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 volumes" +
          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 volumes" % diff +
          where + "and launched_at is not null")
    mysql("select '  MySQL Failures:' as label, count(*) from volumes" +
          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)

    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:
            vols = volumes(test_host)

            # terminate available volumes and check on launched volumes
            for v in vols:
                status = v.status.split(' ')[0]
                if status == 'available':
                    if launch_times.has_key(v.id):
                        print '%s,%s,%s, %d' % (launch_times[v.id], v.id, v.host, time.time()-launch_times[v.id])
                        del launch_times[v.id]
                    euca_conn.delete_volume(v.id)
                    total_succeeded = total_succeeded + 1
                # TODO(vish): verify volumes and check for failures

            # launching new volumes
            launch_count = min(random.randint(0, max_volume_count-len(vols)), max_launch_count)
            if launch_count > 0:
                for x in range(launch_count):
                    if test_host:
                        zone = 'test:%s' % test_host
                    else:
                        zone = None
                    v = euca_conn.create_volume(size=volume_size, zone=zone)
                    launch_times[v.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))
