#!/usr/bin/env python

import os
import sys
import subprocess
import time

pcs = set(open(os.path.expanduser('~/.labpcs')).read().strip().split())
command = sys.argv[1]

def sshrun(pc, command, **kw):
    return subprocess.Popen(['ssh', '-x', '-o', 'StrictHostKeyChecking=no', pc, command], **kw)

todo = sys.argv[2:]
if not todo:
    todo = [''] * len(pcs)
elif len(todo) == 1 and todo[0].startswith('py:'):
    todo = eval(todo[0][3:])

print 'PCs to use'
print '----------'
print ', '.join(sorted(pcs))
print
print
print 'Arguments'
print '---------'
print todo

# in seconds, how frequently to dump information about the running
# commands (e.g. how long they've been running, their PC names, etc.)
log_frequency = 10
last_print = time.time()

procs = []
while todo or procs:
    finished = [x for x in procs if x[1].poll() is not None]
    # deal with finished
    if finished:
        for pc, p, com, t in finished:
            # no route to host exit status = 255
            if p.poll() == 255:
                arg = com[len(command) + 1:]
                print 'no route to ', pc, '-- forgetting about it'
                print 'restarting', arg
                pcs.remove(pc)
                todo.insert(0, arg)
            else:
                print 'finished: %s on %s with status: %s' % (pc, com, p.poll())
    # filter out all finished
    procs = [x for x in procs if x[1].poll() is None]

    if (time.time() - last_print) > log_frequency:
        print 'in running:'
        for pc, p, c, t in procs:
            print '%s: %s  (%sm)' % (pc, c, int(round((time.time() - t) / 60, 0)))
        print
        last_print = time.time()

    # allocate new commands to any PCs not in use
    in_use = set(x[0] for x in procs)
    if in_use != pcs and todo:
        pc = (pcs - in_use).pop()
        arg = todo.pop(0)
        c = '%s %s' % (command, arg)
        procs.append((pc, sshrun(pc, c), c, time.time()))
        print 'starting:', (pc, c)
    else:
        time.sleep(1)