#! /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.

function print_usage {
  cat <<EOF
Usage: accumulo-cluster <command> (<argument> ...)

Commands:
  create-config       Creates cluster config
  restart             Restarts the Accumulo cluster
  start               Starts Accumulo cluster
  stop                Stops Accumulo cluster
  kill                Kills Accumulo cluster
  start-non-tservers  Starts all services except tservers
  start-tservers      Starts all tservers on cluster
  stop-tservers       Stops all tservers on cluster
  start-here          Starts all services on this node
  stop-here           Stops all services on this node
EOF
  exit 1
}

function invalid_args {
  echo -e "Invalid arguments: $1\n"
  print_usage 1>&2
  exit 1
}

function verify_config {
  if [[ -f ${conf}/slaves ]]; then
    echo "ERROR: A 'slaves' file was found in ${conf}/"
    echo "Accumulo now reads tablet server hosts from 'tservers' and requires that the 'slaves' file not be present to reduce confusion."
    echo "Please rename the 'slaves' file to 'tservers' or remove it if both exist."
    exit 1
  fi

  if [[ ! -f ${conf}/tservers ]]; then
    echo "ERROR: A 'tservers' file was not found at ${conf}/tservers"
    echo "Please make sure it exists and is configured with tablet server hosts."
    exit 1
  fi

  unset master1
  if [[ -f "${conf}/masters" ]]; then
    master1=$(egrep -v '(^#|^\s*$)' "${conf}/masters" | head -1)
  fi

  if [[ -z "${monitor}" ]] ; then
    monitor=$master1
    if [[ -f "${conf}/monitor" ]]; then
      monitor=$(egrep -v '(^#|^\s*$)' "${conf}/monitor" | head -1)
    fi
    if [[ -z "${monitor}" ]] ; then
      echo "Could not infer a Monitor role. You need to either define \"${conf}/monitor\"," 
      echo "or make sure \"${conf}/masters\" is non-empty."
      exit 1
    fi
  fi
  if [[ ! -f "${conf}/tracers" ]]; then
    if [[ -z "${master1}" ]] ; then
      echo "Could not find a master node to use as a default for the tracer role."
      echo "Either set up \"${conf}/tracers\" or make sure \"${conf}/masters\" is non-empty."
      exit 1
    else
      echo "$master1" > "${conf}/tracers"
    fi
  fi
  if [[ ! -f "${conf}/gc" ]]; then
    if [[ -z "${master1}" ]] ; then
      echo "Could not infer a GC role. You need to either set up \"${conf}/gc\" or make sure \"${conf}/masters\" is non-empty."
      exit 1
    else
      echo "$master1" > "${conf}/gc"
    fi
  fi
}

function get_ip() {
  ip_addr=$(ip addr | grep 'state UP' -A2 | tail -n1 | awk '{print $2}' | cut -f1  -d'/')
  if [[ $? != 0 ]]; then
    ip_addr=$(getent ahosts "$(hostname -f)" | grep DGRAM | cut -f 1 -d ' ')
  fi
  echo "$ip_addr"
}

function start_service() {
  host="$1"
  service="$2"

  if [[ $host == "localhost" || $host == $(hostname -f) || $host == $(hostname -s) || $host == $(get_ip) ]]; then
    "${bin}/accumulo-service" "$service" start
  else
    $SSH "$host" "bash -c '${bin}/accumulo-service \"$service\" start'"
  fi
}

function start_tservers() {
  echo -n "Starting tablet servers ..."
  count=1
  for server in $(egrep -v '(^#|^\s*$)' "${conf}/tservers"); do
    echo -n "."
    start_service "$server" tserver &
    if (( ++count % 72 == 0 )) ;
    then
      echo
      wait
    fi
  done
  echo " done"
}

function start_all() {
  unset DISPLAY

  start_service "$monitor" monitor 

  if [ "$1" != "--no-tservers" ]; then
    start_tservers
  fi

  for host in $(egrep -v '(^#|^\s*$)' "${conf}/masters"); do
    start_service "$host" master
  done

  for host in $(egrep -v '(^#|^\s*$)' "${conf}/gc"); do
    start_service "$host" gc
  done

  for host in $(egrep -v '(^#|^\s*$)' "${conf}/tracers"); do
    start_service "$host" tracer
  done
}

function start_here() {

  local_hosts="$(hostname -a 2> /dev/null) $(hostname) localhost 127.0.0.1 $(get_ip)"
  for host in $local_hosts; do
    if grep -q "^${host}\$" "${conf}/tservers"; then
      start_service "$host" tserver
      break
    fi
  done

  for host in $local_hosts; do
    if grep -q "^${host}\$" "${conf}/masters"; then
      start_service "$host" master
      break
    fi
  done

  for host in $local_hosts; do
    if grep -q "^${host}\$" "${conf}/gc"; then
      start_service "$host" gc
      break
    fi
  done

  for host in $local_hosts; do
    if [ "$host" == "$monitor" ]; then
      start_service "$host" monitor 
      break
    fi
  done

  for host in $local_hosts; do
    if grep -q "^${host}\$" "${conf}/tracers"; then
      start_service "$host" tracer 
      break
    fi
  done
}

function end_service() {
  host="$1"
  service="$2"
  end_cmd="$3"
  if [[ $host == localhost || $host = "$(hostname -s)" || $host = "$(hostname -f)" || $host = $(get_ip) ]] ; then
    "${bin}/accumulo-service" "$service" "$end_cmd"
  else
    $SSH "$host" "bash -c '${bin}/accumulo-service \"$service\" \"$end_cmd\"'"
  fi
}

function stop_service() {
  end_service "$1" "$2" "stop"
}

function kill_service() {
  end_service "$1" "$2" "kill"
}

function stop_tservers() {
  tserver_hosts=$(egrep -v '(^#|^\s*$)' "${conf}/tservers")

  echo "Stopping unresponsive tablet servers (if any)..."
  for host in ${tserver_hosts}; do
    stop_service "$host" tserver & 
  done

  sleep 10

  echo "Stopping unresponsive tablet servers hard (if any)..."
  for host in ${tserver_hosts}; do
    kill_service "$host" tserver & 
  done

  echo "Cleaning tablet server entries from zookeeper"
  ${accumulo_cmd} org.apache.accumulo.server.util.ZooZap -tservers
}

function kill_all() {
  echo "Killing Accumulo cluster..."

  for master in $(grep -v '^#' "${conf}/masters"); do
    kill_service "$master" master
  done

  for gc in $(grep -v '^#' "${conf}/gc"); do
    kill_service "$gc" gc
  done

  kill_service "$monitor" monitor

  for tracer in $(egrep -v '(^#|^\s*$)' "${conf}/tracers"); do
    kill_service "$tracer" tracer
  done

  for host in $(egrep -v '(^#|^\s*$)' "${conf}/tservers"); do
    kill_service "$host" tserver
  done

  echo "Cleaning all server entries in ZooKeeper"
  ${accumulo_cmd} org.apache.accumulo.server.util.ZooZap -master -tservers -tracers
}


function stop_all() {
  echo "Stopping Accumulo cluster..."
  if ! ${accumulo_cmd} admin stopAll
  then
    echo "Invalid password or unable to connect to the master"
    echo "Initiating forced shutdown in 15 seconds (Ctrl-C to abort)"
    sleep 10
    echo "Initiating forced shutdown in  5 seconds (Ctrl-C to abort)"
  else
    echo "Accumulo shut down cleanly"
    echo "Utilities and unresponsive servers will shut down in 5 seconds (Ctrl-C to abort)"
  fi

  sleep 5

  # Look for processes not killed by 'admin stopAll'
  for end_cmd in "stop" "kill" ; do
    for master in $(grep -v '^#' "${conf}/masters"); do
      end_service "$master" master $end_cmd
    done

    for gc in $(grep -v '^#' "${conf}/gc"); do
      end_service "$gc" gc $end_cmd
    done

    end_service "$monitor" monitor $end_cmd

    for tracer in $(egrep -v '(^#|^\s*$)' "${conf}/tracers"); do
      end_service "$tracer" tracer $end_cmd
    done
  done

  # stop tserver still running
  stop_tservers

  echo "Cleaning all server entries in ZooKeeper"
  ${accumulo_cmd} org.apache.accumulo.server.util.ZooZap -master -tservers -tracers
}

function stop_here() {
  # Determine hostname without errors to user
  hosts_to_check=($(hostname -a 2> /dev/null | head -1) $(hostname -f))

  if egrep -q localhost\|127.0.0.1 "${conf}/tservers"; then
    ${accumulo_cmd} admin stop localhost
  else
    for host in "${hosts_to_check[@]}"; do
      if grep -q "$host" "${conf}/tservers"; then
        ${accumulo_cmd} admin stop "$host"
      fi
    done
  fi

  for host in "${hosts_to_check[@]}"; do
    for end_cmd in "stop" "kill"; do
      for svc in tserver gc master monitor tracer; do
        end_service "$host" $svc $end_cmd
      done
    done
  done
}

function main() {

  if [[ -z $1 ]]; then
    invalid_args "<command> cannot be empty"
  fi

  # Resolve base directory
  SOURCE="${BASH_SOURCE[0]}"
  while [ -h "${SOURCE}" ]; do
     bin="$( cd -P "$( dirname "${SOURCE}" )" && pwd )"
     SOURCE="$(readlink "${SOURCE}")"
     [[ "${SOURCE}" != /* ]] && SOURCE="${bin}/${SOURCE}"
  done
  bin="$( cd -P "$( dirname "${SOURCE}" )" && pwd )"
  basedir=$( cd -P "${bin}"/.. && pwd )
  conf="${ACCUMULO_CONF_DIR:-${basedir}/conf}"

  accumulo_cmd="${bin}/accumulo"
  SSH='ssh -qnf -o ConnectTimeout=2'

  case "$1" in
    create-config)
      echo "localhost" > "$conf/gc"
      echo "localhost" > "$conf/masters"
      echo "localhost" > "$conf/monitor"
      echo "localhost" > "$conf/tracers"
      echo "localhost" > "$conf/tservers"
      ;;
    restart)
      verify_config
      stop_all
      # Make sure the JVM has a chance to fully exit
      sleep 1
      start_all
      ;;
    start)
      verify_config
      start_all
      ;;
    stop)
      verify_config
      stop_all
      ;;
    kill)
      verify_config
      kill_all
      ;;
    start-non-tservers)
      verify_config
      start_all --no-tservers
      ;;
    start-tservers)
      verify_config
      start_tservers
      ;;
    start-here)
      verify_config
      start_here
      ;;
    stop-tservers)
      verify_config
      stop_tservers
      ;;
    stop-here)
      verify_config
      stop_here
      ;;
    *)
      invalid_args "'$1' is an invalid <command>"
      ;;
  esac
}

main "$@"