#!/usr/bin/env bash

set -euo pipefail

print_help() {
	cat <<EOF
Usage: smagic <command> [options]

Commands:
  ps [-a]           List Slurm jobs (default: current user)
  start [options]   Submit a long-running job (sbatch ... --wrap "hostname && sleep infinity")
  run [options]     Run an interactive Bash using srun
  bash <job>        Attach an interactive shell to a job (by ID or name)
  kill <job>...     Cancel one or more jobs (by ID or name)
  rkill <begin> <last>
                    Cancel your jobs whose ID is in [begin, last] (inclusive)

Global options:
  -h, --help        Show this help message and exit
EOF
}

get_jobid_for_target() {
	local target="$1"
	local jobid=""

	while read -r jid jname; do
		if [[ "${jid}" == "${target}" || "${jname}" == "${target}" ]]; then
			jobid="${jid}"
			break
		fi
	done < <(squeue -h -u "${USER}" -o "%i %j" 2>/dev/null || true)

	# If nothing matched but the target is numeric, treat it as a job ID directly.
	if [[ -z "${jobid}" && "${target}" =~ ^[0-9]+$ ]]; then
		jobid="${target}"
	fi

	if [[ -n "${jobid}" ]]; then
		printf '%s\n' "${jobid}"
		return 0
	fi

	echo "smagic: no job found for target '${target}'" >&2
	return 1
}

run_and_echo() {
	# Echo the singularity-related Slurm command before executing it.
	echo "+ $*"
	exec "$@"
}

has_wrap_arg() {
	for arg in "$@"; do
		if [[ "${arg}" == "--wrap" || "${arg}" == --wrap=* ]]; then
			return 0
		fi
	done
	return 1
}

subcommand="${1-}"

case "${subcommand}" in
	""|-h|--help)
		print_help
		exit 0
		;;
esac

shift || true

case "${subcommand}" in
	ps)
		ps_format='%.9i %.9P %.32j %.8u %.2t %.11M %.5D %R'
		case "${1-}" in
			-a|--all)
				shift || true
				run_and_echo squeue --format="${ps_format}" "$@"
				;;
			*)
				run_and_echo squeue --format="${ps_format}" -u "${USER}" "$@"
				;;
		esac
		;;

	start)
		case "${1-}" in
			-h|--help)
				# Delegate help to sbatch but rewrite the Usage header to mention smagic.
				sbatch --help | sed 's/^Usage: sbatch /Usage: smagic start /'
				exit 0
				;;
			*)
				if has_wrap_arg "$@"; then
					run_and_echo sbatch "$@"
				else
					run_and_echo sbatch "$@" --wrap "hostname && sleep infinity"
				fi
				;;
		esac
		;;

	run)
		case "${1-}" in
			-h|--help)
				srun --help | sed 's/^Usage: srun /Usage: smagic run /'
				exit 0
				;;
			*)
				run_and_echo srun --pty "$@" /bin/bash
				;;
		esac
		;;

	bash)
		if [[ $# -lt 1 ]]; then
			echo "Usage: smagic bash <jobid-or-name>" >&2
			exit 1
		fi

		target="$1"

		jobid="$(get_jobid_for_target "${target}")" || exit 1
		run_and_echo srun --overlap --jobid "${jobid}" --pty /bin/bash
		;;

	kill)
		case "${1-}" in
			""|-h|--help)
				echo "Usage: smagic kill <jobid-or-name> [<jobid-or-name> ...]" >&2
				exit 1
				;;
			*)
				# Snapshot the user's jobs once instead of calling squeue per target.
				jobs_snapshot="$(squeue -h -u "${USER}" -o "%i %j" 2>/dev/null || true)"
				jobids=()
				for target in "$@"; do
					jobid=""
					while read -r jid jname; do
						if [[ "${jid}" == "${target}" || "${jname}" == "${target}" ]]; then
							jobid="${jid}"
							break
						fi
					done <<< "${jobs_snapshot}"

					# If nothing matched but the target is numeric, treat it as a job ID directly.
					if [[ -z "${jobid}" && "${target}" =~ ^[0-9]+$ ]]; then
						jobid="${target}"
					fi

					if [[ -z "${jobid}" ]]; then
						echo "smagic: no job found for target '${target}'" >&2
						exit 1
					fi
					jobids+=("${jobid}")
				done
				run_and_echo scancel "${jobids[@]}"
				;;
		esac
		;;

	rkill)
		begin="${1-}"
		last="${2-}"
		if [[ -z "${begin}" || -z "${last}" ]]; then
			echo "Usage: smagic rkill <begin> <last>" >&2
			exit 1
		fi
		if ! [[ "${begin}" =~ ^[0-9]+$ && "${last}" =~ ^[0-9]+$ ]]; then
			echo "smagic: rkill requires numeric job IDs" >&2
			exit 1
		fi
		if (( begin > last )); then
			echo "smagic: rkill begin (${begin}) must be <= last (${last})" >&2
			exit 1
		fi

		jobids=()
		while read -r jid; do
			if (( jid >= begin && jid <= last )); then
				jobids+=("${jid}")
			fi
		done < <(squeue -h -u "${USER}" -o "%i" 2>/dev/null || true)

		if [[ ${#jobids[@]} -eq 0 ]]; then
			echo "smagic: no jobs found in range [${begin}, ${last}]" >&2
			exit 1
		fi
		run_and_echo scancel "${jobids[@]}"
		;;

	*)
		echo "Unknown subcommand: ${subcommand}" >&2
		print_help >&2
		exit 1
		;;
esac
