#!/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 is local and doesn't already exist)"
  echo -e "  formatWorker  \t Format Tachyon worker storage"
  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 "  journalCrashTest\t Test the Master Journal System in a crash scenario. Try 'tachyon journalCrashTest -help' for more help."
  echo -e "                  \t NOTE: This command will stop the existing server and creates a new one!"
  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 -e "  validateConf  \t Validate Tachyon conf 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`

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

  TACHYON_CONF_DIR=$bin/../conf
  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
}

# bootstrap creates tachyon-env.sh
if [[ "${COMMAND}" == "bootstrap-conf" ]]; then
  bootstrapConf "$@"
  exit 0
fi

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|BasicNonByteBuffer|BasicRawTable> <STORE|NO_STORE> <SYNC_PERSIST|NO_PERSIST>"

  if [[ $# -ne 3 ]]; 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=""
  case "$1" in
    Basic)
      file+="/BasicFile_$2_$3"
      class="tachyon.examples.BasicOperations"
      ;;
    BasicNonByteBuffer)
      file+="/BasicNonByteBuffer_$2_$3"
      class="tachyon.examples.BasicNonByteBufferOperations"
      ;;
    BasicRawTable)
      file+="/BasicRawTable_$2_$3"
      class="tachyon.examples.BasicRawTableOperations"
      ;;
    *)
      echo "Unknown test: $1 with type $2 $3" 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} $2 $3

  # If test process received a signal abort early
  # Exit Codes for signals terminating a JVM are 128 + signal
  RET=$?
  if [[ ${RET} -gt 128 ]]; then
    echo "Test process was terminated due to signal"
    exit ${RET}
  fi
  return ${RET}
}

function journalCrashTest {
  MASTER_ADDRESS=${TACHYON_MASTER_ADDRESS}
  if [[ -z ${TACHYON_MASTER_ADDRESS} ]]; then
    MASTER_ADDRESS=localhost
  fi

  local class="tachyon.examples.JournalCrashTest"

  ${JAVA} -cp ${CLASSPATH} -Dtachyon.home=${TACHYON_HOME} -Dtachyon.master.hostname=${MASTER_ADDRESS} -Dtachyon.logs.dir=${TACHYON_LOGS_DIR} -Dtachyon.logger.type=USER_LOGGER ${TACHYON_JAVA_OPTS} ${class} $@
  RET=$?

  # If test process received a signal abort early
  # Exit Codes for signals terminating a JVM are 128 + signal
  if [[ ${RET} -gt 128 ]]; then
    echo "Test process was terminated due to signal"
    exit ${RET}
  fi
  return ${RET}
}

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

  keyword=$1
  count=0
  for pid in `ps -Aww -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 on `hostname`"
}

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 | grep -v '^#'`

  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
}

PARAMETER=""

if [[ ${COMMAND} == "format" ]]; then
  if [[ $# -eq 1 ]]; then
    if [[ $1 == "-s" ]]; then
      if [[ -e ${TACHYON_UNDERFS_ADDRESS} ]]; then
        # if ufs is local filesystem and already exists
        exit 0
      else
        declare -a schemes=(hdfs s3 s3n glusterfs swift)
        for scheme in ${schemes[@]}; do
          if [[ ${TACHYON_UNDERFS_ADDRESS} == ${scheme}://* ]]; then
            # if ufs is not local filesystem, don't format
            exit 0
          fi
        done

        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} == "formatWorker" ]]; then
  echo "Formatting Tachyon Worker @ `hostname -f`"
  CLASS=tachyon.Format
  PARAMETER=worker
elif [[ ${COMMAND} == "tfs" ]]; then
  CLASS=tachyon.shell.TfsShell
elif [[ ${COMMAND} == "loadufs" ]]; then
  CLASS=tachyon.client.UfsUtils
elif [[ ${COMMAND} == "runTest" ]]; then
  runTest $@
  exit $?
elif [[ ${COMMAND} == "runTests" ]]; then
  declare -a writeTypes=(MUST_CACHE CACHE_THROUGH THROUGH)
  declare -a readTypes=(CACHE_PROMOTE CACHE NO_CACHE)

  failed=0
  for readType in ${readTypes[@]}; do
    for writeType in ${writeTypes[@]}; do
      for example in Basic BasicNonByteBuffer BasicRawTable; do
        echo ${LAUNCHER} ${bin}/tachyon runTest ${example} ${readType} ${writeType}
        ${LAUNCHER} ${bin}/tachyon runTest ${example} ${readType} ${writeType}
        st=$?

        # If test process received a signal abort early
        # Exit Codes for signals terminating a JVM are 128 + signal
        if [[ ${st} -gt 128 ]]; then
          # Received some kind of signal
          echo "Test process was terminated due to signal"
          exit ${st}
        fi
        failed=$(( $failed + $st ))
      done
    done
  done

  if [[ ${failed} -gt 0 ]]; then
    echo "Number of failed tests: ${failed}" 1>&2
  fi
  exit ${failed}
elif [[ ${COMMAND} == "journalCrashTest" ]]; then
  journalCrashTest $@
  exit $?
elif [[ ${COMMAND} == "killAll" ]]; then
  killAll $@
  exit 0
elif [[ ${COMMAND} == "copyDir" ]]; then
  copyDir $@
  exit 0
elif [[ ${COMMAND} == "thriftGen" ]]; then
  rm -rf ${bin}/../common/src/main/java/tachyon/thrift
  thrift --gen java -out ${bin}/../common/src/main/java/. ${bin}/../common/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
elif [[ ${COMMAND} == "validateConf" ]]; then
  CLASS=tachyon.ValidateConf
else
  printUsage
  exit 1
fi

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