#!/bin/bash

if [ "${BASH_VERSION%%.*}" -lt 4 ]; then
	>&2 echo "$0: requires Bash 4 or newer, exiting..."
	exit 1
fi

script_name=$(basename "$0")
ssh="ssh -o ControlMaster=auto -o ControlPath=$HOME/.ssh/${RANDOM::4}_%C -o ControlPersist=yes"
force=0  # Changed using -f/--force

usage() {
	cat <<- EOF
	Usage: $script_name [ --force ] [ --uninstall ] [ HOST [ HOST ] ... ]

	Resets specified WEKA cluster members back to STEM mode to allow them to then
	be re-added to a cluster or added to a new cluster. Alternatively, it can
	uninstall WEKA from hosts completely.

	Prior to these operations, it will unmount all WekaFS file systems from the
	specified hosts.

	BOTH OPERATIONS ARE DESTRUCTIVE! This tool is only host-aware, not
	cluster-aware; if using this script for host removal take great care to
	ensure that this does not irreverisbly break any clusters.

	To reset hosts back to STEM mode:

	    $script_name weka1 weka2 weka3 weka4 weka5

	Brace expansion can alternatively be used. To specify a range:

	    $script_name weka{1..5}.example.com

	Or specific hosts with a common base hostname:

	    $script_name weka{1,3,5}.example.com

	To also uninstall WEKA on these hosts (also demonstrating that FQDNs are not
	required):

	    $script_name --uninstall weka{1..5}

	To suppress the initial prompt detailing what will be done:

	    $script_name --yes --uninstall weka{1..5}

	To also suppress prompts for confirmation before working on each host:

	    $script_name --yes --force --uninstall weka{1..5}

	Options:
	    -f, --force
	        Do not prompt for confirmation before working on each host.

	    -u, --uninstall
	        Uninstall WEKA from each host.

	    -y, --yes
	        Do not detail the operation that will be done and the target hosts,
	        do not prompt to continue.

	    -h, --help
	        Show this help/usage message.
	EOF
}

cleanup() {
	for host in $hosts; do
		$ssh -O exit "$host"
	done
}

weka_in_path() {
	$ssh "$1" 'which weka' 2>&1 > /dev/null
}

unmount_all_weka_fs() {
	$ssh "$1" 'sudo umount --types wekafs --all'
}

unmount_all_weka_fs_all_hosts() {
	local rc=0
	for host in $hosts; do
		echo "Unmounting WekaFS file systems on $host..."
		if ! (unmount_all_weka_fs "$host" &); then
			>&2 echo "Unable to unmount WEKA file systems on $1"
			rc=1
		fi
	done
	wait

	# Sleep required otherwise multiple SSH sessions try to become master if
	# there aren't any file sytsems to unmount
	sleep 0.25
	if [ "$rc" -eq 1 ]; then
		>&2 echo "Aborting $mode process for all hosts"
		exit 1
	fi
}

stop_weka() {
	$ssh "$1" 'sudo weka local stop --force'
}

remove_weka_containers() {
	$ssh "$1" 'sudo weka local rm --all --force'
}

create_stem_container() {
	$ssh "$1" 'sudo weka local setup container --name default'
}

_reset() {
	echo "Stopping WEKA on $1..."
	if ! stop_weka "$1"; then
		>&2 echo "Unable to stop WEKA on host $1"
		>&2 echo "Aborting reset process on $1"
		return 1
	fi

	echo "Removing WEKA containers on $1..."
	if ! remove_weka_containers "$1"; then
		>&2 echo "Unable to remove WEKA containers on $1"
		>&2 echo "Aborting reset process on $1"
		return 1
	fi

	echo "Recreating STEM mode container on $1"
	if ! create_stem_container "$1"; then
		>&2 echo "Unable to create STEM mode container on host $1"
		>&2 echo "Aborting reset process on $1"
		return 1
	fi
}

reset() {
	printf "\nResetting WEKA on $1\n"
	if [ "$force" -eq 0 ]; then
		printf 'Do you wish to continue? [y|N] '
		read -n 1 confirmation

		case $confirmation in
			[Yy]) echo ;;
			*) echo "Skipping host $1..."; return 0 ;;
		esac
	fi

	if ! weka_in_path "$1"; then
		>&2 echo "WEKA not in \$PATH on $1, not proceeding with reset"
		return 1
	fi

	# If -f/--force not set, don't parallelise/background
	if [ "$force" -eq 0 ]; then
		_reset "$1"
	else
		_reset "$1" &
	fi
}

_uninstall() {
	echo "Uninstalling Weka on $1..."
	if ! $ssh "$1" 'sudo weka agent uninstall --force'; then
		>&2 echo "Unable to uninstall Weka on $1"
		>&2 echo "Aborting uninstall process on $1"
		return 1
	fi
	echo "WEKA uninstalled from $1"
}

uninstall() {
	printf "\nUninstalling WEKA on $1\n"
	if [ "$force" -eq 0 ]; then
		printf 'Do you wish to continue? [y|N] '
		read -n 1 confirmation

		case $confirmation in
			[Yy]) echo ;;
			*) echo "Skipping host $1..."; return 0 ;;
		esac
	fi

	if ! weka_in_path "$1"; then
		>&2 echo "WEKA not in \$PATH on $1, not proceeding with uninstall"
		return 1
	fi

	# If -f/--force not set, don't parallelise/background
	if [ "$force" -eq 0 ]; then
		_uninstall "$1"
	else
		_uninstall "$1" &
	fi
}

main() {
	arg_error=0
	print_usage=0

	mode='reset'
	no_clarify_operation=0
	hosts=''

	raw_mode_support=0
	raw=0

	# Check for raw mode, remove it from $@
	# We require --raw as the first argument:
	# - To avoid future clashes with "<UNDERLYING_UTILITY>" options
	# - As we don't parse further input
	# - As parsing further input would require knowing if an opt takes args
	# $arg will get passed directly through to "<UNDERLYING_UTILITY>" later
	if [ "$1" = '--raw' ] && [ "$raw_mode_support" -ne 0 ]; then
		raw=1
		shift
	fi

	# Raw mode
	if [ "$raw" -eq 1 ]; then
		if [ -z "$@" ]; then
			echo 'Raw mode. No further opts/args provided.'
		else
			echo "Raw mode. Input provided: $@"
		fi

	# Normal mode
	else
		# Convert long options to short options, stick them back on the end of $@
		for arg in $@; do
			case $arg in
				--force) set -- "$@" '-f' ;;
				--help) set -- "$@" '-h' ;;
				--raw)
					if [ "$raw_mode_support" -eq 0 ]; then
						>&2 echo "$0: --raw must be first argument"
					else
						>&2 echo "$0: illegal long option -- ${arg:2}"
					fi
					arg_error=1
					;;
				--uninstall) set -- "$@" '-u' ;;
				--yes) set -- "$@" '-y' ;;
				--*)
					>&2 echo "$0: illegal long option -- ${arg:2}"
					arg_error=1
					;;
				*) set -- "$@" "$arg" ;;
			esac
			shift
		done

		# Cycle through options, setting corresponding variables as options are
		# found
		while getopts 'fhpuy' opts; do
			case $opts in
				f) force=1 ;;
				h) print_usage=1 ;;
				u) mode='uninstall' ;;
				y) no_clarify_operation=1 ;;
				?) arg_error=1 ;;
			esac
		done
		shift $((OPTIND - 1))
		[ "$arg_error" -ne 0 ] && >&2 usage && exit 1
		[ "$print_usage" -eq 1 ] && usage && exit 0

		# If not at least one argument; using $@ instead results in `binary
		# operator expected` errors if an option is used (e.g. `-f`)
		if [ -z "$1" ]; then
			>&2 echo "$0: no hosts specified"
			>&2 usage
			exit 1
		fi

		hosts=$@
		trap cleanup SIGINT SIGTERM EXIT

		if [ "$no_clarify_operation" -eq 0 ]; then
			echo "Will $mode WEKA on following hosts:"
			echo "$hosts" | tr ' ' '\n'
			printf '\nDo you wish to continue? [y|N] '
			read -n 1 confirmation
			echo
			case $confirmation in
				[Yy]) ;;
				*) echo 'Exiting...'; exit 0 ;;
			esac
		else
			echo "${mode^}ing WEKA on following hosts:"
			echo "$hosts" | tr ' ' '\n'
		fi

		unmount_all_weka_fs_all_hosts

		case $mode in
			reset) for host in $hosts; do reset "$host"; done; wait ;;
			uninstall) for host in $hosts; do uninstall "$host"; done; wait ;;
		esac
	fi
}

main "$@"

# vim: set filetype=sh:
