#!/usr/bin/env bash

###################################################################
# About :
#
# axiom-fleet lets you spin up fleets of axiom instances in one or multiple regions.
# You can specify the name of a fleet (fleet prefix) or have axiom choose for you.
#
# Examples: 
#
# axiom-fleet # Spin up three instances, let axiom decide on the fleet prefix"
# axiom-fleet javis -i 10 # Spin up 10 instances with a fleet prefix of javis, this will create 10 instances named javis01 to javis10."
# axiom-fleet jerry -i 25 --regions dal13,lon06,fra05,tok05,syd05 # Spin up 25 instances using round robbin region distribution"

###########################################################################################################
# Header
#
AXIOM_PATH="$HOME/.axiom"
source "$AXIOM_PATH/interact/includes/vars.sh"
source "$AXIOM_PATH/interact/includes/functions.sh"
source "$AXIOM_PATH/interact/includes/system-notification.sh"
BASEOS="$(uname)"
case $BASEOS in
'Darwin')
    PATH="$(brew --prefix coreutils)/libexec/gnubin:$PATH"
    ;;
*) ;;
esac

###########################################################################################################
# Declare defaut variables
#
spend=false
hours=false
amount=false
prompt=false
region=false
image=false
cycle_regions=false
region_array=()
gen_name=""
provider="$(jq -r ".provider" "$AXIOM_PATH/axiom.json")"
image_name="$(jq -r '.imageid' "$AXIOM_PATH"/axiom.json)"
region="$(jq -r '.region' "$AXIOM_PATH"/axiom.json)"
image_id=""
instances=false
time=320
init_timeout=4
init_sleep=4

###########################################################################################################
# DO Region Transfer
# Transfer image to region if requested in that region yet does not exist. DO only
#
region_transfer(){
if [[ "$provider" == "do" ]]; then
 avail_image_id_regions=$(doctl compute image get "$image_id" -o json| jq -r '.[] | .regions[]')
 requested_image_id_regions="$regionargs"
 if [[ "$avail_image_id_regions" != *"$requested_image_id_regions"* ]]; then
  echo -e "${Color_Off}]"
  echo -e "${BYellow}instance ${Blue}"$name"${BYellow} requested image in region ${BRed}$regionargs${BYellow}, however image ${BRed}$image_name${BYellow} only exists in ${BRed}$(echo $avail_image_id_regions | tr '\n' ',')"
  echo -e "${BYellow}axiom will attempt to transfer image to new region, this can take a few minutes please be patient...${Color_Off}"
  doctl compute image-action transfer $image_id --region $regionargs --wait
   if [ $? -eq 0 ]; then
    echo -e "\033[32mimage transfer succeeded\033[0m"
    echo -e "${BWhite}waiting 90 seconds before continuing...${Color_Off}"
    sleep 90
   else
    echo -e "${BRed}image transfer failed... you might need to pick a different region. '${BWhite}axiom-region ls${BRed}' to list regions. '${BWhite}axiom-images ls${BRed}' to list images${Color_Off}"
   fi
 echo -n -e "${BWhite}Instances: ${Color_Off}[ ${Blue}"
 fi
echo -n -e ${Blue}
fi 
}

###########################################################################################################
# Help Menu
# 
function usage() {
        echo -e "${BWhite}Description:"
        echo -e "  Spin up fleets of axiom instances in one or multiple regions."
        echo -e "  Specify the name of your fleet (fleet prefix) or have axiom choose for you."
        echo -e "${BWhite}Examples:${Color_Off}"
        echo -e "  ${Blue}axiom-fleet${Color_Off} # Spin up three instances, let axiom decide on the fleet prefix"
        echo -e "  ${Blue}axiom-fleet -i 10${Color_Off} # Spin up 10 instances with random fleet prefix"
        echo -e "  ${Blue}axiom-fleet jerry -i 25 --regions dal13,lon06,fra05,tok05,syd05${Color_Off} # Spin up 25 instances named jerry01 to jerry25 using Round-robin region distribution"
        echo -e "${BWhite}Usage:${Color_Off}"
        echo -e "  -i/--instances <integer>"
        echo -e "    The number of instances to spin up"
        echo -e "  -r/--regions <regions> (optional)"
        echo -e "    Supply comma-separated regions to cycle through (default is region in ~/.axiom/axiom.json)"
        echo -e "  --debug (optional)"            
        echo -e "    Run with set -xv, warning: very verbose"
        echo -e "  --help (optional)"
        echo -e "    Display this help menu"
}

###########################################################################################################
# Parse command line arguments 
#
i=0
for arg in "$@"
do
    i=$((i+1))
    if [[  ! " ${pass[@]} " =~ " ${i} " ]]; then
        set=false
        if [[ "$i" == 1 ]]; then
            instance="$1"
            set=true
            pass+=($i)
        fi
        if [[ "$arg" == "--debug" ]]; then
            set -xv
            set=true
            pass+=($i)
        fi
        if [[ "$arg" == "-i" ]] || [[ "$arg" == "--instances" ]]; then
            n=$((i+1))
            instances=true
            amount=$(echo ${!n})
            set=true
            pass+=($i)
            pass+=($n)
        fi
        if [[ "$arg" == "--regions" ]] || [[ "$arg" == "-r" ]]; then
            n=$((i+1))
            cycle_regions=$(echo ${!n})
            set=true
            pass+=($i)
            pass+=($n)
        fi
        if [[ "$arg" == "--image" ]] ; then
            n=$((i+1))
            image=$(echo ${!n})
            set=true
            pass+=($i)
            pass+=($n)
        fi
        if [[ "$arg" == "--help" ]] || [[ "$arg" == "-h" ]] || [[ "$arg" == "help" ]]; then
            usage
            exit
            set=true
            pass+=($i)
        fi
        if  [[ "$set" != "true" ]]; then
            args="$args $arg"
        fi
    fi
done

###########################################################################################################
# Display Help Menu
#
if [[ "$*" == "--help" ]] || [[ "$*" == "-h" ]] ; then
usage
exit
fi

###########################################################################################################
# If -i /--instances isnt used, default to three instances
#
if [[ "$amount" == "false" ]]; then
 amount=3
fi

###########################################################################################################
# Generate name TODO clean up
#
if [ -z "$1" ] || [[ $1 =~ "-i" ]] || [[ $1 =~ "--instances" ]]|| [[ $1 =~ "-r" ]] || [[ $1 =~ "--regions" ]] || [[ $1 =~ "--image" ]] ; then                                                
 gen_name="${names[$RANDOM % ${#names[@]} ]}"
else
 if [ ! -z "$1" ]; then
  gen_name="$1"
 fi
fi

###########################################################################################################
# Change init_sleep as needed
#
if [[ "$provider" == "linode" ]]; then
 init_sleep=6
fi

###########################################################################################################
# Initialize fleet
#
if [[ "$cycle_regions" == "false" ]]; then

# Chance to cancel initialization of fleet
#
echo -e "${BWhite}Initializing new fleet '$gen_name' with $amount instances...${Color_Off}"
echo -e "${BWhite}INITIALIZING IN 5 SECONDS, CTRL+C to quit... ${Color_Off}"
sleep 5

image_id=$(get_image_id "$image_name")
total=$(query_instances "$gen_name*" | tr " " "\n" | sed 's/[^0-9]*//g'| sort -nr | head -n1)
total="${total#0}"
start="${start#0}"
start=$((total))
amount=$(($amount+$start))
start=$((start+1))
total_spend_per_instance_rounded="0"
slug=$(cat "$AXIOM_PATH"/axiom.json | jq -r .default_size)  >/dev/null 2>&1

echo -n -e "${BWhite}Instances: ${Color_Off}[ ${Blue}"
o=0
for i in $(seq -f "%02g" $start $amount)
do
 time=$((time+3))
 name="$gen_name$i"
 echo -n -e "$name "
 args=""
"$AXIOM_PATH"/interact/axiom-init "$name" --quiet --size "$slug" --image-id "$image_id" --no-select --region  "$region" &
sleep $init_sleep
o=$((o+1))
done
echo -n -e "${Color_Off} ]\n"

while [[ $time -gt 0 ]]; do
    echo -ne ">> T-Minus $time to fleet $gen_name initialization...\033[0K\r"
    sleep 1
    : $((time--))
done
fi

###########################################################################################################
# Round Robin distribution logic here ( i.e. regions_to_cycle)
#
if [[ "$cycle_regions" != "false" ]]; then

###########################################################################################################
# Chance to cancel initialization of fleet
#
echo -e "${BWhite}Initializing new fleet '$gen_name' with $amount instances...${Color_Off}"
echo -e "${BWhite}Cycling through following regions:$cycle_regions...${Color_Off}"
echo -e "${BWhite}INITIALIZING IN 5 SECONDS, CTRL+C to quit... ${Color_Off}"
sleep 5
image_id=$(get_image_id "$image_name")
total=$(query_instances "$gen_name*" | tr " " "\n" | sed 's/[^0-9]*//g'| sort -nr | head -n1)
total="${total#0}"
start="${start#0}"
start=$((total))
amount=$(($amount+$start))
start=$((start+1))
total_spend_per_instance_rounded="0"
slug=$(cat "$AXIOM_PATH"/axiom.json | jq -r .default_size)  >/dev/null 2>&1

IFS=',' read -r -a region_array <<< "$cycle_regions"

# Basically repeat items, if the fleet is 20 hosts and they only supply 3 regions, loop over them until it reaches 20
total_regions=$(echo "${region_array[@]}" |  tr ' ' '\n' | wc -l | awk '{ print $1 }')
regions_to_cycle=()
 k=0
 while [[ "$(echo ${regions_to_cycle[@]} | tr ' ' '\n' | wc -l | awk '{ print $1}')" -lt "$amount" ]]
 do
  regions_to_cycle+=("${region_array[k]}")
  if [[ "$k" -lt "$total_regions" ]]; then
   k=$((k+1))
  else
   k=0
  fi
done

# Remove null element from array and reindex
#
for i in "${!regions_to_cycle[@]}"; do
  [ -n "${regions_to_cycle[$i]}" ] || unset "regions_to_cycle[$i]" && regions_to_cycle=( "${regions_to_cycle[@]}" )
done

echo -n -e "${BWhite}Instances: ${Color_Off}[ ${Blue}"
o=0
for i in $(seq -f "%02g" $start $amount)
do
time=$((time+3))
name="$gen_name$i"
echo -n -e "$name "
args=""
regionargs="${regions_to_cycle[o]}"

# final check to make sure region is set
#
if [ -z ${regionargs:+x} ]; then
regionargs="$(jq -r '.region' "$AXIOM_PATH"/axiom.json)"
fi

###########################################################################################################
# DO Region Transfer
# Transfer image to region if requested in that region yet does not exist. DO only
#
if [[ "$provider" == "do" ]]; then
  region_transfer
fi

# create instance
#
"$AXIOM_PATH"/interact/axiom-init "$name" --quiet --size "$slug" --image-id "$image_id" --no-select --region  "$regionargs" &
sleep $init_sleep
o=$((o+1))
done
echo -n -e "${Color_Off} ]\n"
fi

while [[ $time -gt 0 ]]; do
 echo -ne ">> T-Minus $time to fleet $gen_name initialization...\033[0K\r"
 sleep 1
 : $((time--))
done
"$AXIOM_PATH"/interact/axiom-select "$gen_name*"

echo -e "${BGreen}Fleet started succesfully!\nTo delete your fleet, just run '${Blue}axiom-rm \"$gen_name*\" -f${BGreen}'${Color_Off}"
