#!/usr/bin/env bash
set -e;

cd "$(dirname "$(readlink -f "$0")")" || { echo "Error: Failed to cd to script directory" >&2; exit 1; };

imagename="${1}";

if [[ -z "${NO_COLOR}" ]]; then
	RS="\033[0m"    # reset
	HC="\033[1m"    # hicolor
	
	FRED="\033[31m" # foreground red
	FGRN="\033[32m" # foreground green
	FYEL="\033[33m" # foreground yellow
	FBLE="\033[34m" # foreground blue
fi

help() {
	echo -e "${FGRN}${HC}nomad_smartrestart${RS}, ${HC}v0.1${RS}
	By Starbeamrainbowlabs

${FBLE}${HC}USAGE${RS}
	path/to/nomad_smartrestart.sh {IMAGE_NAME}

${FBLE}${HC}ENVIRONMENT VARIABLES${RS}
	${FYEL}URL_NOMAD${RS}	The URL of the nomad server to talk to (using the Consul domain instead of a direct IP is recommended)"
}

log_msg() {
	echo -e "${HC}>>> nomad_smartrestart: ${RS}$*" >&2;
}
log_error() {
	echo -e "${HC}>>> nomad_smartrestart: ${FRED}$*${RS}";
}

if [[ -z "${imagename}" ]]; then
	log_error "Error: No image name specified." >&2;
	help;
	exit 1;
fi

loc_nomad_txt="./images/${imagename}/nomad.txt";

if [[ ! -f "${loc_nomad_txt}" ]]; then
	log_msg "No nomad.txt present for ${HC}${FGRN}${imagename}${RS}, nothing to do";
	help;
	exit 0;
fi


URL_NOMAD="${URL_NOMAD:-http://nomad.service.mooncarrot.space:4646}";

###############################################################################


# Finds allocations for a given job / task name combo.
# $1	The job name
# $2	The task name
nomad_allocs_find() {
	local jobname="${1}"
	local taskname="${2}"
	
	if [[ -z "${taskname}" ]]; then
		curl -sS "${URL_NOMAD}/v1/allocations" | JOB_NAME="${jobname}" jq --raw-output '.[] | select(.JobID == env.JOB_NAME) | .ID'
	else
		curl -sS "${URL_NOMAD}/v1/allocations" | JOB_NAME="${jobname}" TASK_NAME="${taskname}" jq --raw-output '.[] | select(.JobID == env.JOB_NAME and (.TaskStates|keys|.[]|. == env.TASK_NAME)) | .ID'
	fi
}

# Restarts a Nomad allocation in-place.
# $1	The allocation id to restart.
nomad_alloc_restart() {
	local alloc_id="${1}"
	
	curl -sS "${URL_NOMAD}/v1/client/allocation/${alloc_id}/restart";
	echo
}


###############################################################################


count_tasks=0;
count_restarted=0;
while read -r jobname taskname; do
	if [[ -z "${jobname}" ]]; then
		log_error "Error: No jobname specified." >&2;
		continue;
	fi
	if [[ -z "${taskname}" ]]; then
		log_error "Error: No taskname specified." >&2;
		continue;
	fi
	
	log_msg "Checking ${FYEL}${HC}${jobname}:${taskname}${RS}";
	
	while read -r alloc_id jobname taskname; do
		alloc_id_short="$(echo "${alloc_id}" | head -c7)";
		log_msg "Restarting allocation ${HC}${FBLE}${alloc_id_short}${RS} in-place";
		
		nomad_alloc_restart "${alloc_id}";
		count_restarted="$((count_restarted+1))";
	done < <(nomad_allocs_find "${jobname}" "${taskname}");
	count_tasks="$((count_tasks+1))";
done <"${loc_nomad_txt}"

log_msg "Restarted ${HC}${count_restarted}${RS} allocations for ${HC}${count_tasks}${RS} tasks in-place";