#!/usr/bin/bash

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

### CONFIGURATION {{
RC_FILE="${RC_FILE:-${XDG_CONFIG_HOME:-$HOME/.config}/wnl/wnlrc}"
EXIT_AFTER_NONZERO=0
SHELL_CMD="$SHELL -c"
HOOK_PRE=
HOOK_POST=
HOOK_STARTUP=
HOOK_EXIT=
BANNER_PRE_ENABLE=1
BANNER_POST_ENABLE=1
SHELL_INTEGRATION_ENABLE=1
AUTOSLOT_MAX=10
load_rc_file
### }} CONFIGURATION

pidfile() {
	local slot=$1
	echo "${XDG_RUNTIME_DIR:?XDG_RUNTIME_DIR not set}/wnl_slot_${slot}.pid"
}

### ARG PARSING {{
# check if $1 is a number
if [ -n "$1" ] && [ "$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
else
	# otherwise, find the first available slot
	AUTOSLOT=1
	for ((i=1; i <= "$AUTOSLOT_MAX"; i++)); do
		if [ ! -f "$(pidfile $i)" ]; then
			SLOT=$i
			break
		fi
	done
	if [ -z "$SLOT" ]; then
		echo Error: unable to find an available slot
		exit 1
	fi
fi
export SLOT

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

### CONSTANTS {{
PIDFILE=$(pidfile "$SLOT")
# 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'
### }} CONSTANTS

# define colors
# https://unix.stackexchange.com/a/10065
if test -t 1; then
    # see if it supports colors...
    ncolors=$(tput colors)
    if test -n "$ncolors" && test "$ncolors" -ge 8; then
        FMT_BOLD="$(tput bold)"; export FMT_BOLD
        FMT_UNDERLINE="$(tput smul)"; export FMT_UNDERLINE
        FMT_STANDOUT="$(tput smso)"; export FMT_STANDOUT
        FMT_NORMAL="$(tput sgr0)"; export FMT_NORMAL
        FMT_BLACK="$(tput setaf 0)"; export FMT_BLACK
        FMT_RED="$(tput setaf 1)"; export FMT_RED
        FMT_GREEN="$(tput setaf 2)"; export FMT_GREEN
        FMT_YELLOW="$(tput setaf 3)"; export FMT_YELLOW
        FMT_BLUE="$(tput setaf 4)"; export FMT_BLUE
        FMT_MAGENTA="$(tput setaf 5)"; export FMT_MAGENTA
        FMT_CYAN="$(tput setaf 6)"; export FMT_CYAN
        FMT_WHITE="$(tput setaf 7)"; export FMT_WHITE
        FMT_GREY="$(tput setaf 8)"; export FMT_GREY
    fi
fi

handle_usr1() {
	SIGNAL_START=1
}

handle_usr2(){
	SIGNAL_STOP=1
}

handle_exit() {
	# echo removing pidfile
	rm -f "$PIDFILE"
	kill "$SLP_PID"
	hook_exit
}

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

banner_post() {
	test "$BANNER_POST_ENABLE" -eq "0" && return
	if [ "$EXIT_CODE" -ne 0 ]; then
		bracket_color="${FMT_RED}"
		exit_code_color="${FMT_RED}${FMT_BOLD}"
	else
		bracket_color="${FMT_CYAN}"
		exit_code_color="${FMT_NORMAL}"
	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}${FMT_NORMAL} ${FMT_GREY}at${FMT_NORMAL} $(date +%X) ${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
}

hook_pre() {
	test -n "$HOOK_PRE" && $SHELL_CMD "$HOOK_PRE"
}

hook_post() {
	test -n "$HOOK_POST" && $SHELL_CMD "$HOOK_POST"
}

hook_startup() {
	test -n "$HOOK_STARTUP" && $SHELL_CMD "$HOOK_STARTUP"
}

hook_exit() {
	test -n "$HOOK_EXIT" && $SHELL_CMD "$HOOK_EXIT"
}

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

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

cmd_title() {
	# 56 is roughly the width of the banner_post (which in longer than banner_pre), minus CMD_TITLE
	CMD_TITLE_MAXLENGTH=$(("$(tput cols)" - 56))
	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 heuristically decides what signal should be sent to kill SHELL_CHILD
kill_signal() {
	cmd=$1
	if [[ "$cmd" =~ ^[a-zA-Z]*sh\ -c ]]; then
		echo TERM
	elif [[ "$cmd" =~ ^make\  ]]; then
		echo TERM
	else
		echo INT
	fi
}

## MAIN {{
# make pidfile with locking
exec 17<>"$PIDFILE"
if ! flock -n -x 17; then
    echo "Error: unable to acquire lock on $PIDFILE" >&2
    exit 1
fi
echo $$ > "$PIDFILE"
test -n "$AUTOSLOT" && echo wnl starting with slot "$SLOT"

trap handle_usr1 SIGUSR1
trap handle_usr2 SIGUSR2
trap handle_exit EXIT

sleep infinity &
SLP_PID=$!
SIGNAL_START=0
SIGNAL_STOP=0
hook_startup
while true; do 
	# reload rc file every run
	load_rc_file
	# 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
		# if RUNNING_CMD is defined and it's running, then kill it
		if [ -v RUNNING_CMD ] && kill -0 "$RUNNING_CMD" 2>/dev/null; then
			# Kill the actual process the user specified, rather than trying to kill the shell 
			# _running_ that command.
			# kill -INT "$RUNNING_CMD"
			SHELL_CHILD=$(pgrep -P "$RUNNING_CMD")
			if [ -n "$SHELL_CHILD" ]; then
				kill -"$(kill_signal "$CMD_RAW")" "$SHELL_CHILD"
			fi
			wait -n "$RUNNING_CMD"
			# export EXIT_CODE for consumption by hook
			export EXIT_CODE=$?
			if ! kill -0 "$RUNNING_CMD" 2>/dev/null; then
				banner_post
				hook_post
			fi
		fi
	elif [ "$SIGNAL_START" -eq 1 ]; then
		# 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
			$SHELL_CMD "${CMD}" >&1 2>&2 &
			RUNNING_CMD=$!
		fi
	fi
	SIGNAL_START=0
	SIGNAL_STOP=0
	# 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 ] && kill -0 "$RUNNING_CMD" 2>/dev/null; then
		wait -n "$RUNNING_CMD"
		# export EXIT_CODE for consumption by hook
		export EXIT_CODE=$?
		# only summarize 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
			banner_post
			hook_post
		fi
	else
		wait
	fi
	# restart infinite sleep, just in case it got killed somehow (e.g. by the user)
	if ! kill -0 "$SLP_PID" 2>/dev/null; then
		sleep infinity &
		SLP_PID=$!
	fi
done
## }} MAIN
