#!/usr/bin/env bash

readonly SCRIPT_NAME="limine-snapper-sync"
readonly LIMINE_LOCK_FILE="/tmp/limine-global.lock"
readonly RESTORE_LOCKFILE="/tmp/limine-snapper-restore.lock"
export HOOK_CALLER="${HOOK_CALLER:-$SCRIPT_NAME}"
export HOOK_CMDLINE="$*"

ENABLE_MUTEX=true
ENABLE_HOOKS=true
ENABLE_DEBOUNCE=false
RESTORE_MODE=false
ARGS=()

for arg in "$@"; do
	case "$arg" in
	--no-mutex)
		ENABLE_MUTEX=false
		;;
	--mutex)
		ENABLE_MUTEX=true
		;;
	--no-hooks)
		ENABLE_HOOKS=false
		;;
	--hooks)
		ENABLE_HOOKS=true
		;;
	--debounce)
		ENABLE_DEBOUNCE=true
		;;
	--restore | --restore-kernels)
		RESTORE_MODE=true
		ARGS+=("$arg")
		;;
	*)
		ARGS+=("$arg")
		;;
	esac
done

_color_reset=""
_color_yellow=""
_color_red=""
_color_prompt=""
colors="$(tput colors 2>/dev/null || echo 0)"
if ((colors >= 8)); then
	_color_reset="\033[0m"
	_color_yellow="\033[1;33m"
	_color_red="\033[1;31m"
	_color_prompt="\033[0;33m"
fi

warning_msg() {
	echo -e "${_color_yellow}WARNING: $1${_color_reset} ${2:-}" >&2
}

error_msg() {
	echo -e "${_color_red}ERROR: $1${_color_reset} ${2:-}" >&2
}

if ((EUID != 0)); then
	error_msg "${SCRIPT_NAME} must be run with root privileges."
	exit 1
fi

mutex_lock() {
	local name=$1
	exec 200>"${LIMINE_LOCK_FILE}" || {
		rm -f "${LIMINE_LOCK_FILE}"
		exec 200>"${LIMINE_LOCK_FILE}"
	}
	flock --timeout=30 200 || {
		warning_msg "Mutex lock timeout on ${name}."
		return 1
	}
}

mutex_unlock() {
	flock --unlock 200
}

cleanup_all() {
	$RESTORE_MODE && rm -f "${RESTORE_LOCKFILE}"
	$ENABLE_MUTEX && mutex_unlock
}

run_boot_hooks() {
	local dir="$1"
	local hook rc
	for hook in /etc/boot/hooks/"$dir".d/*; do
		[[ -f $hook && -x $hook && $hook != *.disabled ]] || continue
		rc=0
		"$hook" || rc=$?
		((rc == 0)) && continue
		logger -p err -t "$SCRIPT_NAME" "$dir hook failed (exit code $rc): $hook"
		if ((rc >= 100)); then
			error_msg "$dir hook failed (fatal, exit code $rc): $hook"
			return 1
		fi
		warning_msg "$dir hook failed (exit code $rc): $hook"
	done
	return 0
}

prompt_reboot_after_restore() {
	local lockfile="$1" post_failed="${2:-0}"
	[[ -f "$lockfile" && "$(<"$lockfile")" == "restored" ]] || return 1

	if [[ $post_failed -ne 0 ]] || ! command -v reboot &>/dev/null; then
		echo -e "\n${_color_prompt}Restore complete. Please reboot manually. Press Enter to exit:${_color_reset} \c"
		read -r answer
		return 0
	fi
	echo -e "\n${_color_prompt}Restore complete. Reboot now? [Y/n]:${_color_reset} \c"
	read -r answer
	case "$answer" in
	"" | [Yy]*) reboot ;;
	esac
}

# Debounce multiple concurrent requests.
# Only the first caller becomes owner and waits until no new request, then it proceeds.
# All other callers only refresh the lock-timestamp and should exit immediately.
# This improves performance by grouping many snapshot deletions into one sync instead of running many syncs.
debounce_request_lock() {
	local debounce_ms=800 # debounce window in milliseconds
	local lock="/dev/shm/limine-snapper-debounce.lock"
	local now_ms last_update lock_age fd

	# Ensure the lock file exists
	: >"$lock"

	# Open lock file descriptor and try to become owner
	exec {fd}>"$lock" || return 1
	if flock -n "$fd"; then
		# The first caller becomes owner
		while :; do
			# Poll interval
			sleep 0.2
			now_ms=$(date +%s%3N)
			# Read last timestamp of the lock file
			if ! last_update=$(date -r "$lock" +%s%3N 2>/dev/null); then
				# Reading timestamp failed -> stop waiting and proceed
				flock -u "$fd"
				rm -f "$lock"
				return 0
			fi
			lock_age=$((now_ms - last_update))
			# If no updates arrived within debounce window -> proceed
			if ((lock_age > debounce_ms)); then
				# release lock
				flock -u "$fd"
				rm -f "$lock"
				return 0
			fi
		done
	else
		# Other callers are not owner -> refresh timestamp and exit
		: >"$lock"
		return 1
	fi
}

# Main logic
if [[ -e "${RESTORE_LOCKFILE}" ]]; then
	warning_msg "limine-snapper-restore is already running."
	exit 1
fi

trap cleanup_all EXIT

if $RESTORE_MODE; then
	# Create lock file
	: >"$RESTORE_LOCKFILE"
fi

if $ENABLE_DEBOUNCE; then
	debounce_request_lock || exit 0
fi

if $ENABLE_MUTEX; then
	mutex_lock "${SCRIPT_NAME}"
fi

if $ENABLE_HOOKS && ! run_boot_hooks pre; then
	echo "Aborting $HOOK_CALLER: pre hook failed"
	exit 2
fi

/usr/lib/limine/limine-snapper-sync "${ARGS[@]}"
exit_code="$?"

post_failed=0
if $ENABLE_HOOKS; then
	run_boot_hooks post || post_failed=$?
fi

if $RESTORE_MODE; then
	prompt_reboot_after_restore "$RESTORE_LOCKFILE" "$post_failed"
fi

exit "$exit_code"
