#!/usr/bin/env bash

LAUNCHER=
# If debugging is enabled propagate that through to sub-shells
if [[ "$-" == *x* ]]; then
  LAUNCHER="bash -x"
fi

function printUsage {
  echo "Usage: tachyon COMMAND "
  echo "where COMMAND is one of:"
  echo -e "  format [-s]   \t Format Tachyon (if -s specified, only format if underfs doesn't exist)"
  echo -e "  bootstrap-conf\t Generate a config file if one doesn't exist"
  echo -e "  tfs           \t Command line input for generic filesystem user client."
  echo -e "  loadufs       \t Load existing files in underlayer filesystem into Tachyon."
  echo -e "  runTest       \t Run a end-to-end test on a Tachyon cluster."
  echo -e "  runTests      \t Run all end-to-end tests on a Tachyon cluster."
  echo -e "  killAll <WORD>\t Kill processes containing the WORD."
  echo -e "  copyDir <PATH>\t Copy the PATH to all worker nodes."
  echo -e "  clearCache    \t Clear OS buffer cache of the machine."
  echo -e "  thriftGen     \t Generate all thrift code."
  echo -e "  version       \t Print Tachyon version and exit."
  echo "Commands print help when invoked without parameters."
}

if [ "$#" == 0 ]; then
  printUsage
  exit 1
fi

COMMAND=$1
shift

bin=`cd "$( dirname "$0" )"; pwd`

DEFAULT_LIBEXEC_DIR="$bin"/../libexec
TACHYON_LIBEXEC_DIR=${TACHYON_LIBEXEC_DIR:-$DEFAULT_LIBEXEC_DIR}
. $TACHYON_LIBEXEC_DIR/tachyon-config.sh

function runTest {
  Usage="Usage: tachyon runTest <Basic|BasicRawTable|BasicCheckpoint|BasicNonByteBuffer> <MUST_CACHE|TRY_CACHE|CACHE_THROUGH|THROUGH|ASYNC_THROUGH>"

  if [ "$#" -ne 2 ]; then
    echo $Usage
    exit 1
  fi

  MASTER_ADDRESS=$TACHYON_MASTER_ADDRESS
  if [ -z $TACHYON_MASTER_ADDRESS ] ; then
    MASTER_ADDRESS=localhost
  fi

  local file="/default_tests_files"
  local class=""
  local args="$2"
  case "$1" in
    Basic)
      file+="/BasicFile_$2"
      class="tachyon.examples.BasicOperations"
      ;;
    BasicRawTable)
      file+="/BasicRawTable_$2"
      class="tachyon.examples.BasicRawTableOperations"
      ;;
    BasicCheckpoint)
      file+="/BasicCheckpoint"
      class="tachyon.examples.BasicCheckpoint"
      args="10"
      ;;
    BasicNonByteBuffer)
      file+="/BasicNonByteBuffer_$2"
      class="tachyon.examples.BasicNonByteBufferOperations"
      ;;
    *)
      echo "Unknown test: $1 with type $2" 1>&2
      echo "$Usage"
      exit 1
      ;;
  esac
  $LAUNCHER $bin/tachyon tfs rmr "$file"
  $JAVA -cp $CLASSPATH $TACHYON_JAVA_OPTS "$class" tachyon://$MASTER_ADDRESS:19998 "$file" "$args"
}

function killAll {
  if [ "$#" -ne 1 ]
  then
    echo "Usage: tachyon killAll <WORD>"
    exit
  fi

  keyword=$1
  count=0
  for pid in `ps -A -o pid,command | grep -i "[j]ava" | grep $keyword | awk '{print $1}'`; do
    kill -15 "$pid" > /dev/null 2>&1
    local cnt=30
    while kill -0 "$pid" > /dev/null 2>&1; do
      if [ $cnt -gt 1 ]; then
        # still not dead, wait
        cnt=`expr $cnt - 1`
        sleep 1
      else
        # waited long enough, kill the process
        echo "Process did not complete after 30 seconds, killing."
        kill -9 $pid 2> /dev/null
      fi
    done
    count=`expr $count + 1`
  done
  echo "Killed $count processes"
}

function copyDir {
  if [[ "$#" -ne 1 ]] ; then
    echo "Usage: tachyon copyDir <path>"
    exit 1
  fi

  bin=`cd "$( dirname "$0" )"; pwd`
  WORKERS=`cat $TACHYON_CONF_DIR/workers`

  DIR=`readlink -f "$1"`
  DIR=`echo "$DIR"|sed 's@/$@@'`
  DEST=`dirname "$DIR"`

  SSH_OPTS="-o StrictHostKeyChecking=no -o ConnectTimeout=5"

  echo "RSYNC'ing $DIR to workers..."
  for worker in $WORKERS; do
      echo $worker
      rsync -e "ssh $SSH_OPTS" -az "$DIR" "$worker:$DEST" & sleep 0.05
  done
  wait
}

function bootstrapConf {
  if [ $# -ne 1 ]; then
    echo "Usage $0 bootstrap-conf TACHYON_MASTER_HOSTNAME"
    exit 1
  fi

  if [ ! -e "$TACHYON_CONF_DIR/tachyon-env.sh" ]; then
    if [[ `uname -a` == Darwin* ]]; then
      # osx sed wants an argument to -i, use .bootstrap.bk
      TACHYON_SED="sed -i .bootstrap.bk"
      TOTAL_MEM=`sysctl hw.memsize | cut -d ' ' -f2`
      TOTAL_MEM=$[TOTAL_MEM / 1024]
      TOTAL_MEM=$[TOTAL_MEM * 2 / 3]
    else
      TACHYON_SED="sed -i"
      TOTAL_MEM=`awk '/MemTotal/{print $2}' /proc/meminfo`
      TOTAL_MEM=$[TOTAL_MEM * 2 / 3]
    fi

    # Create a default config that can be overridden later
    cp "$TACHYON_CONF_DIR/tachyon-env.sh.template" "$TACHYON_CONF_DIR/tachyon-env.sh"
    $TACHYON_SED "s/TACHYON_MASTER_ADDRESS=localhost/TACHYON_MASTER_ADDRESS=$1/" "$TACHYON_CONF_DIR/tachyon-env.sh"
    $TACHYON_SED "s/TACHYON_WORKER_MEMORY_SIZE=1GB/TACHYON_WORKER_MEMORY_SIZE=${TOTAL_MEM}KB/" "$TACHYON_CONF_DIR/tachyon-env.sh"
  fi
}

PARAMETER=""

if [ "$COMMAND" == "format" ]; then
  if [ $# -eq 1 ]; then
    if [ "$1" == "-s" ]; then
      if [ -e $TACHYON_UNDERFS_ADDRESS ] || [[ $TACHYON_UNDERFS_ADDRESS == hdfs://* ]] || [[ $TACHYON_UNDERFS_ADDRESS == s3://* ]] || [[ $TACHYON_UNDERFS_ADDRESS == glusterfs://* ]]; then
        # already exists, hdfs, or s3, don't format
        exit 0
      else
        shift # remove -s param
      fi
    else
      echo "Usage $0 format [-s]"
      exit 2
    fi
  elif [ $# -gt 1 ]; then
    echo "Usage $0 format [-s]"
    exit 2
  fi

  if [ -z $TACHYON_MASTER_ADDRESS ] ; then
    TACHYON_MASTER_ADDRESS=localhost
  fi

  $LAUNCHER $bin/tachyon-workers.sh $bin/tachyon formatWorker

  echo "Formatting Tachyon Master @ $TACHYON_MASTER_ADDRESS"
  CLASS=tachyon.Format
  PARAMETER=master
elif [ "$COMMAND" == "bootstrap-conf" ]; then
  bootstrapConf "$@"
  exit 0
elif [ "$COMMAND" == "formatWorker" ]; then
  echo "Formatting Tachyon Worker @ `hostname -f`"
  CLASS=tachyon.Format
  PARAMETER=worker
elif [ "$COMMAND" == "tfs" ]; then
  CLASS=tachyon.command.TFsShell
elif [ "$COMMAND" == "loadufs" ]; then
  CLASS=tachyon.util.UfsUtils
elif [ "$COMMAND" == "runTest" ]; then
  runTest "$@"
  exit $?
elif [ "$COMMAND" == "runTests" ]; then
  declare -a opArr=(MUST_CACHE TRY_CACHE CACHE_THROUGH THROUGH ASYNC_THROUGH)

  failed=0
  for op in ${opArr[@]}
  do
    for example in Basic BasicRawTable BasicCheckpoint BasicNonByteBuffer
    do
      echo $LAUNCHER $bin/tachyon runTest $example $op
      $LAUNCHER $bin/tachyon runTest $example $op
      st=$?
      failed=$(( $failed + $st ))
    done
  done

  if [ $failed -gt 0 ]; then
    echo "Number of failed tests: $failed" 1>&2
  fi
  exit $failed
elif [ "$COMMAND" == "killAll" ]; then
  killAll "$@"
  exit 0
elif [ "$COMMAND" == "copyDir" ]; then
  copyDir "$@"
  exit 0
elif [ "$COMMAND" == "thriftGen" ]; then
  rm -rf $bin/../core/src/main/java/tachyon/thrift
  thrift --gen java -out $bin/../core/src/main/java/. $bin/../core/src/thrift/tachyon.thrift
  exit 0
elif [ "$COMMAND" == "clearCache" ]; then
  sync; echo 3 > /proc/sys/vm/drop_caches ;
  exit 0
elif [ "$COMMAND" == "version" ]; then
  CLASS=tachyon.Version
else
  printUsage
  exit 1
fi

$JAVA -cp $CLASSPATH -Dtachyon.home=$TACHYON_HOME -Dtachyon.master.hostname=$TACHYON_MASTER_ADDRESS -Dtachyon.logger.type=USER_LOGGER $TACHYON_JAVA_OPTS $CLASS $PARAMETER $@
