#!/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=1

## 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 disable the SGE queue of the current host (default)
or the one passed as first argument.

EXAMPLES :
    - Disable SGE's queue of the current host
        ${PROGNAME}

    - Disable SGE's queue of the "marvin.domain.tld" host
        ${PROGNAME} marvin.domain.tld
	EOF

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

	local_message="${1}"

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

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

	local_sge_hostname="${1}"
	local_sge_queue_name="${2}"
	## List all queues with 'disable' state and filter to the expected queue name
	### Then count the number of lines between "queuename" (from header) and the expected queue
	### To avoid to count pending jobs for this queue
	local_sge_queue_test=$(qstat -f -qs d -q "${local_sge_queue_name:=/dev/null}@${local_sge_hostname:=/dev/null}" \
		| awk 'BEGIN {count = 0} /queuename/,/main.q/ {count++} END { print count }')

	case "${local_sge_queue_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_sge_hostname:=/dev/null} queue (command return ${local_sge_queue_test} lines).${RESET}"
			exit 2
			;;
esac

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

	return "${return_queue_enable}"

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

	local_sge_hostname="${1}"
	local_sge_queue_name="${2}"
	## List all queues with 'disable' state and filter to the expected queue name
	### Then count the number of lines between "queuename" (from header) and the expected queue
	### To avoid to count pending jobs for this queue
	local_sge_queue_test=$(qstat -f -qs d -q "${local_sge_queue_name:=/dev/null}@${local_sge_hostname:=/dev/null}" \
		| awk 'BEGIN {count = 0} /queuename/,/main.q/ {count++} END { print count }')

	case "${local_sge_queue_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_sge_queue_name:=/dev/null}@${local_sge_hostname:=/dev/null} queue (command return ${local_sge_queue_test} lines).${RESET}"
			exit 3
			;;
esac

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

	return "${return_queue_disable}"

}
# }}}
disable_sge_queue() {                                           # {{{

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

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

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

	return "${return_disable_queue}"

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

	manage_args "${ARGS}"

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

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

	## If the queue is already disable
	### Exit
	is_queue_disable "${sge_hostname}" "${sge_queue_name}" \
		&& exit 0

	## If the queue is enable
	### Try to disable it
	is_queue_enable "${sge_hostname}" "${sge_queue_name}" \
		&& disable_sge_queue "${sge_hostname}" "${sge_queue_name}"

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

}
# }}}

main

exit 0