import py
import conftest
import sys

task = """
import py

pypywc = py.path.svnwc("/tmp/bolz/pypy-0.7.x")
pypywc.update()

testresult = py.path.svnwc("/tmp/bolz/pypy-0.7.x/testresult")
testdir = py.path.local("/tmp/bolz/pypy-0.7.x/lib-python/2.4.1/test")
cmd = "python ../../../pypy/test_all.py %s -T 10000 >& out"


def process_task(testname):
     test_cmd = cmd % testname
     testdir.chdir()
     py.std.os.system(test_cmd)
     unknowns = [wc for wc in testresult.listdir() if wc.status().unknown != []]
     for unknowndir in unknowns:
         for unknownfile in unknowndir.status().unknown:
             unknownfile.add()
         unknowndir.commit("tests")
     return testname

mynumber = channel.receive()
task_data = channel.receive()

while 1:
     result = process_task(task_data)
     channel.send((mynumber, result))
     next = channel.receive()
     if next == "stop":
          break
     elif next == "ok":
          pass
     else:
          raise ValueError, "don't know what to do"
     task_data = channel.receive()
"""

source = py.code.Source(task)
tests = [test.basename for test in conftest.testmap if test.core]
py.std.random.shuffle(tests)

print "running tests", tests

nodes = ["bolz@mu.physi.uni-heidelberg.de",
         "bolz@k0.physi.uni-heidelberg.de",
         "bolz@b0bar.physi.uni-heidelberg.de"]

class TaskDistributor(object):
     def __init__(self, nodes, tests):
          self.i = 0
          self.nodes = nodes
          self.channels = []
          for node in nodes:
               print "setting up", node, tests[self.i]
               gw = py.execnet.SshGateway(node)
               channel = gw.newchannel(self.new_task)
               gw.remote_exec(source, channel=channel)
               self.channels.append(channel)
               channel.send(len(self.channels) - 1)
               channel.send(tests[self.i])
               self.i += 1

     def new_task(self, message):
          number, results = message
          print "got result from", nodes[number], results
          channel = self.channels[number]
          print "a1"
          if self.i < len(tests):
              print "sending new task", tests[self.i]
              channel.send("ok")
              channel.send(tests[self.i])
              self.i += 1
          else:
              print "stopping"
              channel.send("break")
              channel.close()

if __name__ == '__main__':
    td = TaskDistributor(nodes, tests)
    try:
        while 1:
            py.std.time.sleep(10)
    except KeyboardInterrupt:
        res = [chan.close() for chan in td.channels]
        print "closing", res