#!/afs/ihep.ac.cn/soft/juno/JUNO-ALL-SLC6/Pre-Release/J17v1r1-Pre2/ExternalLibs/Python/2.7.6/bin/python
import argparse
import subprocess
import os,sys
import time
assert os.environ.has_key("DistJETPATH")
sys.path.append(os.environ['DistJETPATH'])

DistTempDir = os.environ['HOME']+'/.DistJET'

def getArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c","--conf",dest="config_path", type=str,help="configure file path")
    parser.add_argument("--server",dest="server_host", type=str,help="where to start name service")
    parser.add_argument("-m","--app", dest="app_module", type=str,help="the app you want to run")
    parser.add_argument("--app-conf",dest="app_config", type=str,help="the configure file of Application")
    parser.add_argument("-n","--worker-num",dest="worker_num", type=int,help="the number of worker you want to start")
    parser.add_argument("-d",dest="log_level",action="store_true", default=False, help="debug mode")
    parser.add_argument("-s","--screen", dest="screen",action="store_true", default=False,help="print log output into screen")
    parser.add_argument("--backend",dest="backend",choices=["condor","local"],default="condor",help="Backend choice")
    parser.add_argument('-t', dest="test",action="store_true",default=False,help="For test")
    args = parser.parse_args()
    if (not args.app_module) or (not args.worker_num):
        print parser.error("less of app or worker number")
    return args

def gen_submit(num):
    subfile = "job_submit"
    while subfile in os.listdir("."):
        subfile+="_1"
    with open(subfile,'w+') as sub:
        sub.write("Universe = vanilla\n")
        sub.write("Executable = %s"%os.environ['DistJETPATH']+'/bin/ssh/run-sshd.sh\n')
        sub.write("Accounting_Group = juno\n")
        sub.write("getenv = True\n")
        #sub.write("initialdir = tmp/res\n")
        sub.write('Requirements = Target.OpSysAndVer =?= "SL6" && regexp("jnws*", Name)\n')
        sub.write("Queue %s\n"%str(num))

#clear config
if os.path.exists(DistTempDir+'/config.ini'):
    os.remove(DistTempDir+'/config.ini')

# check runtime environment
import psutil
pid = [p.info['pid'] for p in psutil.process_iter(attrs=['pid','name']) if 'hydra_nameserver' in p.info['name']]
if not pid:
    subprocess.Popen(['hydra_nameserver &'], shell=True)

time.sleep(1)
args = getArgs()
arg_dict = {}

# app module
app_module = os.path.abspath(args.app_module)
arg_dict['app_module'] = app_module
#config path
config_path = args.config_path
if config_path is None or config_path == "":
    config_path="null"
else:
    config_path = os.path.abspath(config_path)
arg_dict['config_path'] = config_path
#loglevel
loglevel = "info"
if args.log_level:
    loglevel = "debug"
arg_dict['loglevel'] = loglevel

#app config
if args.app_config is None or args.app_config == "":
    app_config = "null"
else:
    app_config = os.path.abspath(args.app_config) 
arg_dict['app_config'] = app_config
#server_host
server = "localhost"
arg_dict['server'] = server
#Rundir
curr_dir = os.path.abspath(os.getcwd())
runno = 0
while os.path.exists(curr_dir+"/Rundir_"+str(runno)):
	runno = runno+1
rundir = curr_dir+"/Rundir_"+str(runno)
os.mkdir(rundir)
os.chdir(rundir)

# set config file
import python.Util.Config as Config
if config_path != "null":
    Config.set_inipath(config_path)
Config.Config.setCfg('logconsole',args.screen)
Config.Config.setCfg('rundir',rundir)
Config.Config.setCfg('log_level',loglevel)

Config.Config()

if loglevel == "debug":
    print arg_dict
if args.test:
    exit(0)
# apply computing resource
import Backend.backend as bnd
backend = bnd.Backend()
if args.backend is None or args.backend=="condor":
    hostfile = None
    if backend.setBackend("htcondor"):
        host_list = backend.apply(args.worker_num)
        print host_list
        hostfile = open(DistTempDir+'/hostfile','w+')
        for k,v in host_list.items():
            hostfile.write(str(k)+':'+str(v)+'\n')
        hostfile.close()
    else:
        exit(0)

elif args.backend == "local":
    backend.setBackend('local')
    '''
	hostfile = open(DistTempDir+'/hostfile','w+')
    hostfile.truncate()
    hostfile.write(os.environ['HOSTNAME']+':'+str(args.worker_num))
    hostfile.flush()
    hostfile.close()
    '''
    host_list={'localhost':args.worker_num}



if loglevel == "debug":
    print os.path.abspath('.')


print 'mpiexec -nameserver %s python %s %s %s %s %s %s %s'%(server,os.environ['DistJETPATH']+'/bin/master.py',app_module, config_path,loglevel,app_config ,args.screen,rundir)
master = subprocess.Popen(['mpiexec -nameserver %s python %s %s %s %s %s %s %s'%(server,os.environ['DistJETPATH']+'/bin/master.py',app_module, config_path,loglevel,app_config ,args.screen,rundir)],shell=True)

while(True):
    if os.path.exists(DistTempDir+'/config.ini'):
        break
    time.sleep(5)
time.sleep(1)

#print 'mpiexec -n %s -nameserver %s -f %s python %s %s %s %s %s %s'%(args.worker_num,server,os.environ['HOME']+'/.DistJET/hostfile',os.environ['DistJETPATH']+'/bin/worker.py','1','null',loglevel,args.screen,rundir)
#worker = subprocess.Popen(['mpiexec -n %s -nameserver %s -f %s python %s %s %s %s %s %s'%(args.worker_num,server,os.environ['HOME']+'/.DistJET/hostfile',os.environ['DistJETPATH']+'/bin/worker.py','1','null',loglevel, args.screen, rundir)],shell=True)
for host,num in host_list.items():
    print 'mpiexec -n %s -nameserver %s -hosts %s python %s %s %s %s %s %s'%(num,server,host,os.environ['DistJETPATH']+'/bin/worker.py','1','null',loglevel,args.screen,rundir)
    subprocess.Popen(['mpiexec -n %s -nameserver %s -hosts %s python %s %s %s %s %s %s'%(num,server,host,os.environ['DistJETPATH']+'/bin/worker.py','1','null',loglevel, args.screen, rundir)],shell=True)

try:
    master.wait()
except KeyboardInterrupt:
    #clean temp file
    if os.path.exists(os.environ['HOME']+'/.DistJET/config.ini'):
        os.remove(os.environ['HOME']+'/.DistJET/config.ini')

#clean ssh jobs
backend.release()





