#!/bin/sh

# Vars {{{
readonly PROGNAME=$(basename "${0}")
readonly PROGDIR=$(readlink -m $(dirname "${0}"))
readonly ARGS="${*}"
readonly NBARGS="${#}"
## Test if DEBUG is already defined (by parent script,…)
[ -z "${DEBUG}" ] && readonly DEBUG=0

## Colors
readonly PURPLE='\033[1;35m'
readonly RED='\033[0;31m'
readonly RESET='\033[0m'
readonly COLOR_DEBUG="${PURPLE}"
# }}}

manage_args() {                                                 # {{{

	case "${NBARGS}" in
		0 )
			sge_hostname="$(hostname -f)"
			;;
		1 )
			sge_hostname="${1}"
			;;
		* )
			printf '%b\n' "${RED}Expect 1 or 0 argument.${RESET}"
			printf '%b\n' "---"
			usage
			exit 1
			;;
	esac

}
# }}}
usage() {                                                       # {{{

	cat <<- EOF
usage: $PROGNAME [hostname]

Try to enable all SGE queues of the current host (default)
or the one passed as first argument.

EXAMPLES :
    - Enable SGE's queue(s) of the current host
        ${PROGNAME}

    - Enable SGE's queue(s) of "marvin.domain.tld" host
        ${PROGNAME} marvin.domain.tld
	EOF

}
# }}}
debug_message() {                                               # {{{

	local_debug_message="${1}"

	## Print message if DEBUG is enable (=0)
	[ "${DEBUG}" -eq "0" ] && printf '\e[1;35m%-6b\e[m\n' "DEBUG − ${PROGNAME} : ${local_debug_message}"

}
# }}}
is_queue_enable() {                                            # {{{

	local_queue_enable_hostname="${1}"
	local_queue_enable_name="${2}"
	## List all queues with 'disable' state and filter to the expected queue name
	## with a fake_user to avoid pending jobs for this queue
	### And count returned lines
	local_queue_enable_test=$(qstat -f -qs d -q "${local_queue_enable_name:=/dev/null}@${local_queue_enable_hostname:=/dev/null}" -u fake_user \
		| wc -l)

	case "${local_queue_enable_test}" in
		0 ) ## No result so the queue is enable
			local_sge_queue_state="enable"
			return_queue_enable="0"
			;;
		3 ) ## Results (header + queue name) so the queue is disable
			local_sge_queue_state="disable"
			return_queue_enable="1"
			;;
		* ) ## Unexpected result
			printf '%b\n' "${RED}Not able to determine the state of ${local_sge_queue_name:=/dev/null}@${local_queue_enable_hostname:=/dev/null} queue (command return ${local_queue_enable_test} lines).${RESET}"
			exit 2
			;;
esac

	## Simple debug message to valid current variables
	debug_message "is_queue_enable − \
SGE queue: ${RED}${local_queue_enable_name:=/dev/null}${COLOR_DEBUG} \
state is: ${RED}${local_sge_queue_state:=/dev/null}${COLOR_DEBUG}."

	return "${return_queue_enable}"

}
# }}}
is_queue_disable() {                                            # {{{

	local_queue_disable_hostname="${1}"
	local_queue_disable_name="${2}"
	## List all queues with 'disable' state and filter to the expected queue name
	## add a fake_user to avoid pending jobs for this queue
	### And count returned lines
	local_queue_disable_test=$(qstat -f -qs d -q "${local_queue_disable_name:=/dev/null}@${local_queue_disable_hostname:=/dev/null}" -u fake_user \
		| wc -l)

	case "${local_queue_disable_test}" in
		0 ) ## No result so the queue is enable
			local_sge_queue_state="enable"
			return_queue_disable="1"
			;;
		3 ) ## Results (header + queue name) so the queue is disable
			local_sge_queue_state="disable"
			return_queue_disable="0"
			;;
		* ) ## Unexpected result
			printf '%b\n' "${RED}Not able to determine the state of ${local_queue_disable_name:=/dev/null}@${local_queue_disable_hostname:=/dev/null} queue (command return ${local_queue_disable_test} lines).${RESET}"
			exit 3
			;;
esac

	## Simple debug message to valid current variables
	debug_message "is_queue_disable − \
SGE queue: ${RED}${local_queue_disable_name:=/dev/null}${COLOR_DEBUG} \
state is: ${RED}${local_sge_queue_state:=/dev/null}${COLOR_DEBUG}."

	return "${return_queue_disable}"

}
# }}}
is_all_queue_enable() {                                            # {{{

	local_all_queue_enable_hostname="${1}"
	local_all_queue_enable_name="${2}"

	## By default, all queues are enable
	return_all_queue_enable="0"

	## Test all queues one by one
	for loop_enable_queue in ${local_all_queue_enable_name}; do

		### If a queue is not enable
		#### Change the return value
		is_queue_enable "${local_all_queue_enable_hostname}" "${loop_enable_queue}" \
			|| return_all_queue_enable="1"

	done

	return "${return_all_queue_enable}"

}
# }}}
enable_sge_queue() {                                           # {{{

	local_sge_hostname="${1}"
	local_sge_queue_name="${2}"

	## Simple debug message to valid current variables
	debug_message "enable_sge_queue − \
Try to enable SGE queue: ${RED}${local_sge_queue_name:=/dev/null}@${local_sge_hostname:=/dev/null}${COLOR_DEBUG}."

	## SGE command to enable the queue
	qmod --enable "${local_sge_queue_name}@${local_sge_hostname}" \
		&& return_enable_queue="${?}"

	return "${return_enable_queue}"

}
# }}}
main() {                                                        # {{{

	manage_args "${ARGS}"

	sge_queues_name="$(qhost -h "${sge_hostname:=/dev/null}" -q -xml \
		| grep "queue name" \
		| cut -d"'" -f2 )"
	sge_queues_name_print="$(qhost -h "${sge_hostname:=/dev/null}" -q -xml \
		| grep "queue name" \
		| cut -d"'" -f2 \
		| tr -s '\n' ' ' )"

	## Simple debug message with color to valid current variables
	debug_message "main − Try to manage \
SGE queue(s): ${RED}${sge_queues_name_print:=/dev/null}${COLOR_DEBUG}\
for host: ${RED}${sge_hostname:=/dev/null}${COLOR_DEBUG}."

	## If the queue(s) are already enable
	### Exit
	is_all_queue_enable "${sge_hostname}" "${sge_queues_name}" \
		&& exit 0

	## Test all queues one by one
	for loop_queue in ${sge_queues_name}; do

		## If the queue is disable
		### Try to enable it
		is_queue_disable "${sge_hostname}" "${loop_queue}" \
			&& enable_sge_queue "${sge_hostname}" "${loop_queue}"

		## If the queue is still disable
		### Exit with error
		is_queue_disable "${sge_hostname}" "${loop_queue}" \
			&& printf '%b\n' "${RED}ERROR ${loop_queue}@${sge_hostname} is still disable.${RESET}" \
			&& exit 4

	done

}
# }}}

main

exit 0