#!/usr/bin/env python

'''
This is a wrapper command for running multiple applications.

The general form of this comand is as follows:
$ stagerun.py [ : APPLICATION_COMMAND ] *

More details of options will be printed with '-h' option:
$ stagerun.py -h

For the most recent code, please check out here:
https://github.com/jychoi-hpc/stagerun
'''

import sys
import argparse
import subprocess
import os
import time
import logging
import copy
from threading  import Thread
from Queue import Queue, Empty

queue = Queue()
use_thread = False
use_color = False
use_tag = False

def dequeue_output(color=False):
    while True:
        stdnum, out = queue.get()
        if stdnum==2 and use_color:
            print '\x1b[6;37;41m'+out+'\x1b[0m'
        else:
            print out
        sys.stdout.flush()

def enqueue_output(stdnum, out, tag):
    if out is not None:
        while True:
            output = out.readline().rstrip()
            if output:
                queue.put((stdnum, tag + output.rstrip()))

def runthread(stdnum, out, tag):
    global use_thread
    if use_thread is False:
        t = Thread(target=dequeue_output)
        t.daemon = True
        t.start()
        use_thread = True
        
    t = Thread(target=enqueue_output, args=(stdnum, out, tag))
    t.daemon = True
    t.start()

def cmdlist(argv):
    '''
    return list of list
    '''
    cmds = list()
    _args = list()
    for x in argv:
        if x == ':':
            cmds.append(_args)
            _args = list()
        else:
            _args.append(x)

    cmds.append(_args)
    _args = list()

    return cmds

def getstds(args):
    if args.oe is None:
        f_stdout = open(args.stdout, 'w') if args.stdout is not None else subprocess.PIPE
        f_stderr = open(args.stderr, 'w') if args.stderr is not None else subprocess.PIPE
    else:
        f = open(args.oe, 'w')
        f_stdout = f
        f_stderr = f
    return (f_stdout, f_stderr)

def main():
    logging.basicConfig(level=logging.DEBUG, format='%(levelname)s: %(message)s')
    def usage():
        print ('USAGE: stagerun.py <OPTIONS> [ : <APP_COMMAND> ]*')
        print ('====================')
        parser_svr.print_help()
        print ('--------------------')
        parser_cmd.print_help()
        print ('--------------------')
        sys.exit()

    ##import ipdb; ipdb.set_trace()
    nclient = 0
    parser_svr = argparse.ArgumentParser(prog='OPTIONS', add_help=False)
    parser_svr.add_argument('--mpicmd', help='mpi command', default='mpirun')
    parser_svr.add_argument('--opt', help='options for mpi command', action='append', default=[])
    parser_svr.add_argument('--dryrun', help='dryrun', action='store_true')
    parser_svr.add_argument('--sleep', help='sleep time between executions', type=int, default=3)
    parser_svr.add_argument('--serial', help='serial execution', action='store_true')
    parser_svr.add_argument('--usecolor', help='use color', action='store_true')
    parser_svr.add_argument('--tag-output', help='tag output', action='store_true')

    parser_cmd = argparse.ArgumentParser(prog='APP_COMMAND', add_help=False)
    parser_cmd.add_argument('--np', '-n', help='num. of processes', type=int, default=1)
    parser_cmd.add_argument('CMDS', help='app commands', nargs=argparse.REMAINDER)
    parser_cmd.add_argument('--stdout', '-o', help='stdout')
    parser_cmd.add_argument('--stderr', '-e', help='stderr')
    parser_cmd.add_argument('--oe', help='merging stdout and stderr')
    parser_cmd.add_argument('--nompi', action='store_true', help='no mpi')
    parser_cmd.add_argument('--opt', help='options for mpi command', action='append', default=[])
    parser_cmd.add_argument('--cwd', help='work directory')
    parser_cmd.add_argument('--env', '-x', help='env key=value', action='append', default=[])
    parser_cmd.add_argument('--norun', help='no run', action='store_true')
    parser_cmd.add_argument('--tag', help='stdout tag')
    parser_cmd.add_argument('--waitfile', help='check file before executing current command')
    parser_cmd.add_argument('--checkfile', help='check file before executing next command')

    cmds = cmdlist(sys.argv[1:])
    if len(cmds) < 2:
        usage()

    args, _unknown = parser_svr.parse_known_args(cmds[0])
    if len(_unknown) > 0:
        usage()

    global use_color
    global use_tag
    use_color = args.usecolor
    use_tag = args.tag_output

    args_cmd_list = list()
    for cmd in cmds[1:]:
        args_cmd, _unknown = parser_cmd.parse_known_args(cmd)
        setattr(args_cmd, 'UNKNOWN', _unknown)
        nclient += args_cmd.np if not args_cmd.norun else 0
        args_cmd_list.append(args_cmd)

    # Run client
    env = os.environ.copy()
    plist = list()
    for a in args_cmd_list:
        if not a.nompi:
            cl_cmd = ' '.join(['%(mpicmd)s -n %(np)d %(opt)s' % {'mpicmd':args.mpicmd, 'np':a.np, 'opt':' '.join(a.opt)}, ' '.join(a.UNKNOWN), ' '.join(a.CMDS)])
        else:
            cl_cmd = ' '.join([' '.join(a.CMDS)])

        if not a.norun:
            logging.debug('CMD: %s' % cl_cmd)
        
        if not args.dryrun and not a.norun:
            if args.serial and len(plist)>0:
                p = plist.pop()
                logging.debug('Waiting PID: %r' % p.pid)
                p.wait()

            if a.waitfile is not None:
                while not os.path.exists(a.waitfile):
                    logging.debug('Waiting file: %r' % a.waitfile)
                    time.sleep(1)

            cl_env = copy.deepcopy(env)
            for kv in a.env:
                k, v = kv.split('=')
                cl_env[k] = v

            f_stdout, f_stderr = getstds(a)
            p = subprocess.Popen(cl_cmd.split(), env=cl_env,
                    stdout=f_stdout, stderr=f_stderr,
                    close_fds=True, cwd=a.cwd)
            logging.debug('PID: %r (%r)' % (p.pid, cl_cmd))
            plist.append(p)

            tag = '[%d]:'%p.pid if a.tag is None else a.tag
            runthread(1, p.stdout, tag)
            runthread(2, p.stderr, tag)

            if a.checkfile is not None:
                while not os.path.exists(a.checkfile):
                    logging.debug('Waiting file: %r' % a.checkfile)
                    time.sleep(1)

            time.sleep(args.sleep)

    if not args.dryrun and len(plist)>0:
        for p in plist:
            logging.debug('Waiting PID: %r' % p.pid)
            p.wait()
    logging.info('Done.')

if __name__ == '__main__':
    main()
