#!/usr/bin/env python

USAGE="""
fx-rpc -- run programs written with garry's makeshift sockets api thing

more help to come
"""

import os
import signal
import sys
import util
import time
import random

# note: we're not using fx for parameter parsing right now...

if len(sys.argv) <= 1:
  print USAGE
  os.exit(1);

n = int(sys.argv[1])
if n <= 0:
  print "error: n must be positive"
  os.exit(1)

peers = sys.argv[2]
executable = sys.argv[3]
cwd = os.getcwd()

port = random.randrange(10000, 64000)

peerlist = util.readlines(peers)

pids = {}
log = sys.stderr
rootpid = -1
rootexitcode = 0

try:
  for i in range(n):
    machinename = peerlist[i]
    print >> log, "****** CONNECTING TO %s ******" % machinename
    remote_command = ["cd", cwd, "&&", executable] + sys.argv[4:] + [
        "--rpc/n=%d"%n, "--rpc/rank=%d"%i,
        "--rpc/port=%d"%port, "--rpc/peers=%s"%peers]
    argv = ["ssh", "-x", "-T", "-n", "-o", "StrictHostKeyChecking=no",
        machinename, " ".join(remote_command)]

    print >> log, " ".join(argv)
    pid = os.fork()
    if not pid:
      os.execvp(argv[0], argv)
    pids[pid] = i

  while pids:
    (pid, status) = os.wait()
    rank = pids.pop(pid)
    exitcode = util.getstatus(status)
    print >> log, "*********** Rank %d returned, exit code %d" % (pid, exitcode)
    if rank == 0:
      rootexitcode = exitcode

finally:
  for pid in pids:
    try:
      os.kill(pid, signal.SIGINT)
      print >> log, "Kill successful for %d" % pid
    except:
      try:
        os.kill(pid, signal.SIGTERM)
      except:
        try:
          os.kill(pid, signal.SIGKILL)
        except:
          print >> log, "Kill failed for %d" % pid

sys.exit(exitcode)