#!/usr/bin/env python

from __future__ import print_function

import os
import sys
import time
import thread
import tempfile
import subprocess as sp

"""
What this script does:

Run a command through qsub with all stdout/stderr directed to the current
terminal, the return code of the command also goes to the terminal. The 
script will not terminate until either the command being run or the qsub
job terminates. 

How to run this script:

qsub_run.py <qsub options> -- "<command to run>"

The <command to run> can be multiple commands separated by ;, |, && etc.
qsub_run.py collects <command to run> and wrapps it in a bash script. 

<command to run> may need to be surrounded by " to prevent the shell from 
interpreting any &&, ; etc. 

"""

def make_script(cmd):

    # FIXME: find a better tmp dir. It needs to be somewhere
    # accessible to the compute nodes.
    f, f_name = tempfile.mkstemp(dir='/short/v45/nah599/tmp')

    os.write(f, '#!/bin/bash\n')
    os.write(f, 'source /etc/profile.d/nf_sh_modules\n')
    os.write(f, ' '.join(cmd))
    os.close(f)
    os.chmod(f_name, 0755)

    return f_name


def main():

    assert('--' in sys.argv[1:])
    args = sys.argv[1:]
    sep_idx = args.index('--')
    qsub_cmd = args[:sep_idx]
    cmd_to_run = args[sep_idx + 1:]

    # Check for -o or -e option on the qsub command line.
    # If then do not exist then set up.
    fo_name = None
    fe_name = None
    if (not '-o' in qsub_cmd) and (not '-e' in qsub_cmd):
        fo, fo_name = tempfile.mkstemp(dir='/short/v45/nah599/tmp')
    
        qsub_cmd.append('-joe')
        qsub_cmd.append('-o')
        qsub_cmd.append(fo_name)

    # Put -Wblock=true in the args list.
    qsub_cmd.append('-Wblock=true')

    # Make the command script. 
    cmd_filename = make_script(cmd_to_run)

    # Submit the command. 
    cmd = ['qsub'] + qsub_cmd + ['--'] + [cmd_filename]
    print('Executing {}'.format(cmd))
    ret = sp.call(cmd)

    # The command has finished. Read output and write stdout.
    # We don't know when output has stopped so just keep trying
    # until it is all gone. 
    empty_reads = 0
    while True:
        out = os.read(fo, 1024*1024)
        print(out)

        if out == '':
            empty_reads += 1
        else:
            empty_reads = 0

        if empty_reads > 10:
            break

        time.sleep(2)

    # Delete temporary files
    #os.remove(fo_name)
    #os.remove(fe_name)
    #os.remove(cmd_filename)

    return ret

if __name__ == '__main__':
    sys.exit(main())
