#!/bin/bash
# Redirect stdout ( > ) into a named pipe ( >() ) running "tee"
set -e

mkdir -p workflow/planning
COMMAND_LINE="${0} ${@}"
echo "${COMMAND_LINE}" > workflow/planning/pycbc-submit-dax-command.txt
echo "Executing on `hostname` in directory `pwd` at `date`" > workflow/planning/pycbc-submit-dax-log.txt

exec > >(tee -a workflow/pycbc-submit-dax-log.txt)

# Without this, only stdout would be captured - i.e. your
# log file would not contain any error messages.
# SEE answer by Adam Spiers, which keeps STDERR a seperate stream -
# I did not want to steal from him by simply adding his answer to mine.
exec 2>&1

LOCAL_PEGASUS_DIR=""
PEGASUS_PROPERTIES=""
NO_CREATE_PROXY=0
NO_QUERY_DB=0
NO_GRID=""
SUBMIT_DAX="--submit"
HTML_ENTITIES="{\"\'\": '&#39;', '(': '&#40;', ')': '&#41;', '+': '&#43;', '\"': '&quot;'}"

# These will be changed by the bundle builder
DATA_INLINE=False

function expand_pegasus_files() {
(
base64 -d <<DATA_END
PEGASUS_FILE_DATA
DATA_END
) | tar -zx
}

echo "# Properties set on command line" > extra-properties.conf
rm -f _reuse.cache
touch _reuse.cache
rm -f *-extra-site-properties.xml

GETOPT_CMD=`getopt -o p:P:K:Q:n:l:G:h --long pegasus-properties:,append-pegasus-property:,no-create-proxy,no-query-db,no-submit,local-dir:,no-grid,help -n 'pycbc_submit_dax' -- "$@"`
eval set -- "$GETOPT_CMD"

while true ; do
  case "$1" in
    -p|--pegasus-properties)
      case "$2" in
        "") shift 2 ;;
        *) PEGASUS_PROPERTIES=$2 ; shift 2 ;;
      esac ;;
    -P|--append-pegasus-property)
      case "$2" in
        "") shift 2 ;;
        *) echo $2 >> extra-properties.conf ; shift 2 ;;
      esac ;;
    -K|--no-create-proxy) NO_CREATE_PROXY=1 ; shift ;;
    -Q|--no-query-db) NO_QUERY_DB=1 ; shift ;;
    -n|--no-submit) SUBMIT_DAX="" ; shift ;;
    -l|--local-dir)
      case "$2" in
        "") shift 2 ;;
        *) LOCAL_PEGASUS_DIR=$2 ; shift 2 ;;
      esac ;;
    -G|--no-grid) NO_GRID="--forward nogrid" ; shift ;;
    -h|--help)
      echo "usage: pycbc_submit_dax [-h] [optional arguments]"
      echo
      echo "optional arguments:"
      echo "  -h, --help              show this help message and exit"
      echo "  -p, --pegasus-properties FILE use the specified file as"
      echo "                               the pegasus properties file"
      echo "  -P, --append-pegasus-property STRING add the extra property"
      echo "                                          specified by the argument"
      echo "  -K, --no-create-proxy   Do not run ligo-proxy-init and assume"
      echo "                             that the user has a valid grid proxy"
      echo "  -n, --no-submit         Plan the DAX but do not submit it"
      echo "  -l, --local-dir         Directory to put condor files under"
      echo "  -Q, --no-query-db       Don't query the pegasus DB."
      echo "  -G, --no-grid           Disable checks for grid proxy and"
      echo "                             GLOBUS_LOCATION in pegasus-plan"
      echo
      echo "If the environment variable TMPDIR is set then this is prepended to the "
      echo "path to the temporary workflow execute directory passed to pegasus-plan."
      echo "If the --local-dir option is not given."
      echo
      echo "If the environment variable PEGASUS_FILE_DIRECTORY is set then the"
      echo "script will look there for configuration, "
      echo "otherwise the script will look for this directory by querying the"
      echo "pycbc.workflow module."
      echo
      exit 0 ;;
    --) shift ; break ;;
    *) echo "Internal error!" ; exit 1 ;;
  esac
done

if [ $NO_CREATE_PROXY == 0 ]; then
  # Force the user to create a new grid proxy
  LIGO_USER_NAME=""
  while true; do
    read -p "Enter your LIGO.ORG username in (e.g. albert.einstein): " LIGO_USER_NAME
    echo
    if [ ! -z $LIGO_USER_NAME ] ; then
      break
    fi
  done
  unset X509_USER_PROXY
  ligo-proxy-init -p $LIGO_USER_NAME || exit 1
else
  if [ ! -z ${X509_USER_PROXY} ] ; then
    if [ -f ${X509_USER_PROXY} ] ; then
      cp -a ${X509_USER_PROXY} /tmp/x509up_u`id -u`
    fi
  unset X509_USER_PROXY
  fi
fi

if [ -z "${NO_GRID}" ] ; then
  #Check that the proxy is valid
  set +e
  grid-proxy-info -exists
  RESULT=$?
  set -e
  if [ ${RESULT} -eq 0 ] ; then
    PROXY_TYPE=`grid-proxy-info -type | tr -d ' '`
    if [ x${PROXY_TYPE} == 'xRFC3820compliantimpersonationproxy' ] ; then
      grid-proxy-info
    else
      cp /tmp/x509up_u`id -u` /tmp/x509up_u`id -u`.orig
      grid-proxy-init -hours 276 -cert /tmp/x509up_u`id -u`.orig -key /tmp/x509up_u`id -u`.orig
      rm -f /tmp/x509up_u`id -u`.orig
      grid-proxy-info
    fi
  else
    echo "Error: Could not find a valid grid proxy to submit workflow."
    exit 1
  fi
fi

#Make a directory for the submit files
SUBMIT_DIR=`mktemp --tmpdir=${LOCAL_PEGASUS_DIR} -d pycbc-tmp.XXXXXXXXXX`

#Make sure the directory is world readable
chmod 755 $SUBMIT_DIR

# find the site-local template directory
if [ $DATA_INLINE == "True" ] ; then
  expand_pegasus_files
  PEGASUS_FILE_DIRECTORY=${PWD}/pegasus_files
elif [ -z $PEGASUS_FILE_DIRECTORY ] ; then
  PEGASUS_FILE_DIRECTORY=`python -c 'from pycbc.workflow import PEGASUS_FILE_DIRECTORY;print(PEGASUS_FILE_DIRECTORY)'`
fi

# Plan the workflow
echo "Generating concrete workflow"

# cache the pegasus config
if [ -z ${PEGASUS_PROPERTIES} ] ; then
  cp $PEGASUS_FILE_DIRECTORY/pegasus-properties.conf ./pegasus-properties.conf
else
  cp ${PEGASUS_PROPERTIES} ./pegasus-properties.conf
fi

echo >> pegasus-properties.conf
cat extra-properties.conf >> pegasus-properties.conf


# Ian's bash-fu is not good enough to wrap this in the nice error handling that
# it deserves!
STORED_PLANNER_ARGS=`cat additional_planner_args.dat`

pegasus-plan --conf ./pegasus-properties.conf --dir $SUBMIT_DIR $SUBMIT_DAX $NO_GRID ${STORED_PLANNER_ARGS}

echo

rm -f submitdir
ln -sf $SUBMIT_DIR submitdir

echo "pegasus-status --verbose --long $SUBMIT_DIR/work \$@" > status
chmod 755 status

echo "pegasus-analyzer -r -v $SUBMIT_DIR/work \$@" > debug
chmod 755 debug

echo "pegasus-remove $SUBMIT_DIR/work \$@" > stop
chmod 755 stop

if [ -z "${NO_GRID}" ] ; then
  cat << EOF > start
#!/bin/bash

if [ -f /tmp/x509up_u\`id -u\` ] ; then
  unset X509_USER_PROXY
else
  if [ ! -z \${X509_USER_PROXY} ] ; then
    if [ -f \${X509_USER_PROXY} ] ; then
      cp -a \${X509_USER_PROXY} /tmp/x509up_u\`id -u\`
    fi
  fi
  unset X509_USER_PROXY
fi

# Check that the proxy is valid
grid-proxy-info -exists
RESULT=\${?}
if [ \${RESULT} -eq 0 ] ; then
  PROXY_TYPE=\`grid-proxy-info -type | tr -d ' '\`
  if [ x\${PROXY_TYPE} == 'xRFC3820compliantimpersonationproxy' ] ; then
    grid-proxy-info
  else
    cp /tmp/x509up_u\`id -u\` /tmp/x509up_u\`id -u\`.orig
    grid-proxy-init -hours 276 -cert /tmp/x509up_u\`id -u\`.orig -key /tmp/x509up_u\`id -u\`.orig
    rm -f /tmp/x509up_u\`id -u\`.orig
    grid-proxy-info
  fi
else
  echo "Error: Could not find a valid grid proxy to submit workflow."
  exit 1
fi

EOF
fi

echo "pegasus-run $SUBMIT_DIR/work \$@" > start

chmod 755 start

if [ -z ${SUBMIT_DAX} ] ; then
  echo
  echo "WARNING: DAX planned but not submitted. No dashboard entry has been created and"
  echo "         the workflow section of results page will not show a dashboard URL."
  echo "         You must run this script without the --no-submit option if this is a"
  echo "         production run."
  echo
  exit 0
fi

if [ ! -e ${HOME}/.pegasus/workflow.db ] ; then
  echo "WARNING: Could not find Pegasus dashboard database in ${HOME}/.pegasus"
  echo "         Workflow has been submitted but the results page will not contain"
  echo "         a link to the dashboard page. If this is a production workflow,"
  echo "         please remove the workflow, check for the origin of this error,"
  echo "         and re-submit the workflow by re-running this script."
  echo
  exit 1
fi

WORKFLOW_ID_STRING=""
WORKFLOW_DB_CMD="sqlite3 -csv ${HOME}/.pegasus/workflow.db \"select submit_hostname,wf_id,wf_uuid from master_workflow where submit_dir = '${SUBMIT_DIR}/work';\""
DB_TRY=0

# force a condor reschedule to get the workflow running
set +e
condor_reschedule &> /dev/null
set -e

if [ $NO_QUERY_DB == 0 ]; then

    /bin/echo "Querying Pegasus database for workflow stored in ${SUBMIT_DIR}/work"
    /bin/echo -n "This may take up to 120 seconds. Please wait..."
    rm -f pegasus_db.log
    touch pegasus_db.log
    # querying the database sometimes fails, so allow retries
    set +e
    until [ $DB_TRY -ge 15 ]
    do
      /bin/echo -n "."
      WORKFLOW_ID_STRING=`eval $WORKFLOW_DB_CMD 2>> pegasus_db.log`
      if [ $? -eq 0 ] && [ ! -z $WORKFLOW_ID_STRING ] ; then
        /bin/echo " Done."
        DB_QUERY_SUCCESS=0
        break
      else
        DB_QUERY_SUCCESS=1
      fi
      DB_TRY=$(( $DB_TRY + 1 ))
      for s in `seq ${DB_TRY}`
      do
        /bin/echo -n "."
        sleep 1
      done
    done
    set -e

    if [ ${DB_QUERY_SUCCESS} -eq 1 ] ; then
      echo; echo
      /bin/echo "Query failed: ${WORKFLOW_DB_CMD}"
      cat pegasus_db.log
    else
      rm -f pegasus_db.log
    fi

    if [ -z $WORKFLOW_ID_STRING ] ; then
      echo "WARNING: Could not find the workflow in the Pegasus dashboard database."
      echo "         Workflow has been submitted but the results page will not contain"
      echo "         a link to the dashboard page. If this is a production workflow,"
      echo "         please remove the workflow, check for the origin of this error,"
      echo "         and re-submit the workflow by re-running this script."
      echo
      exit 1
    fi

fi

WORKFLOW_ID_ARRAY=(${WORKFLOW_ID_STRING//,/ })
DASHBOARD_URL="https://${WORKFLOW_ID_ARRAY[0]}/pegasus/u/${USER}/r/${WORKFLOW_ID_ARRAY[1]}/w?wf_uuid=${WORKFLOW_ID_ARRAY[2]}"

echo ${DASHBOARD_URL} > workflow/pegasus-dashboard-url.txt

shopt -s nullglob
DASHBOARD_GLOB=(results/*_workflow/*DASHBOARD*.html)
DASHBOARD_PATH=${DASHBOARD_GLOB[0]}
if [[ -e ${DASHBOARD_PATH} ]] ; then
  DASHBOARD_COPY=workflow/`basename ${DASHBOARD_PATH}`.orig
  if [ -e ${DASHBOARD_COPY} ] ; then
    cp ${DASHBOARD_COPY} ${DASHBOARD_PATH}
  else
    cp ${DASHBOARD_PATH} ${DASHBOARD_COPY}
  fi

  COMMAND_LINE=$(python -c "import sys; from xml.sax.saxutils import (escape, unescape); print(escape(unescape(sys.stdin.read()), entities=${HTML_ENTITIES}))" <<< ${SITE_PROFILE})
  sed -i.bak -e "s+PYCBC_SUBMIT_DAX_ARGV+${COMMAND_LINE}+g" ${DASHBOARD_PATH}
  sed -i.bak -e "s+PEGASUS_DASHBOARD_URL+${DASHBOARD_URL}+g" ${DASHBOARD_PATH}
fi
shopt -u nullglob

echo "Workflow submission completed successfully."
echo
echo "The Pegasus dashboard URL for this workflow is:"
echo "  ${DASHBOARD_URL}"
echo
echo "Note that it make take a while for the dashboard entry to appear while the workflow"
echo "is parsed by the dashboard. The delay can be on the order of one hour for very large"
echo "workflows."
echo

exit 0
