#!/usr/bin/env python2.3
# -*- python -*-
from __future__ import division

__version__ = "$Revision: 1.50 $"

import commands
import errno
from optparse import make_option, OptionParser
import os
import popen2
import re
from select import select
import shutil
import sys

import poly
from poly._util import die

POLYJOB_CMD = "polyjob"
POLYWATCH_CMD = "polywatch"

BSUB_CMD = "bsub"
POLYCAT_CMDLINE = """bsub -q "small" -w 'ended(%s)' -J "%s" polycat %s %s %s"""
POLYCAT_JOBNAME = "polycat_%s"
DEFAULT_DIRNAME = "polydata"

SUPERGLOBAL_FILENAME = "superglobals.pickle"

def polycat_jobname(jobname):
    return jobname.replace("[", "_").replace("]", "_") + "polycat"

def _run_bsub(cmdline, function):
    status, output, stderr = function(cmdline)

    if status:
        print output
        print >>sys.stderr, stderr
        
        die("bsub failed with status %d" % status)
    else:
        if stderr:
            print >>sys.stderr, stderr
        return output

def run_bsub(cmdline):
    return _run_bsub(cmdline, poly._getstatus_stdout_stderr)

def main(args):
    long_args = []
    bsub_args = [BSUB_CMD]
    
    for index, arg in enumerate(args):
        if not arg.startswith("--") or arg == "--":
            long_args = args[:index]
            bsub_args.extend(args[index:])
            break
    else:
        long_args = args
    
    parser = OptionParser(usage="%prog [POLYSUB-OPTION]... -J JOBNAME [BSUB-OPTION]... -- COMMAND [COMMAND-ARG]...",
                          version="%%prog %s" % __version__,
                          option_list=[
        make_option("", "--output",
                    action="store", dest="stdout_filename", default=None,
                    help="write standard output to FILE", metavar="FILE"),
        make_option("", "--error",
                    action="store", dest="stderr_filename", default=None,
                    help="write standard error to FILE", metavar="FILE"),
        make_option("", "--superglobal",
                    action="store_true", dest="superglobal", default=False,
                    help="generate superglobals interactively before submitting batch job array"),
        make_option("", "--remove",
                    action="store_true", dest="remove", default=False,
                    help="remove existing desination files before running"),
        make_option("", "--force",
                    action="store_true", dest="force", default=False,
                    help="force run")
        ])
    options, long_args = parser.parse_args(long_args)
    assert not long_args
    
    destination_host = poly._try_each(KeyError,
                                      lambda: os.environ["POLY_STORAGE_HOST"],
                                      lambda: os.environ["HOSTNAME"],
                                      lambda: commands.getoutput("hostname"))

    try:
        destination = os.environ["POLY_DESTINATION"]
        destination_split = destination.split(":")
        if len(destination_split) != 2:
            die("Environment variable POLY_DESTINATION must contain only one colon")

        # assumes destination_dir is available from submitting host, but its parents will be created if it is not
        destination_dir = destination_split[1]
    except KeyError:
        destination_dir = os.path.join(os.environ["HOME"], DEFAULT_DIRNAME)
        destination = os.environ["POLY_DESTINATION"] = "%s:%s" % (destination_host, destination_dir)

    try:
        program_args = bsub_args[bsub_args.index("--")+1:]
        bsub_args[bsub_args.index("--")] = POLYJOB_CMD
    except ValueError:
        die("job specification did not include --")

    try:
        os.environ["POLYWATCH_HOSTS"]
        use_polywatch = True
    except KeyError:
        use_polywatch = False
        
    try:
        args_index_jobname = bsub_args.index("-J")+1 
        jobname = bsub_args[args_index_jobname]

        if use_polywatch:
            try: # job slot specified; fix jobname
                jobname = jobname[:jobname.index("%")]
                use_polywatch = False
            except ValueError: # job slot limit not specified; fix bsub arg
                bsub_args[args_index_jobname] = jobname + "%1" # job slot limit
            
    except ValueError:
        die("job specification did not include -J jobname")

    try:
        os.environ["BASH_ENV"] # if you don't want it, set it to " " or /dev/null
    except KeyError:
        os.environ["BASH_ENV"] = "~/.bashrc"

    jobdir = os.path.join(destination_dir, jobname)
    if options.stdout_filename is None:
        options.stdout_filename = jobdir + os.extsep + "out"
    elif options.stdout_filename == "":
        options.stdout_filename = "/dev/null"
        
    if options.stderr_filename is None:
        if options.stdout_filename.endswith(os.extsep + "out"):
            options.stderr_filename = options.stdout_filename[:-3] + "err" # extsep not needed
        else:
            options.stderr_filename = options.stdout_filename + os.extsep + "err" # append to end
    elif options.stderr_filename == "":
        options.stderr_filename = "/dev/null"
            
    if options.remove:
        try:
            shutil.rmtree(jobdir)
        except OSError:
            pass
        for filename in options.stdout_filename, options.stderr_filename:
            try:
                os.remove(filename)
            except OSError:
                pass

    if not options.force:
        try:
            os.makedirs(jobdir) # if it already exists this will throw an exception
        except OSError, err:
            if err.errno == errno.EEXIST:
                die("destination store %s already exists; use %s --remove to remove it" % (jobdir, sys.argv[0]))
            else:
                die("can't create destination store %s: %s" % (jobdir, err.strerror))
        
    for dirname in ["hostinfo", "stdout", "stderr", "tar"]:
        os.makedirs(os.path.join(jobdir, dirname))

    if options.superglobal:
        superglobal_filename = os.path.join(jobdir, SUPERGLOBAL_FILENAME)
        
        os.environ["POLY_SUPERGLOBAL_SAVE"] = superglobal_filename
        os.spawnvpe(os.P_WAIT, program_args[0], program_args, os.environ)
        del os.environ["POLY_SUPERGLOBAL_SAVE"]
        
        os.environ["POLY_SUPERGLOBAL_LOAD"] = superglobal_filename

    bsub_output = run_bsub(bsub_args)
    print "Main", bsub_output
    jobid = bsub_output[bsub_output.index("<")+1:bsub_output.index(">")]

    # XXX: runs through subshell. change this, and expanduser() the filenames
    polycat_cmd = POLYCAT_CMDLINE % (jobid, polycat_jobname(jobname), jobdir, options.stdout_filename, options.stderr_filename)

    print "Summary ", run_bsub(polycat_cmd)

    sys.stdout.flush()
    if use_polywatch:
        os.system(" ".join([POLYWATCH_CMD, jobid]))
    
if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))

