#!/bin/bash
# -*- mode: shell-script; indent-tabs-mode: t; sh-basic-offset: 8; sh-indentation: 8; sh-indent-for-case-alt: + -*-

script_path="$(dirname ${0})"
script_name="$(basename ${0})"
pbench_bin="$(realpath -e ${script_path}/..)"

# source the base script
. "${pbench_bin}"/base

# The sole purpose of this script is to ensure a tool of your choosing is used
# during the execution of a benchmark on the optional list of remote hostnames
# specified. For a list of performance tools, use the --help option or look at
# the ${pbench_bin}/tool-scripts directory.

# Defaults
remotes_arg="${full_hostname}"
group="default"
options=""
do_install=1
_do_test_labels=0

function usage() {
	# Usage is formatted for a 80 column width screen.
	printf -- "usage:\n"
	printf -- "${script_name} --name=<tool-name> [--group=<group-name>]"
	printf -- " [--no-install] [--remotes=<remote-host>[,<remote-host>]]"
	printf -- " [--labels=<label>[,<label>]] -- [all tool specific"
	printf -- " options here]\n"
	printf -- "${script_name} --name=<tool-name> [--group=<group-name>]"
	printf -- " [--no-install] [--remotes=@<remotes-file>] -- [all tool"
	printf -- " specific options here]\n\n"
	#             1         2         3         4         5         6         7         8
	# (tab)    1 901234567890123456789012345678901234567890123456789012345678901234567890
	printf -- "\tWhere the list of labels must match the list of remotes, one label per\n"
	printf -- "\thost.\n\n"
	printf -- "\tOne can specify as the argument to the \"--remotes\" option a single\n"
	printf -- "\tremote host, a list of remote hosts (comma-separated, no spaces) or an\n"
	printf -- "\t\"at\" sign (\"@\") followed by a filename.  In this last case, the file\n"
	printf -- "\tshould contain a list of hosts and their (optional) labels.  Each line\n"
	printf -- "\tof the file should contain a host name, optionally followed by a label\n"
	printf -- "\tseparated by a comma (\",\"); empty lines are ignored, and comments are\n"
	printf -- "\tdenoted by a leading hash, or pound (\"#\"), character.\n"
	printf -- "\nAvailable tools:\n"
	for tool in $(find ${pbench_bin}/tool-scripts -maxdepth 1 ! -type d ! -name '*README*' ! -name base-tool ! -name unittests -printf "%P\n" 2> /dev/null | sort); do
		printf -- "\t${tool}\n"
	done
	#                     1         2         3         4         5         6         7         8
	# (no tab)   12345678901234567890123456789012345678901234567890123456789012345678901234567890
	printf -- "\nFor a list of tool specific options, run:\n"
	printf -- "\t${pbench_bin}/tool-scripts/<tool-name> --help\n"
}

# Process options and arguments
opts=$(getopt -q -o hn:g:r:l: --longoptions "help,name:,group:,no-install,remote:,remotes:,label:,labels:,test-labels" -n "getopt.sh" -- "${@}")
if [[ ${?} -ne 0 ]]; then
	printf -- "\n${*}\n" >&2
	printf -- "${script_name}: you specified an invalid option\n\n" >&2
	usage >&2
	exit 1
fi

# Process only the options and arguments before the "--", which are intended
# for pbench-register-tool itself; all agruments after the "--" are for the
# tool specified by the "--name" option.
eval set -- "${opts}"
while true; do
	case "${1}" in
	-n|--name)
		shift
		if [[ -n "${1}" ]]; then
			name="${1}"
			shift
		fi
		;;
	-r|--remote|--remotes)
		shift
		if [[ -n "${1}" ]]; then
			remotes_arg="${1}"
			shift
		fi
		;;
	-l|--label|--labels)
		shift
		if [[ -n "${1}" ]]; then
			labels_arg="${1}"
			shift
		fi
		;;
	-g|--group)
		shift
		if [[ -n "${1}" ]]; then
			group="${1}"
			shift
		fi
		;;
	--no-install)
		shift
		do_install=0
		;;
	--test-labels)
		shift
		_do_test_labels=1
		;;
	-h|--help)
		usage
		exit 0
		;;
	--)
		shift
		break
		;;
	esac
done


if [[ -z ${name} ]]; then
	error_log "Missing required parameter --name=<tool-name>"
	usage >&2
	exit 1
fi
if [[ ! -x "${pbench_bin}/tool-scripts/${name}" ]]; then
	error_log "Could not find ${name} in ${pbench_bin}/tool-scripts/: has this tool been integrated into the pbench-agent?"
	usage >&2
	exit 1
fi

if [[ ${remotes_arg::1} == "@" ]]; then
	# We are dealing with a file containing the list of remotes.

	if [[ -n ${labels_arg} ]]; then
		error_log "--labels=${labels_arg} not allowed with remotes file (${remotes_arg})"
		usage >&2
		exit 1
	fi
	remotes_file=${remotes_arg#*@}
	if [[ ! -e ${remotes_file} ]]; then
		error_log "--remotes=@${remotes_file} specifies a file that does not exist"
		usage >&2
		exit 1
	fi
	declare -a remotes_A
	declare -A labels_A
	idx=0
	lineno=0
	while read line; do
		if [[ ${line::1} == "#" ]]; then
			# Ignore comments
			(( lineno++ ))
			continue
		fi
		IFS=',' read -ra parts <<< "${line}"
		if [[ ${#parts[@]} == 0 ]]; then
			# Ignore empty lines
			(( lineno++ ))
			continue
		elif [[ ${#parts[@]} == 1 ]]; then
			# Assume the only part is a remote host name.
			remotes_A[${idx}]=${parts[0]}
			labels_A[${parts[0]}]=""
		elif [[ ${#parts[@]} == 2 ]]; then
			# Assume the first part is the remote host name, and
			# the second part is the label for that host name.
			remotes_A[${idx}]=${parts[0]}
			labels_A[${parts[0]}]=${parts[1]}
		else
			error_log "--remotes=@${remotes_file} contains an invalid file format, expected lines with \"<hostname>[,<label>]\" at line #${lineno}"
			usage >&2
			exit 1
		fi
		(( lineno++ ))
		(( idx++ ))
	done < ${remotes_file}
else
	# First create the array of remotes and their labels, "remotes_A"
	# "labels_A", from the "remotes_arg" and "labels_arg".
	IFS=',' read -ra remotes_A <<< "${remotes_arg}"
	if [[ -z "${remotes_A}" ]]; then
		error_log "INTERNAL: missing -r|--remote|--remotes=<remote-host>[,<remote-host>] argument for some unknown reason (should not happen)"
		usage >&2
		exit 1
	fi

	declare -A labels_A
	if [[ -n ${labels_arg} ]]; then
		IFS=',' read -ra labels_A0 <<< "${labels_arg}"
		if [[ -n "${labels_A0}" ]]; then
			if [[ ${#remotes_A[@]} -ne ${#labels_A0[@]} ]]; then
				if [[ ${_do_test_labels} -eq 0 ]]; then
					# We emit an error message now if we
					# are not working on behalf of
					# pbench-register-tool-set, since it
					# will handle its own error message on
					# failure.
					error_log "The number of labels given, \"${labels_arg}\", does not match the number of remotes given, \"${remotes_arg}\""
					usage >&2
				fi
				exit 1
			fi
			# Now create an array for labels, where the index is
			# by given hostname from the remotes.
			for (( idx=0; idx<${#labels_A0[@]}; idx++ )); do
				labels_A[${remotes_A[${idx}]}]=${labels_A0[${idx}]}
			done
		fi
	fi
fi

if [[ ${_do_test_labels} -ne 0 ]]; then
	# Used by pbench-register-tool-set to avoid duplicating logic, we
	# have been asked to exit early successfully if all argument
	# processing has been successful so far.
	exit 0
fi

# The remainder of ${@} is now tool-specific options. The best way to not
# mess up arguments seems to be using an array.
idx=0
declare -a tool_opts
while [[ -n "${1}" ]]; do
	tool_opts[${idx}]="${1}"
	let idx=${idx}+1
	shift
done
tool_opt_array_cmd=""
for (( i=0; ${i} < ${#tool_opts[@]}; i++ )); do
	tool_opt_array_cmd="${tool_opt_array_cmd} tool_opts[${i}]=\"${tool_opts[${i}]}\";"
done
debug_log "tool_opts: \"${tool_opts[@]}\""

# At this point all argument processing is complete.

# Determine what our local hostname(s) and IP(s) are.
# FIXME: this does not check for ipv6
local_ips=$(ip a | grep "inet " | awk '{print $2}' | awk -F/ '{print $1}')
local_interfaces="${local_ips} ${hostname} ${full_hostname} localhost"

# If for some reason the user specified a remote host, and that host is a
# local interface, we must remove it to avoid ssh'ing to ourselves.
idx=0
local_hostname=""
local_label=""
declare -a remotes
for (( i=0; ${i} < ${#remotes_A[@]}; i++ )); do
	remote="${remotes_A[${i}]}"
	for local_interface in ${local_interfaces}; do
		if [[ "${remote}" == "${local_interface}" ]]; then
			if [[ -z "${local_hostname}" ]]; then
				if [[ "${remotes_arg}" != "${full_hostname}" ]]; then
					debug_log "The remote host you have provided, ${remote}, matches a local interface, so we will register this tool locally only"
				fi
				local_hostname="${full_hostname}"
				local_label="${labels_A[${remote}]}"
			else
				if [[ "${remotes_arg}" != "${full_hostname}" ]]; then
					debug_log "The remote host you have provided, ${remote}, matches a local interface, and has already been registered locally"
				fi
			fi
			remote=""
			break
		fi
	done
	if [[ -n "${remote}" ]]; then
		remotes[${idx}]=${remote}
		let idx=${idx}+1
	fi
done

rc=0
if [[ -n "${local_hostname}" ]]; then
	if [[ ${do_install} == 1 ]]; then
		# A local tool registration
		debug_log "checking to see if tool is installed..."
		${pbench_bin}/tool-scripts/${name} --install "${tool_opts[@]}"
		rc=${?}
	else
		rc=0
	fi

	if [[ ${rc} == 0 ]]; then
		# Remember the tool and its options since its installation was
		# successful.
		if [[ ! -d "${pbench_run}/tools-$group" ]]; then
			mkdir "${pbench_run}/tools-${group}"
			if [[ ${?} -ne 0 ]]; then
				error_log "Unable to create the necessary tools group directory, ${pbench_run}/tools-${group}"
				exit 1
			fi
		fi
		this_tool_file="${pbench_run}/tools-${group}/${name}"
		/bin/rm -f "${this_tool_file}"
		touch "${this_tool_file}"
		i=0
		# tool options are 1 per line
		for (( i=0; ${i} < ${#tool_opts[@]}; i++ )); do
			echo "${tool_opts[${i}]}" >> "${this_tool_file}"
		done
		_msg="${name} tool is now registered in group ${group}"
		log "${_msg}"
		printf -- "%s\n" "${_msg}"
		# If we have a label, update the label in the tool directory
		if [[ "${local_label}" != "" ]]; then
			this_label_file="${pbench_run}/tools-${group}/label"
			/bin/rm -f "${this_label_file}"
			echo "${local_label}" > "${this_label_file}"
		fi
	else
		error_log "Failed to install tool ${name} (${pbench_bin}/tool-scripts/${name} --install \"${tool_opts[@]}\")"
	fi
fi

# Make sure we capture any local error codes to report failure later of the
# overall script.
exit_rc=${rc}

if [[ ${do_install} == 1 ]]; then
	install_opt=""
else
	install_opt=" --no-install"
fi
for (( i=0; ${i} < ${#remotes[@]}; i++ )); do
	remote="${remotes[${i}]}"
	# Register this tool on a remote host
	label=${labels_A[${remote}]}
	if [[ -z "${label}" ]]; then
		label_opt=""
	else
		label_opt=" --label=${label}"
	fi
	# Can't really pass an array of tool options over ssh, so here we will
	# actually construct the array within the ssh call, then reference that
	# array.
	debug_log "${tool_opt_array_cmd}"' pbench-register-tool --name='"${name}"' --group='"${group}"${install_opt}${label_opt}' -- "${tool_opts[@]}" 2>&1'
	ssh ${ssh_opts} ${remote} "${tool_opt_array_cmd}"' pbench-register-tool --name='"${name}"' --group='"${group}"${install_opt}${label_opt}' -- "${tool_opts[@]}" 2>&1' | sed -e 's/\(.*\)/'[${remote}]'\1/g'
	rc=${?}
	let exit_rc=${exit_rc}+${rc}
	if [[ ${rc} -ne 0 ]]; then
		error_log "Failed to register tool ${name} on remote host ${remote}"
		continue
	fi
	if [[ ! -d "${pbench_run}/tools-${group}" ]]; then
		mkdir "${pbench_run}/tools-${group}"
		if [[ ${?} -ne 0 ]]; then
			error_log "Unable to create the necessary tools group directory, ${pbench_run}/tools-${group}"
			exit 1
		fi
	fi
	this_tool_file="${pbench_run}/tools-${group}/remote@${remote}"
	/bin/rm -f "${this_tool_file}"
	# The specific tool options are stored on the remote host, so there is
	# no need to keep them here, however we do use this file to store the
	# label for this remote host [if any].
	if [[ -z "${label}" ]]; then
		touch "${this_tool_file}"
	else
		echo "${label}" > "${this_tool_file}"
	fi
done

exit ${exit_rc}
