#!/bin/bash

set -u

function load_rc_file() {
	# shellcheck disable=SC1090
	test -f "$RC_FILE" && source "$RC_FILE"
}

SHOW_CONFIG=${SHOW_CONFIG:-false}
### CONFIGURATION {{
RC_FILE=${RC_FILE:-${XDG_CONFIG_HOME:-$HOME/.config}/wnl/wnlrc}
load_rc_file
EXIT_AFTER_NONZERO=${EXIT_AFTER_NONZERO:-0}
SHELL_CMD=${SHELL_CMD:-"$SHELL -c"}
HOOK_PRE=${HOOK_PRE:-}
HOOK_POST=${HOOK_POST:-}
HOOK_STARTUP=${HOOK_STARTUP:-}
HOOK_EXIT=${HOOK_EXIT:-}
BANNER_PRE_ENABLE=${BANNER_PRE_ENABLE:-1}
BANNER_POST_ENABLE=${BANNER_POST_ENABLE:-1}
SHELL_INTEGRATION_ENABLE=${SHELL_INTEGRATION_ENABLE:-1}
AUTOSLOT_MIN=${AUTOSLOT_MIN:-1}
AUTOSLOT_MAX=${AUTOSLOT_MAX:-10}
RUNTIME_DIR=${RUNTIME_DIR:-${XDG_RUNTIME_DIR:-/tmp}}
DOUBLE_TAP_REQUIRED=${DOUBLE_TAP_REQUIRED:-false}
DOUBLE_TAP_TIMEOUT_MS=${DOUBLE_TAP_TIMEOUT_MS:-200}
RESTART_MODE=${RESTART_MODE:-false}
STOPSTART_MODE=${STOPSTART_MODE:-false}
SIGNAL=${SIGNAL:-INT}
SSH_MODE=${SSH_MODE:-false}
SSH_EXECUTABLE="${SSH_EXECUTABLE:-ssh}"
SSH_ARGS="-o StreamLocalBindUnlink=yes ${SSH_ARGS:-}"
WNL_SSH="${WNL_SSH:-}"
### }} CONFIGURATION

function maybe_tput_colors() {
	if test -t "${stdout}"; then
		local ncolors
		ncolors=$(tput colors)
		if test -n "$ncolors" && test "$ncolors" -ge 8; then
			tput "$@"
		else
			:
		fi
	fi
}

### INTERNAL VALUES {{
exec {stdin}<&0
exec {stdout}>&1
exec {stderr}>&2
SIGNAL_START=0
SIGNAL_STOP=0
# https://sw.kovidgoyal.net/kitty/shell-integration/#notes-for-shell-developers
SHELL_INTEGRATION_PROMPT_START='\x1b\x5d133;A\x1b\x5c'
SHELL_INTEGRATION_CMD_START='\x1b\x5d133;C\x1b\x5c'
# define colors
# https://unix.stackexchange.com/a/10065
FMT_BOLD="$(maybe_tput_colors bold)"
export FMT_BOLD
FMT_UNDERLINE="$(maybe_tput_colors smul)"
export FMT_UNDERLINE
FMT_STANDOUT="$(maybe_tput_colors smso)"
export FMT_STANDOUT
FMT_NORMAL="$(maybe_tput_colors sgr0)"
export FMT_NORMAL
FMT_BLACK="$(maybe_tput_colors setaf 0)"
export FMT_BLACK
FMT_RED="$(maybe_tput_colors setaf 1)"
export FMT_RED
FMT_GREEN="$(maybe_tput_colors setaf 2)"
export FMT_GREEN
FMT_YELLOW="$(maybe_tput_colors setaf 3)"
export FMT_YELLOW
FMT_BLUE="$(maybe_tput_colors setaf 4)"
export FMT_BLUE
FMT_MAGENTA="$(maybe_tput_colors setaf 5)"
export FMT_MAGENTA
FMT_CYAN="$(maybe_tput_colors setaf 6)"
export FMT_CYAN
FMT_WHITE="$(maybe_tput_colors setaf 7)"
export FMT_WHITE
FMT_GREY="$(maybe_tput_colors setaf 8)"
export FMT_GREY
### }} INTERNAL VALUES

function diff_time() {
	t1=$(date -d "$1" +%s 2>/dev/null) || {
		echo "Bad date: $1" >&2
		return 1
	}
	t2=$(date -d "$2" +%s 2>/dev/null) || {
		echo "Bad date: $2" >&2
		return 1
	}

	# it turns out that arithmetic in bash can do cool things like assignment and ternaries
	# https://www.gnu.org/software/bash/manual/html_node/Shell-Arithmetic.html
	# "ARITHMETIC EVALUATION" section in bash manpage
	((delta = t1 > t2 ? t1 - t2 : t2 - t1))
	((days = delta / 86400, delta %= 86400))
	((hours = delta / 3600, delta %= 3600))
	((mins = delta / 60, secs = delta % 60))

	if [ "$days" -gt 0 ]; then
		printf "%d:%02d:%02d:%02d\n" "$days" "$hours" "$mins" "$secs"
	elif [ "$hours" -gt 0 ]; then
		printf "%d:%02d:%02d\n" "$hours" "$mins" "$secs"
	else
		printf "%d:%02d\n" "$mins" "$secs"
	fi
}

function banner_pre() {
	START_TIME=$(date +%X)
	test "$BANNER_PRE_ENABLE" -eq "0" && return
	bracket_color="${FMT_GREEN}"
	echo "${bracket_color}[[ ${FMT_NORMAL}${FMT_GREY}running${FMT_NORMAL} $(cmd_title) ${FMT_GREY}at${FMT_NORMAL} ${START_TIME} ${FMT_GREY}in slot${FMT_NORMAL} $SLOT ${bracket_color}]]${FMT_NORMAL}"
}

function banner_post() {
	local stop_time
	stop_time=$(date +%X)
	test "$BANNER_POST_ENABLE" -eq "0" && return
	local fmt_normal=${FMT_NORMAL}
	if [ "$EXIT_CODE" -ne 0 ]; then
		if ${restarting:-false}; then
			local bracket_color="${FMT_GREY}"
			local exit_code_color="${FMT_GREY}${FMT_BOLD}"
			local restarting_notice="(restarting)"
			fmt_normal=${FMT_NORMAL}${FMT_GREY}
		else
			local bracket_color="${FMT_RED}"
			local exit_code_color="${FMT_RED}${FMT_BOLD}"
			local restarting_notice=""
		fi
	else
		local bracket_color="${FMT_GREEN}"
		local exit_code_color="${FMT_NORMAL}"
		local restarting_notice=""
	fi
	shell_integrations_prompt_start
	echo "${bracket_color}[[ ${fmt_normal}${FMT_GREY}finished ${fmt_normal}$(cmd_title)${FMT_GREY} with exit code ${exit_code_color}${EXIT_CODE}${restarting_notice}${fmt_normal} ${FMT_GREY}at${fmt_normal} ${stop_time}(+$(diff_time "$START_TIME" "$stop_time")) ${FMT_GREY}in slot${fmt_normal} $SLOT ${bracket_color}]]${fmt_normal}"
	if [ "$EXIT_CODE" -ne "0" ] && [ "$EXIT_AFTER_NONZERO" -ne "0" ]; then
		exit "$EXIT_CODE"
	fi
}

function banner_startup() {
	local ssh_suffix
	ssh_suffix=$(if test -n "$WNL_SSH"; then echo " on $HOSTNAME"; fi)
	echo wnl starting with slot "$SLOT$ssh_suffix"
}

function hook_pre() {
	# shellcheck disable=SC2291
	test -n "$HOOK_PRE" && { $SHELL_CMD "$HOOK_PRE" || echo HOOK_PRE exited with code $?. continuing.; }
}

function hook_post() {
	# shellcheck disable=SC2291
	test -n "$HOOK_POST" && { $SHELL_CMD "$HOOK_POST" || echo HOOK_POST exited with code $?. continuing.; }
}

function hook_startup() {
	# shellcheck disable=SC2291
	test -n "$HOOK_STARTUP" && { $SHELL_CMD "$HOOK_STARTUP" || echo HOOK_STARTUP exited with code $?. continuing.; }
}

function hook_exit() {
	# shellcheck disable=SC2291
	test -n "$HOOK_EXIT" && { $SHELL_CMD "$HOOK_EXIT" || echo HOOK_EXIT exited with code $?. continuing.; }
}

function shell_integrations_cmd_start() {
	test "$SHELL_INTEGRATION_ENABLE" -ne "0" && echo -ne "$SHELL_INTEGRATION_CMD_START"
}

function shell_integrations_prompt_start() {
	test "$SHELL_INTEGRATION_ENABLE" -ne "0" && echo -ne "$SHELL_INTEGRATION_PROMPT_START"
}

function cmd_title() {
	# 63 is roughly the width of the banner_post (which in longer than banner_pre), minus CMD_TITLE
	CMD_TITLE_MAXLENGTH=$(("$(tput cols)" - 63))
	test $CMD_TITLE_MAXLENGTH -lt 10 && CMD_TITLE_MAXLENGTH=10
	excess_title_length="$(("${#CMD_RAW}" - "${CMD_TITLE_MAXLENGTH}"))"
	if [ "${excess_title_length}" -gt "1" ]; then
		trim_head="${CMD_RAW:0:$(("$CMD_TITLE_MAXLENGTH" / 2))}"
		trim_tail="${CMD_RAW:$(("${#CMD}" - ("$CMD_TITLE_MAXLENGTH" / 2))):"${#CMD}"}"
		CMD_TITLE="${trim_head}…${trim_tail}"
	else
		CMD_TITLE="$CMD_RAW"
	fi
	echo "$CMD_TITLE"
}

# kill_signal decides what signal should be sent to kill RUNNING_CMD
kill_signal() {
        echo $SIGNAL
}

function sleep_started() {
	if ! pgrep -f "$SLP_CMD" >/dev/null; then
		$SLP_CMD &
	fi
}

function sleep_killed() {
	test -v SLP_CMD && pkill -f "$SLP_CMD"
}

function received_command_start() {
	if $DOUBLE_TAP_REQUIRED; then
		now=$(date +%s)
		if [[ -v last_double_tap ]] && ((now <= (last_double_tap + DOUBLE_TAP_TIMEOUT_MS / 1000))); then
			SIGNAL_START=1
			last_double_tap=0
		else
			last_double_tap=$now
		fi
	else
		SIGNAL_START=1
	fi
}

function received_command_stop() {
	SIGNAL_STOP=1
}

function received_command_unknown() {
	echo received unknown command "'$1'"
}

function handle_exit() {
	test -v SOCKFILE && rm -f "$SOCKFILE"
	test -v LOCAL_SOCK && rm -f "$LOCAL_SOCK"
	sleep_killed
	hook_exit
}

trap handle_exit EXIT

trap 'exit 0' INT

function sockfile() {
	local slot=$1
	mkdir -p "${RUNTIME_DIR}"
	echo "${RUNTIME_DIR}/wnl_slot_${slot}.sock"
}

function lock_presockfile() {
	local slot=$1
	# shellcheck disable=SC2155
	local presockfile_to_lock="$(sockfile "$slot")"
	if [ ! -e "$presockfile_to_lock" ]; then
		exec 17<>"$presockfile_to_lock"
		if ! flock -n -x 17; then
			return 2
		fi
		echo $$ >"$presockfile_to_lock"
		SOCKFILE=$(sockfile "$SLOT")
	else
		return 1
	fi
}

function stop_command() {
	# if RUNNING_CMD is defined and it's running, then kill it
	if [ -v RUNNING_CMD ] && kill -0 "$RUNNING_CMD" 2>/dev/null; then
		local KILL_SIGNAL
		KILL_SIGNAL="$(kill_signal "$CMD_RAW")"
		kill -"$KILL_SIGNAL" "$RUNNING_CMD" $(pgrep -P $RUNNING_CMD)
		wait -n "$RUNNING_CMD"
		# export EXIT_CODE for consumption by hook
		export EXIT_CODE=$?
		if ! kill -0 "$RUNNING_CMD" 2>/dev/null; then
			unset RUNNING_CMD
			banner_post 1>&"$stdout" 2>&"$stderr" <&"$stdin"
			hook_post 1>&"$stdout" 2>&"$stderr" <&"$stdin"
		fi
	fi
}

function start_command() {
	# if RUNNING_CMD is not defined or it's not running, then start it
	if ! [ -v RUNNING_CMD ] || ! kill -0 "$RUNNING_CMD" 2>/dev/null; then
		banner_pre
		hook_pre
		shell_integrations_cmd_start
		(
			# this `set -m` crucially enables job control for this backgrounded subshell.
			# without it, SIGINT handling doesn't work.
			set -m
			exec $SHELL_CMD "$EXEC_IF_FISH${CMD}" 1>&"$stdout" 2>&"$stderr" <&"$stdin"
		) &
		RUNNING_CMD=$!
	fi
}

function start_or_stop_command() {
	if ! [ -v RUNNING_CMD ] || ! kill -0 "$RUNNING_CMD" 2>/dev/null; then
		start_command
	else
		stop_command
	fi
}

### ARG PARSING {{
# check if $1 is a number
if [ "${1:-}" -eq "${1:-}" ] 2>/dev/null; then
	# if it is, then use it as the slot
	SLOT=$1
	# and then shift (remove $1 from args) to have the rest turn into CMD
	shift
	lock_presockfile "$SLOT" || { echo "Error: unable to acquire lock on $(sockfile "$SLOT")" >&2 && exit 1; }
elif [[ -v SLOT ]] && [[ "$SLOT" -eq "$SLOT" ]]; then
	# no-op to honor SLOT if it is set
	:
else
	# otherwise, find the first available slot
	for ((i = AUTOSLOT_MIN; i <= AUTOSLOT_MAX; i++)); do
		SLOT=$i
		lock_presockfile "$SLOT" && break
	done
	if [ -z "$SLOT" ]; then
		echo Error: unable to find an available slot
		exit 1
	fi
fi
export SLOT

# possibly enter ssh mode
if [ "${1:-}" = "ssh" ]; then
	SSH_MODE=true
	SOCKFILE=$(sockfile "$SLOT")
	rm -f "$SOCKFILE"
	shift
	# and then possibly get the remote slot
	if [ "${1:-}" -eq "${1:-}" ] 2>/dev/null; then
		REMOTE_SLOT=$1
		shift
	else
		REMOTE_SLOT=$SLOT
	fi
	SSH_ARGS+=" ${1:?must provide arguments to ssh}"
	shift
elif [ "${1:-}" = "--" ]; then
	shift
fi

# escape command to preserve user-provided as faithfully as possible
CMD=$(printf "%q " "$@")
CMD_RAW="$*"
# remove trailing space
CMD="${CMD%" "}"
### }} ARG PARSING

# craft a predictable command line for sleeping, so that `wait` in one job can
# be interrupted by another
sleep_duration="123000"
SOCKFILE=$(sockfile "$SLOT")
# iterate over each byte of $str
for ((i = 0; i < ${#SOCKFILE}; i++)); do
	c=${SOCKFILE:i:1}
	sleep_duration+=$(printf "%d" "'$c")
done
SLP_CMD="sleep 1000000${sleep_duration}000${SLOT}"

# stop_command handling relies on identifying the process being run
# ($RUNNING_CMD). standard shells (e.g. bash) when called with `-c` do an
# immediate `exec`. The pid can then easily be found with `RUNNING_CMD=$!`, as
# shown in start_command
#
# However, fish does a `fork` then `exec`. So pgrepping for the child process
# would becomes necessary. The fish devs are aware of this difference in
# behavior, but will not fix:
# https://github.com/fish-shell/fish-shell/issues/6902
#
# So the workaround is to insert a little `exec` into the command whenever fish
# is the shell in use ¯\_(ツ)_/¯
if [[ "$SHELL_CMD" =~ .*fish\ -c ]]; then
	EXEC_IF_FISH="exec "
else
	EXEC_IF_FISH=""
fi

function get_config_item() {
	local config_key=${1:?must provide config key}
	local config
	config=$(printf "%s" "$(
		IFS="=" read -r _ value < <(cat | grep "${config_key}=")
		echo "$value"
	)")
	echo "$config"
}

function show_config() {
	for c in RC_FILE \
		EXIT_AFTER_NONZERO \
		SHELL_CMD \
		HOOK_PRE \
		HOOK_POST \
		HOOK_STARTUP \
		HOOK_EXIT \
		BANNER_PRE_ENABLE \
		BANNER_POST_ENABLE \
		SHELL_INTEGRATION_ENABLE \
		AUTOSLOT_MIN \
		AUTOSLOT_MAX \
		RUNTIME_DIR \
		DOUBLE_TAP_REQUIRED \
		DOUBLE_TAP_TIMEOUT_MS \
		RESTART_MODE \
		STOPSTART_MODE \
		SLOT \
		SOCKFILE; do
		# https://unix.stackexchange.com/a/397587
		echo $c="${!c}"
	done
	exit 0
}

function run_ssh() {
	local local_config remote_config remote_sock remote_slot
	local_config=$(SHOW_CONFIG=true wnl "$SLOT" --) || {
		echo error retrieving local wnl config
		exit 1
	}
	# shellcheck disable=SC2086
	remote_config=$($SSH_EXECUTABLE $SSH_ARGS -- SHOW_CONFIG=true wnl "$REMOTE_SLOT") || {
		echo error retrieving remote wnl config
		exit 1
	}
	LOCAL_SOCK="$(echo "$local_config" | get_config_item SOCKFILE)"
	# create a short-lived file in the place where SSH will create a socket.
	# this does two things:
	# 1. helps claim the socket so that other wnl instances won't steal it, and
	# 2. provides an indication that the slot will be filled, such that a
	# user's HOOK_STARTUP can take any desired actions
	# https://stackoverflow.com/a/13829090
	set -o noclobber
	{ : >"$LOCAL_SOCK"; } &>/dev/null
	remote_sock="$(echo "$remote_config" | get_config_item SOCKFILE)"
	remote_slot="$(echo "$remote_config" | get_config_item SLOT)"
	local ssh_command_line=("$SSH_EXECUTABLE $SSH_ARGS -t -L $LOCAL_SOCK:$remote_sock")
	banner_startup
	if [[ -n "$CMD_RAW" ]]; then
		ssh_command_line+=("WNL_SSH=true wnl $remote_slot $CMD")
	else
		echo -n "${FMT_BLUE}${FMT_BOLD}"
		$(if command -v cowsay >/dev/null; then echo cowsay; else echo echo; fi) "enter in 'wnl $REMOTE_SLOT <yourcommand>'"
		echo -n "${FMT_NORMAL}"
	fi
	# hook_startup added here to preserve symmetry with hook_exit, which is
	# called on EXIT
	hook_startup
	# shellcheck disable=SC2068
	${ssh_command_line[@]}
}

function main_loop() {
	while read -r operation _; do
		# reload rc file every run
		load_rc_file
		case $operation in
		START)
			received_command_start
			;;
		STOP)
			received_command_stop
			;;
		*)
			received_command_unknown "$operation"
			continue
			;;
		esac
		# echo eval state begin: SIGNAL_STOP: $SIGNAL_STOP, SIGNAL_START: $SIGNAL_START, RUNNING_CMD: $RUNNING_CMD, IS_RUNNING: "$([ -v RUNNING_CMD ] && kill -0 "$RUNNING_CMD" 2>/dev/null && echo yes || echo no)"
		if [ "$SIGNAL_STOP" -eq "1" ]; then
			stop_command
		elif [ "$SIGNAL_START" -eq 1 ]; then
			if $STOPSTART_MODE; then
				start_or_stop_command
			else
				if $RESTART_MODE; then
					restarting=true
					stop_command
					restarting=false
				fi
				start_command
			fi
		fi
		SIGNAL_START=0
		SIGNAL_STOP=0
		sleep_started
		# echo eval state end: SIGNAL_STOP: $SIGNAL_STOP, SIGNAL_START: $SIGNAL_START, RUNNING_CMD: $RUNNING_CMD, IS_RUNNING: "$([ -v RUNNING_CMD ] && kill -0 "$RUNNING_CMD" 2>/dev/null && echo yes || echo no)"
		if [ -v RUNNING_CMD ]; then
			# FIXME: mild race condition here--it can happen that:
			# 1. the above `test -v` succeeds
			# 2. stop_command() successfully `wait`s RUNNING_CMD
			# 3. then, because `wait`ing a PID isn't idempotent, this following
			# `wait` fails
			wait -n "$RUNNING_CMD" $(pgrep -f "$SLP_CMD")
			# export EXIT_CODE for consumption by hook
			export EXIT_CODE=$?
			# only run post-CMD actions if RUNNING_CMD is done, i.e. that the wait wasn't just interrupted to handle a signal
			if ! kill -0 "$RUNNING_CMD" 2>/dev/null; then
				unset RUNNING_CMD
				banner_post
				hook_post
			fi
		fi
	done < <(
		socat -u UNIX-LISTEN:"$SOCKFILE",fork,reuseaddr STDOUT |
			while read -r line; do
				sleep_killed
				printf '%s\n' "$line"
			done
	)
}

function main() {
	rm -f "$SOCKFILE"
	main_loop &
	banner_startup
	hook_startup
	while true; do
		sleep infinity
	done
}

if $SHOW_CONFIG; then
	show_config
elif $SSH_MODE; then
	run_ssh
	exit $?
else
	main
fi
