#!/usr/bin/bash

# determine STUBL install location
# Copied from Apache Ant:
# https://git-wip-us.apache.org/repos/asf?p=ant.git;a=blob;f=src/script/ant;h=b5ed5be6a8fe3a08d26dea53ea0fb3f5fab45e3f
if [ -z "$STUBL_HOME" -o ! -d "$STUBL_HOME" ] ; then
  ## resolve links - $0 may be a link to stubl's home
  PRG="$0"
  progname=`basename "$0"`

  # need this for relative symlinks
  while [ -h "$PRG" ] ; do
    ls=`ls -ld "$PRG"`
    link=`expr "$ls" : '.*-> \(.*\)$'`
    if expr "$link" : '/.*' > /dev/null; then
    PRG="$link"
    else
    PRG=`dirname "$PRG"`"/$link"
    fi
  done

  STUBL_HOME=`dirname "$PRG"`/..

  # make it fully qualified
  STUBL_HOME=`cd "$STUBL_HOME" > /dev/null && pwd`
fi

# setup STUBL environment
. $STUBL_HOME/conf/stubl 

TINS=$STUBL_HOME/bin/slogs_helpers/tins
SLOGPLUS=$STUBL_HOME/bin/slogs_helpers/slogplus

if [ "$1" == "--help" ]; then
  echo "======================================================="
  echo ""
  echo "  slogs - A script that retrieves resource usage and   "
  echo "  accounting information for a user or list of users.  "
  echo "  For each job run after the given start date, the     "
  echo "  following information is gathered from the SLURM     "
  echo "  accounting logs:  " 
  echo "     number of CPUS, wait time, start time, elapsed    "
  echo "     time, Amount of RAM Requested, Avg RAM Used, and  "
  echo "     Max RAM Used "
  echo ""
  echo "  Usage:  slogs [start_date] [users] [sacct_args]      "
  echo ""
  echo "  Notes: start_date should be in MMDDYY format.        "
  echo "         Use a comma-separated list for multiple users."
  echo "         All sacct_args passed along to sacct command. "
  echo ""
  echo "  Extras: Include a \"--plus\" argument to select an   "
  echo "          enhanced output format that includes data on "
  echo "          cpu and memory efficiency.                   "
  echo ""
  echo "          Use the SACCT_XFMT environment variable to   "
  echo "          request additional fields of output.         "
  echo ""
  echo " Defaults:                                             "
  echo "    start_date : 010113 if no start date is given.     "
  echo ""
  echo "    users : Current user if no user name is given.     "
  echo "            \"all\" selects all users                  "
  echo "======================================================="

  exit
fi

# try to figure out if users are in first or second arg
if [ "$1" != "" -a "`echo $1 | tr ',' '\n' | head -n1 | xargs id 2>/dev/null | wc -l`" != 0 ]; then
  dol1=$2
  dol2=$1
else
  dol1=$1
  dol2=$2
fi

dol1Test=`echo $dol1 | grep '\--' | wc -l`
if [ "$dol1" == "" -o "$dol1Test" == "1" ]; then
  StartDate=010113
else
  StartDate=$dol1
  shift
fi

if [ "$dol2" == "" ]; then
  TheGroup=`whoami`
else
  TheGroup=$dol2
  shift
fi

bPlus=`echo " $@ " | grep ' --plus\|^--plus ' | wc -l`
TheArgs=`echo " $@ " | sed 's/^--plus //g' | sed 's/ --plus//g'`

if [ "$TheGroup" == "all" ]; then
    TheGroup=`sacct --start=${StartDate} --allusers $TheArgs -ouser -n | \
             sort -u | \
             sed 's/ //g' | \
             tr '\n' ',' | \
             sed 's/^,//g' | \
             sed 's/,$//g'`
fi
TheList=`echo "$TheGroup" | tr ',' ' '`
ListSize=`echo "$TheGroup" | tr ',' '\n' | wc -l`

EndDate=`date +%m%d%y`

if [ "$bPlus" == "1" ]; then
  export SACCT_XFMT="TotalCPU"
fi

export SACCT_FORMAT="jobid%20,User,ncpus,start,elapsed,ReqMem,AveRSS,MaxRSS"
if [ "$SACCT_XFMT" != "" ]; then
  export SACCT_FORMAT="${SACCT_FORMAT},${SACCT_XFMT}"
fi

if [ "${ListSize}" == "1" ]; then
  TheGroupName=${TheGroup}
else
  TheGroupName=multiple_users
fi

TMPFILE=/tmp/${TheGroupName}_usage_logs_${StartDate}_to_${EndDate}.tmp.$$
PLUSFILE=/tmp/${TheGroupName}_usage_logs_${StartDate}_to_${EndDate}.plus.$$
#OUTFILE=${TheGroup}_usage_logs_${StartDate}_to_${EndDate}.txt

#echo $SACCT_HEADER > $TMPFILE
rm -f $TMPFILE

bHeader=1
for i in $TheList; do
  echo "Retrieving accounting data for user $i ..."
  user9=`echo $i | awk '{ printf("%9s", $1); }'`
  if [ "$bHeader" == "1" ]; then
    sacct --start=${StartDate} --user=$i $TheArgs  | $TINS 21 "$user9" 2 >> $TMPFILE
  else
    sacct --start=${StartDate} --user=$i -n $TheArgs  | $TINS 21 "$user9" 0 >> $TMPFILE
  fi
  bHeader=0
done
#cat $TMPFILE | sed "s/T/ /g" | tee $OUTFILE

if [ "$bPlus" == "1" ]; then
 cat $TMPFILE | sed 's/\([0-9]\)T\([0-9]\)/\1 \2/g' | grep '\.batch\|^               JobID' > $PLUSFILE
 # echo "$SLOGPLUS $PLUSFILE"
 $SLOGPLUS $PLUSFILE
 mv $PLUSFILE `pwd`
else
  cat $TMPFILE | sed 's/\([0-9]\)T\([0-9]\)/\1 \2/g'
fi

rm $TMPFILE

