#!/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):

    f, f_name = tempfile.mkstemp(dir='./')

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

    return f_name

def qsub_finished(run_id):

    qsub_out = ''
    try:
        qsub_out = sp.check_output(['qstat', run_id], stderr=sp.STDOUT)
    except sp.CalledProcessError as err:
        qsub_out = err.output

    if 'Job has finished' in qsub_out or 'Unknown Job Id' in qsub_out:
        return True

    return False

def wait_and_get_qsub_output_filename(run_id):
    """
    qsub has been invoked with the -k option, this means that the output file
    will be $HOME/STDIN.o<run_id>.

    This function waits until the file exists and then returns it.
    """

    fname = os.path.join(os.environ['HOME'], 'STDIN.o' + run_id.split('.')[0])

    while not os.path.exists(fname):
        time.sleep(2)

    return fname

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. These are not
    # supported, instead output goes to stdout.
    if '-o' in qsub_cmd or '-e' in qsub_cmd:
        print('qsub options -o, -e not supported, all output goes to stdout.',
                file=sys.stderr)
        return 1

    qsub_cmd.append('-joe')
    qsub_cmd.append('-koe')

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

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

    fo_name = wait_and_get_qsub_output_filename(run_id)
    fo = open(fo_name)

    # As the command runs read output and write stdout.  We don't know when
    # output has stopped so just keep trying until the command has stopped and
    # there's not more output.
    empty_reads = 0
    while True:
        out = fo.read(1024)

        if out == '':
            empty_reads += 1
        else:
            empty_reads = 0
            print(out)

        job_done = qsub_finished(run_id)

        if empty_reads > 10 and job_done:
            break

        if job_done:
            time.sleep(2)
        else:
            time.sleep(10)

    # Delete temporary files
    fo.close()
    os.remove(fo_name)
    os.remove(cmd_filename)

    return 0

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