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

set -o nounset
set -o errexit  # exit script if any command exits with nonzero value

readonly PROG_NAME=$(basename $0)
readonly PROG_DIR=$(dirname $(realpath $0))
readonly INVOKE_DIR=$(pwd)
readonly ARGS="$@"

# overrideable defaults
AWS=false
PARALLEL=true
MAX_PARALLEL=5
DEBUG=false

readonly USAGE="Usage: $PROG_NAME [-h | --help] [--aws [--no-parallel] [--max-parallel MAX]]"
readonly HELP="$(cat <<EOF
Tool to bring up a vagrant cluster on local machine or aws.

    -h | --help             Show this help message
    --aws                   Use if you are running in aws
    --no-parallel           Bring up machines not in parallel. Only applicable on aws
    --max-parallel  MAX     Maximum number of machines to bring up in parallel. Note: only applicable on test worker machines on aws. default: $MAX_PARALLEL
    --debug                 Enable debug information for vagrant
Approximately speaking, this wrapper script essentially wraps 2 commands:
    vagrant up
    vagrant hostmanager

The situation on aws is complicated by the fact that aws imposes a maximum request rate,
which effectively caps the number of machines we are able to bring up in parallel. Therefore, on aws,
this wrapper script attempts to bring up machines in small batches.

If you are seeing rate limit exceeded errors, you may need to use a reduced --max-parallel setting.

EOF
)"

function help {
    echo "$USAGE"
    echo "$HELP"
    exit 0
}

while [[ $# > 0 ]]; do
    key="$1"
    case $key in
        -h | --help)
            help
            ;;
        --aws)
            AWS=true
            ;;
        --no-parallel)
            PARALLEL=false
            ;;
        --max-parallel)
            MAX_PARALLEL="$2"
            shift
            ;;
        --debug)
            DEBUG=true
            ;;
        *)
            # unknown option
            echo "Unknown option $1"
            exit 1
            ;;
esac
shift # past argument or value
done

# Get a list of vagrant machines (in any state)
function read_vagrant_machines {
    local ignore_state="ignore"
    local reading_state="reading"
    local tmp_file="tmp-$RANDOM"

    local state="$ignore_state"
    local machines=""

    while read -r line; do
        # Lines before the first empty line are ignored
        # The first empty line triggers change from ignore state to reading state
        # When in reading state, we parse in machine names until we hit the next empty line,
        # which signals that we're done parsing
        if [[ -z "$line" ]]; then
            if [[ "$state" == "$ignore_state" ]]; then
                state="$reading_state"
            else
                # all done
                echo "$machines"
                return
            fi
            continue
        fi

        # Parse machine name while in reading state
        if [[ "$state" == "$reading_state" ]]; then
            line=$(echo "$line" | cut -d ' ' -f 1)
            if [[ -z "$machines" ]]; then
                machines="$line"
            else
                machines="${machines} ${line}"
            fi
        fi
    done < <(vagrant status)
}

# Filter "list", returning a list of strings containing pattern as a substring
function filter {
    local list="$1"
    local pattern="$2"

    local result=""
    for item in $list; do
        if [[ ! -z "$(echo $item | grep "$pattern")"  ]]; then
            result="$result $item"
        fi
    done
    echo "$result"
}

# Given a list of machine names, return only test worker machines
function worker {
    local machines="$1"
    local workers=$(filter "$machines" "worker")
    workers=$(echo "$workers" | xargs)  # trim leading/trailing whitespace
    echo "$workers"
}

# Given a list of machine names, return only zookeeper and broker machines
function zk_broker {
    local machines="$1"
    local zk_broker_list=$(filter "$machines" "zk")
    zk_broker_list="$zk_broker_list $(filter "$machines" "broker")"
    zk_broker_list=$(echo "$zk_broker_list" | xargs)  # trim leading/trailing whitespace
    echo "$zk_broker_list"
}

# Run a vagrant command on batches of machines of size $group_size
# This is annoying but necessary on aws to avoid errors due to AWS request rate
# throttling
#
# Example
#   $ vagrant_batch_command "vagrant up" "m1 m2 m3 m4 m5" "2"
#
#   This is equivalent to running "vagrant up" on groups of machines of size 2 or less, i.e.:
#   $ vagrant up m1 m2
#   $ vagrant up m3 m4
#   $ vagrant up m5
function vagrant_batch_command {
    local vagrant_cmd="$1"
    local machines="$2"
    local group_size="$3"

    local count=1
    local m_group=""
    # Using --provision flag makes this command useable both when bringing up a cluster from scratch,
    # and when bringing up a halted cluster. Permissions on certain directores set during provisioning
    # seem to revert when machines are halted, so --provision ensures permissions are set correctly in all cases
    for machine in $machines; do
        m_group="$m_group $machine"

        if [[ $(expr $count % $group_size) == 0 ]]; then
            # We've reached a full group
            # Bring up this part of the cluster
            $vagrant_cmd $m_group
            m_group=""
        fi
        ((count++))
    done

    # Take care of any leftover partially complete group
    if [[ ! -z "$m_group" ]]; then
        $vagrant_cmd $m_group
    fi
}

# We assume vagrant-hostmanager is installed, but may or may not be disabled during vagrant up
# In this fashion, we ensure we run hostmanager after machines are up, and before provisioning.
# This sequence of commands is necessary for example for bringing up a multi-node zookeeper cluster
function bring_up_local {
    vagrant up --no-provision
    vagrant hostmanager
    vagrant provision
}

function bring_up_aws {
    local parallel="$1"
    local max_parallel="$2"
    local machines="$(read_vagrant_machines)"
    case "$3" in
          true)
            local debug="--debug"
            ;;
          false)
            local debug=""
            ;;
    esac
    zk_broker_machines=$(zk_broker "$machines")
    worker_machines=$(worker "$machines")

    if [[ "$parallel" == "true" ]]; then
        if [[ ! -z "$zk_broker_machines" ]]; then
            # We still have to bring up zookeeper/broker nodes serially
            echo "Bringing up zookeeper/broker machines serially"
            vagrant up --provider=aws --no-parallel --no-provision $zk_broker_machines $debug
            vagrant hostmanager --provider=aws
            vagrant provision
        fi

        if [[ ! -z "$worker_machines" ]]; then
            echo "Bringing up test worker machines in parallel"
	    # Try to isolate this job in its own /tmp space. See note
	    # below about vagrant issue
            local vagrant_rsync_temp_dir=$(mktemp -d);
            TMPDIR=$vagrant_rsync_temp_dir vagrant_batch_command "vagrant up $debug --provider=aws" "$worker_machines" "$max_parallel"
            rm -rf $vagrant_rsync_temp_dir
            vagrant hostmanager --provider=aws
        fi
    else
        vagrant up --provider=aws --no-parallel --no-provision $debug
        vagrant hostmanager --provider=aws
        vagrant provision
    fi

    # Currently it seems that the AWS provider will always run rsync
    # as part of vagrant up. However,
    # https://github.com/mitchellh/vagrant/issues/7531 means it is not
    # safe to do so. Since the bug doesn't seem to cause any direct
    # errors, just missing data on some nodes, follow up with serial
    # rsyncing to ensure we're in a clean state. Use custom TMPDIR
    # values to ensure we're isolated from any other instances of this
    # script that are running/ran recently and may cause different
    # instances to sync to the wrong nodes
    for worker in $worker_machines; do
        local vagrant_rsync_temp_dir=$(mktemp -d);
        TMPDIR=$vagrant_rsync_temp_dir vagrant rsync $worker;
        rm -rf $vagrant_rsync_temp_dir
    done
}

function main {
    if [[ "$AWS" == "true" ]]; then
        bring_up_aws "$PARALLEL" "$MAX_PARALLEL" "$DEBUG"
    else
        bring_up_local
    fi
}

main