from optparse import OptionParser
from subprocess import Popen, PIPE
import time
import sys

__all__ = [main, Workforce]

'''
Use Amazon ec2 instances to quickly run a program that takes lines from stdin.

Must be run on an ec2 controller (think ecmm.chitika.net).
Justin Pombrio, 2008
'''

parser = OptionParser()
parser.add_option('-n', '--workers', type='int', dest='workers',
                  help='Start this many ec2 instances')
parser.add_option('-e', '--executable', dest='executable',
                  help='The program to run')
parser.add_option('-a', '--auxilary', dest='auxilary',
                  help='Extra files required by the program')
parser.add_option('-s', '--syncronize', dest='syncronize',
                  help='Syncronize these files up and down')
parser.add_option('-b', '--batch_size', type='int', dest='batch_size', default=10,
                  help='Number of lines to be uploaded at once')
parser.add_option('-i', '--ami_id', dest='ami_id', default='ami-d206e3bb',
                  help='The id of the ami image to use for the workers')
parser.add_option('-q', '--quiet', action='store_const',
                  const=0, dest='verbosity') 
parser.add_option('-v', '--verbose', action='store_const',
                  const=1, dest='verbosity', default=1)
parser.add_option('-vv', '--very_verbose', action='store_const',
                  const=2, dest='verbosity')


image_cols = ['type', 'ami_id', 'name', 'key', 'status', 'ownership']

instance_cols = ['type', 'instance_id', 'ami_id', 'description', 'ip',
                 'status', '???', 'number', 'size', 'starttime']

start_cols = instance_cols

terminate_cols = ['type', 'instance_id', 'status']

auxilary_rsync_opts = '-rz' # recurse into dirs, gzip files

syncronize_rsync_opts = '-urz' # recurse into dirs, gzip files, update only
                               # (do not overwrite a newer version of a file)

command_template = 'ssh -p %s root@%s cd ~/; %s' # % (port, ip, command)

sync_up_template = 'rsync %s -p %s %s root@%s:~/%s' # % (opts, port, file, ip, file)

sync_down_template = 'rsync %s -p %s root@%s:~/%s %s' # % (opts, port, ip, file, file)

delay = 5 # seconds


def group_by(iterable, length):
    iterator = iterable.__iter__()
    group = []

    try:
        while True:
            group = []
            count = 0
            while count < length:
                group.append(iterator.next())
                count += 1
            yield group

    except StopIteration:
        if group:
            yield group

def read_ec2_output(fileobj, cols):
    result = []
    for line in fileobj:
        row_dict = {}
        row = line.rstrip('\n').split('\t')
        for col, element in zip(cols, row):
            row_dict[col] = element.strip()
        result.append(row_dict)

    return result

def run_ec2_command(args, cols):
    process = Popen(args, stdout=PIPE)
    process.wait()
    return read_ec2_output(process.stdout, cols)

def read_info(info, requirements, key):
    result = []
    for row in info:
        row_dict = {}
        for entry in row:
            requirement = requirements.get(entry)
            if requirement is not None and row[entry] != requirement:
                continue
        row_dict[row[key]] = row
        result.append(row_dict)
    return result

def get_ip(internal_ip):
    """Get the ip address out of a string such as ip-00-000-000-00.ec2.internal"""
    ip = internal_ip.split('.')[0]
    if ip.startswith('ip-'):
        ip = ip[3:]
    ip = ip.replace('-', '.')
    if len(ip) != 13:
        raise ValueError("Unable to get the ip of %s." % internal_ip)
    return ip


class Workforce(object):

    def __init__(self):
        self.port = 9999

        # Only the instances started by this program
        self.instances = {} # {instance_id : {attribute : value}}  where attributes are from instance_cols
        # Only available images
        self.images = {} # {ami_id : {attribute : value}}  where attributes are from image_cols

        self.processes = []
        self.upload = []
        self.download = []

    def update_images(self):
        image_info = run_ec2_command(['ec2-describe-images'], image_cols)
        self.images = read_info(image_info, {'type' : 'INSTANCE', 'status' : 'available'}, 'ami_id')

    def update_instances(self):
        instance_info = run_ec2_command(['ec2-describe-instances'] + self.instances.keys(), instance_cols)
        self.instances = read_info(instance_info, {'type' : 'INSTANCE'}, 'instance_id')

    def get_instances(self, status, boolean):
        if boolean:
            return [instance for instance in self.instances if self.instances[instance]['status'] == status]
        else:
            return [instance for instance in self.instances if self.instances[instance]['status'] != status]

    def _start_instances(self, command):
        instance_info = run_ec2_command(command, instance_cols)
        self.instances = read_info(instance_info, {'type' : 'INSTANCE'}, 'instance_id')

    def show_status(instances, urgent_instances=[]):
        if verbosity >= 2:
            for instance in instances:
                sys.stderr.write('Instance %s is %s.'
                                 % (instance, self.instances[instance]['status']))
        if verbosity >= 1:
            for instance in urgent_instances:
                sys.stderr.write('Instance %s is %s.'
                                 % (instance, self.instances[instance]['status']))

    def start(self, num, ami_id, keypair=None):
        self.update_images()
        if self.images.get(ami_id) is None:
            raise ValueError('The ami identifier %s does not exist or is not available!' % ami_id)

        args = ['ec2-run-instances', ami_id, '-n %i' % num]
        if keypair:
            args += ['-k %s' % keypair]

        self._start_instances(args)

        if len(self.instances) != num:
            raise Exception('We cannot find all the instances! '
                            'We recommend that you find and shut them down manually. '
                            'The known instances are: %s.' % ' '.join(self.instances.keys()))

        started = []
        while True:
            self.update_instances()
            unstarted = self.get_instances('running', False)
            if not unstarted:
                break
            just_started = [instance for instance in self.instances
                            if instance not in unstarted and instance not in started]
            started += just_started
            show_status(unstarted, just_started)    
            time.sleep(delay)

    def _stop_instances(self, command):
        process = Popen(command, stdout=PIPE)
        process.wait()
        #self.terminate = read_info(run_ec2_command(command, terminate_cols), {'type' : 'INSTANCE'}, 'instance_id')

    def stop(self):
        args = ['ec2-terminate-instances'] + self.instances.keys()
        self._stop_instances(args)

        terminated = []
        while True:
            self.update_instances()
            unterminated = self.get_instances('terminated', False)
            if not unterminated:
                break
            just_terminated = [instance for instance in self.instances
                               if instance not in unterminated and instance not in terminated]
            self.show_status(unterminated, just_terminated)
            time.sleep(delay)

        if self.verbose >= 1:
            sys.stderr.write('All instances have been terminated.\n')

    def start_process(self, command, instances=self.instances):
        for instance in self.instances:
            ip = get_ip(self.instances[instance]['ip'])
            command = 'bash %s' % command
            process = Popen(command_template % (self.port, ip, command),
                            stdin=PIPE, stdout=sys.stdout, stderr=sys.stderr))
            self.processes.append(process)

    def sync_up(self, files, opts='-rz', instances=self.instances):
        for instance in self.instances:
            ip = get_ip(self.instances[instance]['ip'])
            command = sync_up_template % (opts, self.port, files, ip, files)
            self.upload.append(Popen(command, stdout=PIPE, stderr=sys.stderr))

    def sync_down(self, files, opts='-rz', instances=self.instances):
        for instance in self.instances:
            ip = get_ip(self.instances[instance]['ip'])
            command = sync_down_template % (opts, self.port, ip, files, files)
            self.download.append(Popen(command, stdout=PIPE, stderr=sys.stderr))


def main():
    options, args = parser.parse_args()
    workers = options.workers
    executable = options.executable
    auxilary = options.auxilary
    syncronized = options.syncronized
    batch_size = options.batch_size
    ami_id = options.ami_id
    verbosity = options.verbosity

    workforce = Workforce()

    try:
        # Start
        workforce.start(workers, ami_id)
        time.sleep(delay)

        if verbosity >= 1:
            sys.stderr.write('Uploading auxilary files, syncronized files, and executable.\n')
        workforce.sync_up(auxilary, auxilary_rsync_opts)
        workforce.sync_up(syncronized, syncronized_rsync_opts)
        workforce.sync_up(executable, auxilary_rsync_opts)

        if verbosity >= 1:
            sys.stderr.write('Starting the executable.\n')
        workforce.start_process(executable)

        if verbosity >= 1:
            sys.stderr.write('Feeding input to the executable.\n')
        while True:
            processes = [process for process in workforce.processes
                         if process.returncode is None] # unfinished processes
            batches = group_by(group_by(sys.stdin, batch_size), processes)
            for batch, process in zip(batches, processes):
                for line in batch:
                    process.stdin.write(line)

        # Wait until all the processes are done
        if verbosity >= 1:
            sys.stderr.write('Waiting until all the processes are finished.\n')

        while True:
            if verbosity >= 2:
                print 'waiting...'

            processes = [process for process in workforce.processes
                         if process.returncode is None]
            if not processes:
                break
            time.sleep(delay)

        # Sync_down the syncronized files
        if verbosity >= 1:
            sys.stderr.write('Downloading syncronized files.\n')
        workforce.sync_down(syncronized, syncronized_rsync_opts)

        # Shut down the instances
        if verbosity >= 1:
            sys.stderr.write('Shutting down instances.\n')
        workforce.stop()

    except Exception e:
        sys.stderr.write('ERROR: %s\nTrying to shut down ec2 instances.' % e)
        try:
            workforce.stop()
        except KeyboardInterrupt:
            sys.stderr.write("""We are trying to close the instances!
                                Please don't interrupt.""")
            try:
                workforce.stop()
            except KeyboardInterrupt:
                sys.stderr.write("""If you interrupt one more time, you must close the instances manually!""")
                workforce.stop()
        except:
            sys.stderr.write("""FAILED TO SHUT DOWN INSTANCES!""")

if __name__ == '__main__':
    main()