#!/usr/bin/env python
#
# TODO:
# - document this
# - make it work outside of $GAPROOT/bin
# - add proper error handling, e.g. when gap.sh cannot be found
# - make sure the GAP instances call 'LoadPackage("zeromqinterface");'
# - make sure this actually works (again :-)

import sys, os, subprocess, time, tempfile, shutil, re, glob

args = sys.argv[1:]

if len(args) < 2 or args[0] != "-N":
  print("Usage: zgap -N <nodes> <gap arguments>")
  sys.exit(1)

conf = args[1]
args = args[2:]

if re.match(r"^[0-9]+$", conf):
  numnodes = int(conf)
  numa = False
else:
  numarel = False
  m = re.match(r"^([0-9]+):([0-9]+)-([0-9]+)$", conf)
  if not m:
    m = re.match(r"^([0-9]+):\+([0-9]+)-([0-9]+)$", conf)
    numarel = True
  if not m:
    print("Usage: zgap -N <nodes> <gap arguments>")
    sys.exit(1)
  numnodes, startnode, endnode = tuple(map(int, m.groups()))
  numa = True

dir = tempfile.mkdtemp()
try: os.mkdir(dir)
except: pass

print "[zgap] %s" % dir

nodes = [ None ] * (numnodes + 1)

if numa:
  gap = [ "numactl", "-N", None, "-l", "bin/gap.sh" ]
else:
  gap = [ "bin/gap.sh" ]

if numa:
  numanode = startnode
for i in xrange(numnodes+1):
  node_file = "%s/setup%02d.g" % (dir, i)
  log = open("%s/log%02d.txt" % (dir, i), "w")
  fp = open(node_file, "w")
  init = """
  LoadPackage("ZeroMQInterface");
  ReadPackage("ZeroMQInterface", "gap/znode.g");
  ZSetNodes(%d, [1..%d]);
  ZSetIncoming("ipc://%s/%02d.socket");
  """ % (i, numnodes, dir, i)
  for j in xrange(0, numnodes+1):
    init += """ZSetOutgoing(%d, "+ipc://%s/%02d.socket");\n""" %(j, dir, j)
  fp.write(init)
  fp.close()
  if i > 0:
    if numa:
      gap[2] = str(numanode)
    nodes[i] = subprocess.Popen(gap + [ "-S", "-T", node_file ] + args,
      stdin = subprocess.PIPE, stdout = log, stderr = log)
  log.close()
  if numa:
    numanode += 1
    if numanode > endnode:
      numanode = startnode

main_file = "%s/setup00.g" % dir

subprocess.call(gap + [ main_file ] + args)

for node in nodes[1:]:
  node.stdin.write("QUIT_GAP();\n");
  node.stdin.flush();
time.sleep(0.1);
for node in nodes[1:]:
  node.kill()

for file in (glob.glob(dir+"/setup[0-9][0-9].g") +
             glob.glob(dir+"/log[0-9][0-9].txt") +
      glob.glob(dir+"/[0-9][0-9].socket")):
  try: os.unlink(file)
  except: pass
os.rmdir(dir)