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

# This script will download, setup, start, and stop servers for Kafka, YARN, and ZooKeeper,

if [ -z "$JAVA_HOME" ]; then
  if [ -x /usr/libexec/java_home ]; then
    export JAVA_HOME="$(/usr/libexec/java_home)"
  else
    echo "JAVA_HOME not set. Exiting."
    exit 1
  fi
fi

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
BASE_DIR=$(dirname $DIR)
DEPLOY_ROOT_DIR=$BASE_DIR/deploy
DOWNLOAD_CACHE_DIR=$HOME/.grid/download
COMMAND=$1
SYSTEM=$2

DOWNLOAD_KAFKA=https://archive.apache.org/dist/kafka/0.10.2.1/kafka_2.11-0.10.2.1.tgz
DOWNLOAD_YARN=https://archive.apache.org/dist/hadoop/common/hadoop-2.6.1/hadoop-2.6.1.tar.gz
DOWNLOAD_ZOOKEEPER=http://archive.apache.org/dist/zookeeper/zookeeper-3.4.3/zookeeper-3.4.3.tar.gz

SERVICE_WAIT_TIMEOUT_SEC=20
ZOOKEEPER_PORT=2181
RESOURCEMANAGER_PORT=8032
NODEMANAGER_PORT=8042
KAFKA_PORT=9092

bootstrap() {
  echo "Bootstrapping the system..."
  stop_all
  rm -rf "$DEPLOY_ROOT_DIR"
  mkdir "$DEPLOY_ROOT_DIR"
  install_all
  start_all
  exit 0
}

standalone() {
  echo "Setting up the system..."
  stop_all
  rm -rf "$DEPLOY_ROOT_DIR"
  mkdir "$DEPLOY_ROOT_DIR"
  install_all_without_yarn
  start_all_without_yarn
  exit 0
}

install_all() {
  $DIR/grid install zookeeper
  $DIR/grid install yarn
  $DIR/grid install kafka
}

install_all_without_yarn() {
  $DIR/grid install zookeeper
  $DIR/grid install kafka
}

install_zookeeper() {
  mkdir -p "$DEPLOY_ROOT_DIR"
  install zookeeper $DOWNLOAD_ZOOKEEPER zookeeper-3.4.3
  cp "$DEPLOY_ROOT_DIR/zookeeper/conf/zoo_sample.cfg" "$DEPLOY_ROOT_DIR/zookeeper/conf/zoo.cfg"
}

install_yarn() {
  mkdir -p "$DEPLOY_ROOT_DIR"
  install yarn $DOWNLOAD_YARN hadoop-2.6.1
  cp "$BASE_DIR/conf/yarn-site.xml" "$DEPLOY_ROOT_DIR/yarn/etc/hadoop/yarn-site.xml"
  if [ ! -f "$HOME/.grid/conf/yarn-site.xml" ]; then
    mkdir -p "$HOME/.grid/conf"
    cp "$BASE_DIR/conf/yarn-site.xml" "$HOME/.grid/conf/yarn-site.xml"
  fi
}

install_kafka() {
  mkdir -p "$DEPLOY_ROOT_DIR"
  install kafka $DOWNLOAD_KAFKA kafka_2.11-0.10.2.1
  # have to use SIGTERM since nohup on appears to ignore SIGINT
  # and Kafka switched to SIGINT in KAFKA-1031.
  sed -i.bak 's/SIGINT/SIGTERM/g' $DEPLOY_ROOT_DIR/kafka/bin/kafka-server-stop.sh
  # in order to simplify the wikipedia-stats example job, set topic to have just 1 partition by default
  sed -i.bak 's/^num\.partitions *=.*/num.partitions=1/' $DEPLOY_ROOT_DIR/kafka/config/server.properties
}

install() {
  DESTINATION_DIR="$DEPLOY_ROOT_DIR/$1"
  DOWNLOAD_URL=$2
  PACKAGE_DIR="$DOWNLOAD_CACHE_DIR/$3"
  PACKAGE_FILE="$DOWNLOAD_CACHE_DIR/$(basename $DOWNLOAD_URL)"
  if [ -f "$PACKAGE_FILE" ]; then
    echo "Using previously downloaded file $PACKAGE_FILE"
  else
    echo "Downloading $(basename $DOWNLOAD_URL)..."
    mkdir -p $DOWNLOAD_CACHE_DIR
    curl "$DOWNLOAD_URL" > "${PACKAGE_FILE}.tmp"
    mv "${PACKAGE_FILE}.tmp" "$PACKAGE_FILE"
  fi
  rm -rf "$DESTINATION_DIR" "$PACKAGE_DIR"
  tar -xf "$PACKAGE_FILE" -C $DOWNLOAD_CACHE_DIR
  mv "$PACKAGE_DIR" "$DESTINATION_DIR"
}

start_all() {
  $DIR/grid start zookeeper
  $DIR/grid start yarn
  $DIR/grid start kafka
}

start_all_without_yarn() {
  $DIR/grid start zookeeper
  $DIR/grid start kafka
}

start_zookeeper() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/bin/zkServer.sh ]; then
    cd $DEPLOY_ROOT_DIR/$SYSTEM
    bin/zkServer.sh start
    wait_for_service "zookeeper" $ZOOKEEPER_PORT
    cd - > /dev/null
  else
    echo 'Zookeeper is not installed. Run: bin/grid install zookeeper'
  fi
}

start_yarn() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh ]; then
    $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh start resourcemanager
    wait_for_service "resourcemanager" $RESOURCEMANAGER_PORT
    $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh start nodemanager
    wait_for_service "nodemanager" $NODEMANAGER_PORT
  else
    echo 'YARN is not installed. Run: bin/grid install yarn'
  fi
}

start_kafka() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/bin/kafka-server-start.sh ]; then
    mkdir -p $DEPLOY_ROOT_DIR/$SYSTEM/logs
    cd $DEPLOY_ROOT_DIR/$SYSTEM
    nohup bin/kafka-server-start.sh config/server.properties > logs/kafka.log 2>&1 &
    cd - > /dev/null
    wait_for_service "kafka" $KAFKA_PORT
  else
    echo 'Kafka is not installed. Run: bin/grid install kafka'
  fi
}

wait_for_service() {
  local SERVICE_NAME=$1
  local PORT=$2
  echo "Waiting for $SERVICE_NAME to start..."
  local CURRENT_WAIT_TIME=0
  until $(nc -w 1 localhost $PORT); do
    printf '.'
    sleep 1
    if [ $((++CURRENT_WAIT_TIME)) -eq $SERVICE_WAIT_TIMEOUT_SEC ]; then
      printf "\nError: timed out while waiting for $SERVICE_NAME to start.\n"
      exit 1
    fi
  done
  printf '\n'
  echo "$SERVICE_NAME has started";
}

stop_all() {
  $DIR/grid stop kafka
  $DIR/grid stop yarn
  $DIR/grid stop zookeeper
}

stop_zookeeper() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/bin/zkServer.sh ]; then
    cd $DEPLOY_ROOT_DIR/$SYSTEM
    bin/zkServer.sh stop
    cd - > /dev/null
  else
    echo 'Zookeeper is not installed. Run: bin/grid install zookeeper'
  fi
}

stop_yarn() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh ]; then
    $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh stop resourcemanager
    $DEPLOY_ROOT_DIR/$SYSTEM/sbin/yarn-daemon.sh stop nodemanager
  else
    echo 'YARN is not installed. Run: bin/grid install yarn'
  fi
}

stop_kafka() {
  if [ -f $DEPLOY_ROOT_DIR/$SYSTEM/bin/kafka-server-stop.sh ]; then
    cd $DEPLOY_ROOT_DIR/$SYSTEM
    bin/kafka-server-stop.sh || true # tolerate nonzero exit status if Kafka isn't running
    cd - > /dev/null
  else
    echo 'Kafka is not installed. Run: bin/grid install kafka'
  fi
}

# Check arguments
if [ "$COMMAND" == "bootstrap" ] && test -z "$SYSTEM"; then
  bootstrap
  exit 0
elif [ "$COMMAND" == "standalone" ] && test -z "$SYSTEM"; then
  standalone
  exit 0
elif (test -z "$COMMAND" && test -z "$SYSTEM") \
  || ( [ "$COMMAND" == "help" ] || test -z "$COMMAND" || test -z "$SYSTEM"); then
  echo
  echo "  Usage.."
  echo
  echo "  $ grid"
  echo "  $ grid bootstrap"
  echo "  $ grid standalone"
  echo "  $ grid install [yarn|kafka|zookeeper|all]"
  echo "  $ grid start [yarn|kafka|zookeeper|all]"
  echo "  $ grid stop [yarn|kafka|zookeeper|all]"
  echo
  exit 1
else
  echo "EXECUTING: $COMMAND $SYSTEM"

  "$COMMAND"_"$SYSTEM"
fi