#! /usr/bin/env bash

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

bin_dir=$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )
at_home=$( cd "$( dirname "$bin_dir" )" && pwd )
at_version=2.0.0-SNAPSHOT

function print_usage() {
  cat <<EOF

Usage: accumulo-testing <command> (<argument>)

Possible commands:
  agitator <command>            Runs agitator <command>. Available commands: start, stop
  ci-createtable                Creates Accumulo table with splits for continuous ingest tests
  ci-local <application>        Runs continuous ingest <application> in local java process
                                Applications: ingest, walk, batchwalk, scan
  ci-yarn <num> <application>   Runs continuous ingest <application> in <num> containers in YARN
                                Applications: ingest, walk, batchwalk, scan
  ci-mapred <application>       Runs continuous ingest mapreduce <application>
                                Applications: verify, moru  
  rw-local <module>*            Runs randomwalk <module> in local java process
  rw-yarn <num> <module>*       Runs randomwalk <module> in <num> containers on YARN
                                *Modules located in core/src/main/resources/randomwalk/modules
EOF
}

if [ -f "$at_home/conf/accumulo-testing-env.sh" ]; then
  . "$at_home"/conf/accumulo-testing-env.sh
else
  . "$at_home"/conf/accumulo-testing-env.sh.example
fi

if [ -z "$ACCUMULO_VERSION" ]; then
  echo "ERROR: ACCUMULO_VERSION must be set conf/accumulo-testing-env.sh"
  exit 1
fi
if [ -z "$HADOOP_VERSION" ]; then
  echo "ERROR: HADOOP_VERSION must be set conf/accumulo-testing-env.sh"
  exit 1
fi

at_props="$at_home/conf/accumulo-testing.properties"
if [ ! -f "$at_props" ]; then
  echo "Please create and edit accumulo-testing.properties in $at_home/conf"
  exit 1
fi

log4j_config="$at_home/conf/log4j.properties"
if [ ! -f "$log4j_config" ]; then
  log4j_config="$at_home/conf/log4j.properties.example"
  if [ ! -f "$log4j_config" ]; then
    echo "Could not find logj4.properties or log4j.properties.example in $at_home/conf"
    exit 1
  fi
fi

function build_shade_jar() {
  export at_shaded_jar="$at_home/core/target/accumulo-testing-core-$at_version-shaded.jar"
  export CLASSPATH="$at_shaded_jar:$CLASSPATH"
  if [ ! -f "$at_shaded_jar" ]; then
    echo "Building $at_shaded_jar"
    cd "$at_home" || exit 1
    mvn clean package -P create-shade-jar -D skipTests -D accumulo.version="$ACCUMULO_VERSION" -D hadoop.version="$HADOOP_VERSION" -D zookeeper.version="$ZOOKEEPER_VERSION"
  fi
}

function determine_app_main() {
  ci_package="org.apache.accumulo.testing.core.continuous"
  case "$1" in 
    ingest)
      ci_main="${ci_package}.ContinuousIngest"
      ;;
    walk)
      ci_main="${ci_package}.ContinuousWalk"
      ;;
    batchwalk)
      ci_main="${ci_package}.ContinuousBatchWalker"
      ;;
    scan)
      ci_main="${ci_package}.ContinuousScanner"
      ;;
    *)
      echo "Unknown application: $1"
      print_usage
      exit 1    
  esac
}

function determine_mapred_main() {
  ci_package="org.apache.accumulo.testing.core.continuous"
  case "$1" in 
    verify)
      ci_main="${ci_package}.ContinuousVerify"
      ;;
    moru)
      ci_main="${ci_package}.ContinuousMoru"
      ;;
    *)
      echo "Unknown application: $1"
      print_usage
      exit 1    
  esac
}

function start_agitator() {
  hash pssh 2>/dev/null || { echo >&2 "The agitator requires pssh to be installed. Aborting."; exit 1; }

  mkdir -p "${at_home}/logs"
  log_base="${at_home}/logs/$(date +%Y%m%d%H%M%S)_$(hostname)"
  libexec="${at_home}/libexec"
  master_log="${log_base}_master-agitator"
  tserver_log="${log_base}_tserver-agitator"
  datanode_log="${log_base}_datanode-agitator"
  master_cmd="nohup ${libexec}/master-agitator.pl $AGTR_MASTER_KILL_SLEEP_TIME $AGTR_MASTER_RESTART_SLEEP_TIME"
  tserver_cmd="nohup ${libexec}/tserver-agitator.pl $AGTR_TSERVER_KILL_SLEEP_TIME $AGTR_TSERVER_RESTART_SLEEP_TIME $AGTR_TSERVER_MIN_KILL $AGTR_TSERVER_MAX_KILL"
  datanode_cmd="nohup ${libexec}/datanode-agitator.pl $AGTR_DATANODE_KILL_SLEEP_TIME $AGTR_DATANODE_RESTART_SLEEP_TIME $HADOOP_HOME $AGTR_DATANODE_MIN_KILL $AGTR_DATANODE_MAX_KILL"
  [[ -n $AGITATOR_USER ]] || AGITATOR_USER=$(whoami)

  if [[ $AGITATOR_USER == root ]];  then
    echo "Running master-agitator and tserver-agitator as $AGTR_ACCUMULO_USER using su. Running datanode-agitator as $AGTR_HDFS_USER using su."
    su -c "$master_cmd >${master_log}.out 2>${master_log}.err" & -m - "$AGTR_ACCUMULO_USER"
    su -c "$tserver_cmd >${tserver_log}.out 2>${tserver_log}.err" & -m - "$AGTR_ACCUMULO_USER"
    su -c "$datanode_cmd >${datanode_log}.out 2>${datanode_log}.err" & -m - "$AGTR_HDFS_USER"
  else
    if [[ $AGITATOR_USER == "$AGTR_ACCUMULO_USER" ]]; then
      echo "Running master-agitator and tserver-agitator as $AGITATOR_USER"
      $master_cmd > "${master_log}.out" 2> "${master_log}.err" &
      $tserver_cmd > "${tserver_log}.out" 2> "${tserver_log}.err" &
    else 
      echo "Running master-agitator and tserver-agitator as $AGTR_ACCUMULO_USER using sudo."
      sudo -u "$AGTR_ACCUMULO_USER" $master_cmd > "${master_log}.out" 2> "${master_log}.err" &
      sudo -u "$AGTR_ACCUMULO_USER" $tserver_cmd > "${tserver_log}.out" 2> "${tserver_log}.err" &
    fi
    if [[ $AGITATOR_USER == "$AGTR_HDFS_USER" ]]; then
      echo "Running datanode-agitator as $AGITATOR_USER"
      $datanode_cmd > "${datanode_log}.out" 2> "${datanode_log}.err" &
    else
      echo "Running datanode-agitator as $AGTR_HDFS_USER using sudo."
      sudo -u "$AGTR_HDFS_USER" $datanode_cmd > "${datanode_log}.out" 2> "${datanode_log}.err" &
    fi
  fi

  if ${AGTR_HDFS:-false} ; then
    agitator_log=${log_base}_hdfs-agitator
    sudo -u "$AGTR_HDFS_SUPERUSER" nohup "${libexec}/hdfs-agitator.pl" --sleep "${AGTR_HDFS_SLEEP_TIME}" --hdfs-cmd "${AGTR_HDFS_COMMAND}" --superuser "${AGTR_HDFS_SUPERUSER}" >"${agitator_log}.out" 2>"${agitator_log}.err" &
  fi
}

function stop_agitator() {
  [[ -n $AGITATOR_USER ]] || AGITATOR_USER=$(whoami)
  if [[ $AGITATOR_USER == root ]]; then
    echo "Stopping all processes matching 'agitator.pl' as root"
    pkill -f agitator.pl 2>/dev/null
  elif [[ $AGITATOR_USER == "$AGTR_ACCUMULO_USER" ]];  then
    echo "Stopping all processes matching 'datanode-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f datanode-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'hdfs-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f hdfs-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'agitator.pl' as $AGITATOR_USER"
    pkill -f agitator.pl 2>/dev/null 2>/dev/null
  else
    echo "Stopping all processes matching 'datanode-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f datanode-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'hdfs-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f hdfs-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'agitator.pl' as $AGTR_ACCUMULO_USER"
    sudo -u "$AGTR_ACCUMULO_USER" pkill -f agitator.pl 2>/dev/null
  fi
}

randomwalk_main="org.apache.accumulo.testing.core.randomwalk.Framework"
case "$1" in
ci-createtable)
  build_shade_jar
  java -Dlog4j.configuration="file:$log4j_config" org.apache.accumulo.testing.core.continuous.CreateTable "$at_props" "$ACCUMULO_CLIENT_PROPS"
  ;;
ci-local)
  if [ -z "$2" ]; then
    echo "ERROR: <application> needs to be set"
    print_usage
    exit 1
  fi
  determine_app_main "$2"
  build_shade_jar
  java -Dlog4j.configuration="file:$log4j_config" "$ci_main" "$at_props" "$ACCUMULO_CLIENT_PROPS"
  ;;
ci-yarn)
  if [ -z "$2" ]; then
    echo "ERROR: <num> needs to be set"
    print_usage
    exit 1
  fi
  if [ -z "$3" ]; then
    echo "ERROR: <application> needs to be set"
    print_usage
    exit 1
  fi
  determine_app_main "$3"
  build_shade_jar
  mvn compile -P yarn-test-runner -D hadoop.version="$HADOOP_VERSION" -D exec.args="-t AccumuloCITest-$3 -j $at_shaded_jar -m $ci_main -n $2 -p $at_props -c $ACCUMULO_CLIENT_PROPS -l $log4j_config -a ./accumulo-testing.properties ./accumulo-client.properties"
  ;;
ci-mapred)
  if [ -z "$2" ]; then
    echo "ERROR: <application> needs to be set"
    print_usage
    exit 1
  fi
  determine_mapred_main "$2"
  build_shade_jar
  "$HADOOP_HOME"/bin/yarn jar "$at_shaded_jar" "$ci_main" "$at_props" "$ACCUMULO_CLIENT_PROPS"
  ;;
rw-local)
  if [ -z "$2" ]; then
    echo "ERROR: <module> needs to be set"
    print_usage
    exit 1
  fi
  build_shade_jar
  java -Dlog4j.configuration="file:$log4j_config" "$randomwalk_main" "$at_props" "$ACCUMULO_CLIENT_PROPS" "$2"
  ;;
rw-yarn)
  if [ -z "$2" ]; then
    echo "ERROR: <num> needs to be set"
    print_usage
    exit 1
  fi
  if [ -z "$3" ]; then
    echo "ERROR: <module> needs to be set"
    print_usage
    exit 1
  fi
  build_shade_jar
  mvn compile -P yarn-test-runner -D hadoop.version="$HADOOP_VERSION" -D exec.args="-t AccumuloRWTest-$3 -j $at_shaded_jar -m $randomwalk_main -n $2 -p $at_props -c $ACCUMULO_CLIENT_PROPS -l $log4j_config -a ./accumulo-testing.properties ./accumulo-client.properties $3"
  ;;
agitator)
  case "$2" in
    start)
      start_agitator
      ;;
    stop)
      stop_agitator
      ;;
    *)
      echo "ERROR: unknown command - $2"
      print_usage
      exit 1
  esac
  ;;
*)
  echo "Unknown command: $1"
  print_usage
  exit 1
esac