#!/usr/bin/env python

#   Licensed to the Apache Software Foundation (ASF) under one
#   or more contributor license agreements.  See the NOTICE file
#   distributed with this work for additional information
#   regarding copyright ownership.  The ASF licenses this file
#   to you under the Apache License, Version 2.0 (the
#   "License"); you may not use this file except in compliance
#   with the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing,
#   software distributed under the License is distributed on an
#   #  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#   KIND, either express or implied.  See the License for the
#   specific language governing permissions and limitations
#   under the License.

#
# To use this to execute your Eclipse workspace, install as a symbolic
# link in your JRE/bin directory, with any name other than 'viaducc',
# for instance, java_viaducc,, linking back to the base version of
# viaducc in your ducc_runtime.
#
# Then configre your Eclipse's launch JRE to use java_viaducc in place
# of 'java'

import sys
import os
import re

global ducc_mem_size
global ducc_mem_size_set
global ducc_desc_set
global default_jvm
global ducc_home

ducc_class = 'fixed'
ducc_mach_list = None
java_cmd   = None
ducc_home  = None       # tbd in a minute
ducc_mem_size = None
ducc_mem_size_set = False # whether the user passed a flag to manually set DUCC_MEMORY_SIZE
ducc_desc_set = False # whether the user passed a flag to manually set DUCC_DESCRIPTION
default_jvm = None

# get the default DUCC_HOME right away for message
realpath = os.path.realpath(sys.argv[0])
ndx = realpath.rindex('/')
ndx = realpath.rindex('/', 0, ndx)
ducc_home = realpath[:ndx]

def error(message):
    print "Error:", message
    sys.exit(1)

def usage():

    print "Usage: viaducc [defines] [command and parameters]"
    print "  -or-"
    print "       java_viaducc [defines] [java-class and parameters]"
    print ""
    print "Where defines include:"
    print "   -DDUCC_MEMORY_SIZE=size-in-gb"
    print "       The default is -DDUCC_MEMORY_SIZE="+str(ducc_mem_size)
    print""
    print "   -DDUCC_HOME=alternative-DUCC-runtime"
    print "       The default is -DDUCC_HOME=" + ducc_home
    print ""
    print "   -DDUCC_CLASS=ducc-scheduling-class"
    print "       The default is -DDUCC_CLASS=" + ducc_class
    print ""
    print "   -DDUCC_MACHINE_LIST=list-of-machine-names"
    print "       The default is no machine names.  The machine list string should be"
    print '       blank-delimeted and quoted, for example: -DDUCC_MACHINE_LIST="machine01 machine02 machine03"'
    print ""
    print "   -DDUCC_ENVIRONMENT=environment-settings"
    print "       The default is no additional environment.  The environment string should be"
    print '       blank-delimeted and quoted, for example: -DDUCC_ENVIRONMENT="A=B C=D"'
    if (default_jvm != None ):
        print ""
        print "   -DJAVA_BIN=specific-java-bin-directory"
        print "       For use with java_viaducc.  Details follow."
        print "       The default is -DJAVA_BIN=" + default_jvm
    print ""
    print "   -DDUCC_DESRIPTION=user-description-string"
    print "       The default description is the program name (viaducc or java_viaducc) with arguments"
    print ""
    print "   -DDUCC_ATTACH_CONSOLE"
    print "       Connects the remote process to local console. Remote process terminated if submitting process stops."
    print ""
    print "   When invoked as 'viaducc' the command is executed in a DUCC-managed resource"
    print ""
    print "   When invokded as 'java_viaducc' an appropriate JRE is found and the class"
    print "is executed as a java class in a DUCC-managed resource."
    print ""
    print "Notes for java_viaducc:"
    print ""
    print "   If java_viaducc is installed in a JRE/bin directory, it may be invoked as an"
    print "alternative to java from Eclipse, allowing Eclipse users to execute their workspaces"
    print "in a DUCC-managed resource.  JAVA_BIN should NOT be used in this case.  java_viaducc"
    print "will use the JRE it is installed in."
    print ""
    print "   If java_viaducc is used as a 'normal' command, the JRE is searched for in this order:"
    print "      1. Use the java specified by -DJAVA_BIN="
    print "      2. Use the java configured for DUCC"
    print ""
    print "   If -jar is used it must be the last JVM arg before any command line args"
    
    print
    sys.exit(0)

# where should I get java from?
def read_properties():
    ducc_properties = ducc_home + '/resources/ducc.properties'
    if ( not os.path.exists(ducc_properties) ):
        error("Cannot access DUCC_HOME at " + ducc_home)

    ducc_jvm = None
    # first see if i'm executed out of some JAVA_HOME somewhere
    jpath = os.path.abspath(sys.argv[0])    
    ndx = jpath.rindex('/')
    jdir = jpath[:ndx]
    java = jdir + '/java'
    if (os.path.exists(java)):
        ducc_jvm = java

    # nope, we'll use the default ducc java
    props = open(ducc_properties)
    try:
        for line in props:
            line = line.strip()
            mobj = re.search('[ =:]+', line)    # java props allows apace, =, or : as delimeters
            if ( mobj ):
                key = line[:mobj.start()].strip()
                val = line[mobj.end():].strip()
                if ( (ducc_jvm is None) and (key == 'ducc.jvm') ):       # yes!
                    ducc_jvm = val
                    continue
                if ( key == 'ducc.rm.share.quantum' ):       # yes!
                    ducc_mem = 1   # default quantum may not be the same for all classes
                    continue       # and the actual allocation will be rounded up to quantum
    finally:
        props.close();

    return (ducc_jvm, ducc_mem)

def parse_memory_string(text):
    """Given a string (e.g., "5G", "200m"), Returns the memory requested
    in bytes. For example:

        parse_memory_string('500m') should return
            500 * 1024 * 1024
        parse_memory_string('15G') should return
            15 * 1024 * 1024 * 1024
    """

    # this list should cover us for a while
    suffixes = {
        'k': 1024,
        'm': 1024 * 1024,
        'g': 1024 * 1024 * 1024,
        't': 1024 * 1024 * 1024 * 1024,
    }

    text = text.lower()
    if text[-1] in suffixes:
        multiplier = suffixes[text[-1]]
        text = text[:-1] # chop off suffix
    else:
        multiplier = 1

    return multiplier * int(text)

def parse_java_command_line():
    """Parse the command line and returns a dictionary including
    information on the maximum amount of memory the job might need,
    the Java class name, and its arguments."""
    maximum_memory_required = 1024.0 ** 3  # our default if -Xmx not specified
    args = iter(sys.argv[1:])
    skip_next_arg = False

    # scan args until we find the Java class
    while 1:
        arg = next(args).lower()
        if skip_next_arg:
            skip_next_arg = False
            continue

        if arg.startswith('-'):
            # these are the only options that take an argument and require
            # a space between the flag and the argument
            if arg in ('-cp', '-classpath'):
                skip_next_arg = True
            elif arg.startswith('-xmx'): # maximum Java heap size
                maximum_memory_required = parse_memory_string(arg[4:])
            continue

        break

    return dict(maximum_memory_required=maximum_memory_required,
                java_class_name=arg, java_class_args=list(args))

def main():

    global ducc_mem_size
    global ducc_mem_size_set
    global ducc_desc_set
    global default_jvm
    global ducc_home

    if ( len(sys.argv) == 1 ):
        usage()

    ducc_class = 'fixed'
    ducc_mach_list = None
    java_cmd   = None
    ducc_env   = ''
    enable_attach = False
    ducc_desc  = os.path.basename(sys.argv[0])

    # remember to add the '=' at the end if following value
    p_mem_size  = '-DDUCC_MEMORY_SIZE='
    p_class     = '-DDUCC_CLASS='
    p_mach_list = '-DDUCC_MACHINE_LIST='
    p_jvm_dir   = '-DJAVA_BIN='
    p_ducc_home = '-DDUCC_HOME='
    p_env       = '-DDUCC_ENVIRONMENT='
    p_attach    = '-DDUCC_ATTACH_CONSOLE'
    p_desc      = '-DDUCC_DESCRIPTION='

    need_java = False
    if ( os.path.basename(sys.argv[0]) != 'viaducc' ):
        # if invokded as 'java_viaducc we must inject some kind of jvm as the executable
        # this is more liberal: any link other than 'viaducc' caused injection of the
        # jvm, to allow multiple, different symlinks (in case there are multiple,
        # different versions of ducc, not really reccomended, but supported)
        need_java = True

    ducc_args = []
    for arg in sys.argv[1:]:
 
        if (arg in ('-h', '-?', '-help', '--help') ):
            # have to wait for parsing the args to look up DUCC_HOME so we defer emitting the message
            usage()
       
        elif (arg.startswith(p_mem_size) ):
            ducc_mem_size = arg[len(p_mem_size):]
            ducc_mem_size_set = True

        elif (arg.startswith(p_class) ):
            ducc_class = arg[len(p_class):]

        elif (arg.startswith(p_mach_list) ):
            ducc_mach_list = '"'+arg[len(p_mach_list):]+'"'
            
        elif (arg.startswith(p_jvm_dir) ):
            java_cmd = arg[len(p_jvm_dir):] + '/java'

        elif (arg.startswith(p_ducc_home) ):
            ducc_home = arg[len(p_ducc_home):]

        elif (arg.startswith(p_env) ):
            ducc_env = ' --environment "' + arg[len(p_env):] + '"'

        elif (arg.startswith(p_desc) ):
            ducc_desc = arg[len(p_desc):]
            ducc_desc_set = True

        elif (arg.startswith(p_attach) ):
            enable_attach = True

        else:
            ducc_args.append("'" + arg + "'")

    p = read_properties()
    default_jvm = p[0]
    if (ducc_mem_size is None):
        ducc_mem_size = p[1]

    description = os.path.basename(sys.argv[0])
    if ( need_java and (java_cmd is None) ):
        java_cmd = default_jvm
        if ( (java_cmd is None) or (java_cmd == '') ):
            error("Cannot figure out where java is")

    if need_java:
        java_command_line_info = parse_java_command_line()
        description += ' %s %s' % (java_command_line_info['java_class_name'],
                                    ' '.join(java_command_line_info['java_class_args']))

        import math
        # convert to gigabytes, round up
        maximum_memory_required = java_command_line_info['maximum_memory_required'] / (1024.0 ** 3)
        maximum_memory_required = int(math.ceil(maximum_memory_required))
        if maximum_memory_required > float(ducc_mem_size):
            if ducc_mem_size_set:
                print ("Warning: DUCC memory size specified as %sg but max heap set to %sg.\n"
                       "         Consider increasing -DDUCC_MEMORY_SIZE=<size>\n") % \
                      (ducc_mem_size, maximum_memory_required)
            else:
                print ("Info: DUCC memory size is not specified but asking JVM to set max heap to %sg.\n"
                       "      Setting DUCC memory size to %sg. Use -DDUCC_MEMORY_SIZE=<size> to specify\n") % \
                      (maximum_memory_required, maximum_memory_required)
                ducc_mem_size = maximum_memory_required

    # DUCC_HOME isn't ususally needed but viaducc is slithery and may end up tryng to
    # execute from somebody else's ducc so let's be sure we point at the right one
    os.environ['DUCC_HOME'] = ducc_home
    CMD = ducc_home + '/bin/ducc_process_submit'
    CMD = CMD +       ' --wait_for_completion'
    CMD = CMD +       ' --process_memory_size ' + str(ducc_mem_size)
    CMD = CMD +       ' --scheduling_class ' + ducc_class
    
    if(ducc_mach_list != None):
    	CMD = CMD +       ' --machine_list ' + ducc_mach_list

    if ducc_desc_set:
        CMD = CMD +       ' --description %r' % ducc_desc
    else:
        if need_java:
            CMD = CMD +       ' --description %r' % description # %r so it will be quoted
        else:
            CMD = CMD +        ' --description viaducc'

    if enable_attach:
        CMD = CMD +       ' --attach_console'
        CMD = CMD +       ' --cancel_on_interrupt'
    CMD = CMD +       ' --working_directory ' + os.getcwd()
    CMD = CMD +       ducc_env
 
    if ( need_java ):
        CMD = CMD +   ' --process_executable ' + java_cmd 
        if ( len(ducc_args) > 0):
            CMD = CMD +   ' --process_executable_args "' + ' '.join(ducc_args) + '"'
    else:
        if ( len(ducc_args) == 0 ):
            error("No command specified.")

        CMD = CMD +   ' --process_executable ' + ducc_args[0]
        if ( len(ducc_args[1:]) > 0):
            CMD = CMD +   ' --process_executable_args "' + ' '.join(ducc_args[1:]) + '"'

    print CMD
    sys.stdout.flush()         # UIMA-4168

    rc = os.system(CMD)
    sys.exit(rc);

main()