#!/bin/bash

set -e

BOLD_RED='\033[1;31m'
GREEN='\033[32m'
BOLD_PURPLE='\033[1;35m'
NO_COLOR='\033[0m'

error() {
  echo -e "${BOLD_RED}>>> $1${NO_COLOR}"
}

info() {
  echo -e "${GREEN}>>> $1${NO_COLOR}"
}

error_exit() {
  error "$1"
  exit 1
}

PROVIDERS=("aws" "vb")

usage() {
  echo -e "${BOLD_RED}>>> Usage: $0 [Machine Number] [Provider] [Options]${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tMachine Number:\t Must be at least 2${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tProvider:\t ${PROVIDERS[@]}${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tOptions:${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \t -a, --alias-ip\t Add aws IP aliases(e.g., TachyonMaster) into /etc/hosts${NO_COLOR}"
  exit
}

if [ "$#" -lt 2 ]; then
  usage
fi

TOTAL="$1"
[[ "$TOTAL" -lt 2 ]] && error_exit "Machine number must be at least 2"
shift

TACHYON_MASTERS=$(grep Masters conf/tachyon.yml | cut -d':' -f 2 | tr -d ' ')
[[ "$TOTAL" -lt "$TACHYON_MASTERS" ]] && error_exit "Total number of machines must be at least the same with Masters in conf/tachyon.yml"

PROVIDER="$1"
[[ "${PROVIDERS[*]}" =~ (^| )$PROVIDER( |$) ]] || error_exit "Provider $PROVIDER is not supported yet"
shift

ALIAS_IP="false"
while [[ $# > 0 ]]; do
  key="$1"
  case $key in
    -a|--alias-ip)
      ALIAS_IP="true"
    ;;
    *)
    usage
    ;;
  esac
shift
done

HERE=`dirname "$0"`
cd $HERE

# make sure ec2.yml exists
if [[ "$PROVIDER" == "aws" ]]; then
  if [[ ! -f conf/ec2.yml ]]; then
    error_exit "Please create your own conf/ec2.yml from conf/ec2.yml.template"
  fi
fi

cat > conf/init.yml <<EOF
# Generated by ../start, do not modify
Provider: $PROVIDER
MachineNumber: $TOTAL
EOF

mkdir -p ./files
# files/id_rsa, files/id_rsa.pub
if [[ ! -f ./files/id_rsa ]]; then
  ssh-keygen -f ./files/id_rsa -t rsa -N '' &>/dev/null
fi

# files/workers, order matters, 
# when in Tachyon fault tolerant mode, the last one is often also chosen as Tachyon leader(not sure), 
# when not in fault tolerant mode, the last one is Tachyon master,
# currently, for other systems like HDFS, fault tolerant mode hasn't been enabled, the last one is the single master.
rm -f files/workers
for i in `seq 1 $((TOTAL-TACHYON_MASTERS))`; do
  echo "TachyonWorker${i}" >> files/workers
done
for ((i=2; i<=TACHYON_MASTERS; i++)); do
  echo "TachyonMaster$i" >> files/workers
done
echo "TachyonMaster" >> files/workers

# create virtual machine without provisioning
[[ "$PROVIDER" == "aws" ]] && sgid=$(python bin/init_aws.py | tail -n 1)
# pass security group id to vagrant up
export AWS_SECURITY_GROUP_ID_TACH=$sgid
if [[ "$PROVIDER" == "aws" && "$(python bin/is_aws_spot.py)" == "0" ]]; then
  # register handler to be run when exit 1
  trap "python bin/spot_request.py --cancel" ERR
  python bin/spot_request.py --submit
else
  FULL_PROVIDER_NAME=$PROVIDER
  if [[ "$FULL_PROVIDER_NAME" == "vb" ]]; then
    FULL_PROVIDER_NAME=virtualbox
    ./vbox/buildbox.sh
  fi
  VUP="vagrant up --provider=$FULL_PROVIDER_NAME --no-provision"
  [[ "$PROVIDER" == "aws" ]] && VUP="${VUP} --parallel"
  $VUP
fi

# files/hosts
# IPs are automatically assigned by vagrant
# hack to find out the assigned IPs by looking into vagrant's ssh config to
# output of "ifconfig" is different in RHEL6.5 and RHEL7
#   inet:addr 127.0.0.1  VS. inet 127.0.0.1
# use "ip addr show" instead
rm -f files/hosts
for host in $(cat files/workers); do
  info "looking up IP of $host..."
  IP=$(vagrant ssh $host -c "ip addr show | grep -w inet | tail -1" | grep -w inet | awk '{print $2}' | cut -d'/' -f1)
  info "IP for $host is $IP"
  echo "$IP ${host}" >> files/hosts
done

# provision
vagrant provision

# IP aliasing
# For AWS, file ".vagrant/provisioners/ansible/inventory/vagrant_ansible_inventory"
# contains the hostname and IP information in the following format:
#
#   TachyonWorker1 ansible_ssh_host=54.148.173.226 ansible_ssh_port=22
#   TachyonMaster ansible_ssh_host=54.201.87.214 ansible_ssh_port=22
#
# The IP Aliasing is done by parsing this file to generate the "host ip" mapping

ANSIBLE_INFO_FILE=".vagrant/provisioners/ansible/inventory/vagrant_ansible_inventory"

if [[ "$ALIAS_IP" == "true" ]]; then
  TMP_FILE=./tachyon_hosts_tmp_file
  if [[ -f $ANSIBLE_INFO_FILE ]]; then
    echo -e "\n# Generated by Tachyon vagrant deploy script" > $TMP_FILE
    if [[ "$PROVIDER" == "aws" ]]; then
      regx="^Tachyon.*$"
      while read -r; do
        if [[ $REPLY =~ $regx ]]; then
          host=$(echo $REPLY | cut -d' ' -f1)
          ip=$(echo $REPLY | cut -d' ' -f2 | cut -d'=' -f2)
          echo "$host $ip" >>$TMP_FILE
        fi
      done < $ANSIBLE_INFO_FILE
    fi
    info "Adding following lines into /etc/hosts:"
    cat $TMP_FILE
    info ""
    sudo "$BASH" -c "sudo cat tachyon_hosts_tmp_file >> /etc/hosts"
    rm -f $TMP_FILE
  fi
fi

# show master IP
set +e # turn off the option that exits immediately when a command fails
if [[ "$PROVIDER" == "aws" ]]; then
  masters=$(cat $ANSIBLE_INFO_FILE | grep Master | cut -d' ' -f2 | cut -d'=' -f2)
  main_master=$(cat $ANSIBLE_INFO_FILE | grep 'TachyonMaster ' | cut -d' ' -f2 | cut -d'=' -f2)
else
  masters=$(cat files/hosts | grep Master | cut -d' ' -f1)
  main_master=$(cat files/hosts | grep 'TachyonMaster ' | cut -d' ' -f1)
fi
masters=($(echo $masters | tr '\n' ' ' | tr -d '\r'))

if [[ "$TACHYON_MASTERS" -gt 1 ]]; then
  info "Tachyon is in Fault Tolerant mode, waiting 5 seconds for leader to be selected and web UI to be available"
  sleep 5
  for master in ${masters[@]}; do
    info "Tachyon master: ${master}"
    info "Trying to curl Tachyon web UI at ${master}:19999"
    curl "${master}:19999" >/dev/null
    if [[ "$?" == "0" ]]; then
      info "Master public IP for Tachyon is ${master}, visit ${master}:19999 for Tachyon web UI"
      break
    fi
  done
  info "Master public IP for other softwares is ${main_master}"
else
  info "TachyonMaster public IP is ${main_master}, visit ${main_master}:19999 for Tachyon web UI"
fi
info "visit default port of the web UI of what you deployed"
