#!/bin/bash

CMD=$(basename $0)
DIR=$(cd $(dirname $0) && pwd)

if [ -z "$PBS_JOBID" ]; then
    export JOBID=$$
else
    export JOBID=$PBS_JOBID
fi

INSTALL_DIR=$DIR/condor
LOCAL_DIR=$DIR/pilots/$JOBID
HOSTNAME=$(hostname -f)

echo "JOBID: $JOBID"
echo "DIR: $DIR"
echo "INSTALL_DIR: $INSTALL_DIR"
echo "LOCAL_DIR: $LOCAL_DIR"
echo "HOSTNAME: $HOSTNAME"

function error () {
    echo "$@" >&2
}

# Usage
function usage () {
    error """Usage: $CMD OPTIONS REQUIRED

REQUIRED

   -c | -condorHost <host:port>
      The condor central manager host.
   -t | -wallTime <minutes>
      The glidein runtime in minutes.

OPTIONS

   -it | -idleTime <minutes>
      The maximum idle time for this pilot.
   -d | -debug <level> [-debug <level> ...]
      Turn on extra debugging messages.
   -n | -numCpus
      The number of CPUs to allocate for this pilot.
   -hp | -highport
      The high end of the port range to allow Condor to use.
   -lp | -lowport
      The low end of the port range to allow Condor to use.
   -h | -help
      Display this usage message.
"""
}

CONDOR_HOST=""
WALL_TIME=""
IDLE_TIME="0"
DEBUG=""
NUM_CPUS="1"
HIGHPORT=""
LOWPORT=""

while [ $# -ge 1 ]; do
    case $1 in 
        -c | -condorHost) CONDOR_HOST=$2 ;;
        -t | -wallTime) WALL_TIME=$2 ;;
        -it | -idleTime) IDLE_TIME=$2 ;;
        -n | -numCpus) NUM_CPUS=$2 ;;
        -d | -debug) DEBUG="${DEBUG}$2 " ;;
        -hp | -highport) HIGHPORT=$2 ;;
        -lp | -lowport) LOWPORT=$2 ;;
        -h | -help) usage ; exit 1 ;;
        *) error "Invalid argument: $1" ; exit 1 ;;
    esac
    shift
    shift
done

# Check arguments
if [ -z "${CONDOR_HOST}" ]; then
    error "Must specify -condorHost"
    exit 1
fi
if [ -z "${WALL_TIME}" ]; then
    error "Must specify -wallTime"
    exit 1
fi

# Set default debug level
if [ -z "$DEBUG" ]; then
    DEBUG="D_ALWAYS"
fi

# Calculate wall time in seconds
WALL_TIME_SECONDS=`expr ${WALL_TIME} \* 60`

# Calculate idle time in seconds
IDLE_TIME_SECONDS="0"
if [ ! -z "${IDLE_TIME}" ]; then
    IDLE_TIME_SECONDS=`expr ${IDLE_TIME} \* 60`
fi

# Check install dir
if [ ! -d "${INSTALL_DIR}" ]; then
    error "Install directory does not exist: ${INSTALL_DIR}"
    exit 1
fi

# Set local directory
if ! mkdir -p "${LOCAL_DIR}"; then
    error "Unable to create pilot local dir: ${LOCAL_DIR}"
    exit 1
fi

# Copy config file
if ! cp ${DIR}/pilot_condor_config ${LOCAL_DIR}; then
    error "Unable to copy pilot_condor_config to ${LOCAL_DIR}"
    exit 1
fi

# Create log and execute dirs
if ! mkdir -p "${LOCAL_DIR}/log"; then
    error "Unable to create log dir: ${LOCAL_DIR}/log"
    exit 1
fi
if ! mkdir -p "${LOCAL_DIR}/execute"; then
    error "Unable to create execute dir: ${LOCAL_DIR}/execute"
    exit 1
fi


# Set environment
export CONDOR_CONFIG=${LOCAL_DIR}/pilot_condor_config
# We set LOG and EXECUTE here so they can't be changed.
export _condor_LOG=${LOCAL_DIR}/log
export _condor_EXECUTE=${LOCAL_DIR}/execute
# NUM_CPUS is special because they don't allow you to use an expression to
# define it. So you can't have NUM_CPUS = $(PILOT_NUM_CPUS) in the config
# file.
export _condor_NUM_CPUS=${NUM_CPUS}

# Create lock dir. Should be on /tmp to prevent locking problems.
LOCK=/tmp/condor_pilot_${JOBID}
if ! mkdir -p $LOCK; then
    error "Unable to create lock dir: $LOCK"
    exit 1
fi

# Add condor config entries
printf "\n\n" >> ${CONDOR_CONFIG}
echo "PILOT_CONDOR_CONFIG = ${CONDOR_CONFIG}" >> ${CONDOR_CONFIG}
echo "PILOT_LOCAL_DIR = ${LOCAL_DIR}" >> ${CONDOR_CONFIG}
echo "PILOT_CONDOR_HOST = ${CONDOR_HOST}" >> ${CONDOR_CONFIG}
echo "PILOT_WALL_TIME = ${WALL_TIME}" >> ${CONDOR_CONFIG}
echo "PILOT_WALL_TIME_SECONDS = ${WALL_TIME_SECONDS}" >> ${CONDOR_CONFIG}
echo "PILOT_SBIN = ${INSTALL_DIR}/sbin" >> ${CONDOR_CONFIG}
echo "PILOT_IDLE_TIME = ${IDLE_TIME}" >> ${CONDOR_CONFIG}
echo "PILOT_IDLE_TIME_SECONDS = ${IDLE_TIME_SECONDS}" >> ${CONDOR_CONFIG}
echo "PILOT_DEBUG = ${DEBUG}" >> ${CONDOR_CONFIG}
echo "PILOT_LOCK = ${LOCK}" >> ${CONDOR_CONFIG}
echo "PILOT_HIGHPORT = ${HIGHPORT}" >> ${CONDOR_CONFIG}
echo "PILOT_LOWPORT = ${LOWPORT}" >> ${CONDOR_CONFIG}
echo "PILOT_CCB_ADDRESS = ${CONDOR_HOST}" >> ${CONDOR_CONFIG}
echo "PILOT_JOBID = ${JOBID}" >> ${CONDOR_CONFIG}

# Fork condor_master
echo "Starting condor_master on host ${HOSTNAME}..."
exec ${INSTALL_DIR}/sbin/condor_master -f -r ${WALL_TIME}

