#!/usr/bin/env bash

ENABLE_NOTIFY="no"
SNAPSHOT_ID=""
RESTORE_KERNELS_ONLY="no"

# Parse arguments
while [[ $# -gt 0 ]]; do
	case "$1" in
	--kernels)
		if [[ -z "$2" || "$2" == -* ]]; then
			echo "Error: --kernels requires a Snapper snapshot ID." >&2
			exit 1
		fi
		SNAPSHOT_ID="$2"
		RESTORE_KERNELS_ONLY="yes"
		shift 2
		;;
	--notify)
		ENABLE_NOTIFY="yes"
		shift
		;;
	--)
		shift
		break
		;;
	-*)
		echo "Unknown option: $1" >&2
		echo "Usage: $(basename "$0") [options]"
		echo "  --kernels <ID>    Restore kernel files from the specified Snapper snapshot ID"
		echo "  --notify          Send desktop notifications as a non-root user"
		exit 1
		;;
	*)
		break
		;;
	esac
done

if [[ -n "${SNAPSHOT_ID}" ]]; then
	if [ "$EUID" -ne 0 ]; then
		echo -e "\033[91m Error: Root privileges are required to restore snapshot ID '$SNAPSHOT_ID'\033[0m" >&2
		exit 1
	fi
fi

### Load key=value config
load_key_value_config() {
	local file="$1"
	while IFS= read -r line || [[ -n "$line" ]]; do
		# Skip blank lines and comments
		[[ "$line" =~ ^[[:space:]]*(#.*)?$ ]] && continue

		# Match VAR = value
		if [[ "$line" =~ ^[[:space:]]*([A-Za-z_][A-Za-z0-9_]*)[[:space:]]*=[[:space:]]*(.*)$ ]]; then
			local var="${BASH_REMATCH[1]}"
			local val="${BASH_REMATCH[2]}"
			val="${val%\"}"
			val="${val#\"}"
			export "$var"="$val"
		fi
	done <"$file"
}

# Load configuration
if [[ -f /etc/limine-snapper-sync.conf ]]; then
	load_key_value_config /etc/limine-snapper-sync.conf
fi

if [[ -f /etc/default/limine ]]; then
	load_key_value_config /etc/default/limine
fi

# Variables
TMP_CONFIG="/tmp/limine-snapper-sync.conf"
BASE_CONFIG="/etc/limine-snapper-sync.conf"
DEFAULT_CONFIG="/etc/default/limine"
LOCKFILE="/tmp/limine-snapper-restore.lock"
LIMINE_LOCK_FILE="/tmp/limine-global.lock"
ENABLE_MUTEX="no"
TITLE="Restore this snapshot now!"
MESSAGE="You are currently using this snapshot. Please restore it before rebooting to the normal system."

# Check if you are in snapshot
is_snapshot() {
	cmdline=$(</proc/cmdline)
	if [[ $cmdline =~ rootflags.*subvol=.*?/([0-9]+)/snapshot ]]; then
		return 0
	else
		return 1
	fi
}

# Check if you are running in root
is_root() {
	[[ $EUID -eq 0 ]]
}

# Check if the snapshot is read-only and load temporary config
check_and_load_config() {
	# Read the Btrfs property "ro" from the root partition
	local prop
	prop=$(btrfs property get / ro 2>/dev/null)
	# Check if "true" is included in the result
	if [[ $prop == *true* ]] && is_snapshot && is_root; then
		# Load the temporary config if it exists
		if [[ -f "${TMP_CONFIG}" ]]; then
			if [[ $(stat -c '%U' "${TMP_CONFIG}") == "root" ]]; then
				load_key_value_config "${TMP_CONFIG}"
			else
				echo -e "\033[1;33mWarning: ${TMP_CONFIG} is not owned by root, this file is ignored.\033[0m" >&2
			fi
		elif [[ -f "${DEFAULT_CONFIG}" ]]; then
			cp "${DEFAULT_CONFIG}" "${TMP_CONFIG}"
		elif [[ -f "${BASE_CONFIG}" ]]; then
			cp "${BASE_CONFIG}" "${TMP_CONFIG}"
		fi
	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 || {
		echo -e "\033[1;31m Mutex lock timeout on ${name}.\033[0m" >&2
		return 1
	}
}

mutex_unlock() {
	# Release the lock
	flock --unlock 200
}

# Cleanup lockfile on exit
cleanup() {
	rm -f "$LOCKFILE"
}

trap cleanup EXIT

# Function to check if the script is running in a graphical environment
is_graphical() {
	[[ $XDG_SESSION_TYPE == "x11" || $XDG_SESSION_TYPE == "wayland" ]]
}

get_terminal_command() {
	if [[ -n "$TERMINAL" ]] && command -v "$TERMINAL" &>/dev/null; then
		# Set a default TERMINAL_ARG if it's not specified
		case "$TERMINAL" in
		gnome-terminal)
			echo "$TERMINAL ${TERMINAL_ARG:--- bash -c}"
			;;
		*)
			echo "$TERMINAL ${TERMINAL_ARG:--e}"
			;;
		esac
	elif command -v konsole &>/dev/null; then
		echo konsole -e
	elif command -v gnome-terminal &>/dev/null; then
		echo gnome-terminal -- bash -c
	elif command -v xfce4-terminal &>/dev/null; then
		echo xfce4-terminal -e
	elif command -v qterminal &>/dev/null; then
		echo qterminal -e
	elif command -v mate-terminal &>/dev/null; then
		echo mate-terminal -e
	elif command -v kgx &>/dev/null; then
		echo kgx -e
	elif command -v deepin-terminal &>/dev/null; then
		echo deepin-terminal -e
	elif command -v foot &>/dev/null; then
		echo foot
	elif command -v kitty &>/dev/null; then
		echo kitty
	elif command -v xterm &>/dev/null; then
		echo xterm -e
	else
		return 1
	fi
}

# Notify using dunstify or notify-send
notify_user() {
	local action
	if command -v dunstify &>/dev/null; then
		action=$(dunstify -u "critical" --appname="Snapshot detected!" --icon="$NOTIFICATION_ICON" --action="default,Reply" --action="openRestore,Restore now" "$TITLE" "$MESSAGE" -t 0)
	elif command -v notify-send &>/dev/null; then
		action=$(notify-send -u "critical" --app-name="Snapshot detected!" --icon="$NOTIFICATION_ICON" --action="default=Reply" --action="openRestore=Restore now" "$TITLE" "$MESSAGE" -t 0)
	else
		echo -e "\033[91m notify-send or dunst is not installed.\033[0m" >&2
		exit 1
	fi
	echo "$action"
}

open_terminal() {
	terminal_cmd=$(get_terminal_command)
	if [[ -n $terminal_cmd ]]; then
		if [[ $terminal_cmd == gnome-terminal* ]] || [[ $terminal_cmd == xfce4-terminal* ]] || [[ $terminal_cmd == mate-terminal* ]]; then
			# Note: gnome-terminal, xfce4-terminal and mate-terminal, which are based on GTK, require quoted string.
			$terminal_cmd "${AUTH_METHOD} ${RESTORE_CMD}"
		else
			# Note: Do not use quoted strings, as some terminals (e.g., foot, alacritty, wezterm, kitty) fail to parse them correctly.
			$terminal_cmd ${AUTH_METHOD} ${RESTORE_CMD}
		fi
	else
		echo -e "\033[1;31m No suitable terminal found.\033[0m" >&2
		logger -t limine-snapper-restore -p err "No suitable terminal found."
		exit 1
	fi
}

# Main logic
check_and_load_config

# Determine authentication method
if [[ -z $AUTH_METHOD ]] && ! is_root; then
	if is_graphical && command -v pkexec &>/dev/null; then
		AUTH_METHOD="pkexec"
	elif command -v sudo &>/dev/null; then
		AUTH_METHOD="sudo"
	else
		echo -e "\033[91m Root privileges are required.\033[0m" >&2
		exit 1
	fi
fi

# Handle notification mode
if [[ ${ENABLE_NOTIFY} == "yes" ]]; then
	if ! is_snapshot; then
		echo "You are not in a snapshot."
		exit 0
	fi
fi

# Check if JAVA_BIN_PATH is set and valid
CHECK_JAVA_VERSION="/usr/lib/limine/lss-check-java"
# shellcheck disable=SC2034
SCRIPT_NAME="limine-snapper-restore"
if [[ -f "${CHECK_JAVA_VERSION}" ]]; then
	# shellcheck disable=SC1090
	if source "${CHECK_JAVA_VERSION}"; then
		JAVA_CMD=$(get_java_command) || exit 1
	else
		echo -e "\033[1;31m Error: Failed to source '${CHECK_JAVA_VERSION}'.\033[0m" >&2
		JAVA_CMD="java"
	fi
else
	JAVA_CMD="java"
fi

if [[ "$RESTORE_KERNELS_ONLY" == "yes" ]]; then
	ENABLE_MUTEX="yes"
	RESTORE_CMD="${JAVA_CMD} -jar /usr/share/java/limine-snapper-sync.jar --restore-kernels ${SNAPSHOT_ID}"
else
	ENABLE_MUTEX="no"
	RESTORE_CMD="${JAVA_CMD} -jar /usr/share/java/limine-snapper-sync.jar --restore"
fi

if [[ ${ENABLE_MUTEX} == "yes" ]]; then
	mutex_lock "limine-snapper-restore"
else
	# Create lock file
	: >>"$LOCKFILE"
fi

exit_code=0

# Notify and execute the restore command
if is_graphical && [[ ${ENABLE_NOTIFY} == "yes" ]] && ! is_root; then
	action=$(notify_user)
	if [[ $action == "openRestore" || $action == "default" ]]; then
		open_terminal
	fi
elif is_graphical && ! is_root; then
	open_terminal
elif is_root; then
	${RESTORE_CMD}
	exit_code="$?"
else
	# Fallback for TTY or non-graphical environments
	${AUTH_METHOD} ${RESTORE_CMD}
	exit_code="$?"
fi

if [[ ${ENABLE_MUTEX} == "yes" ]]; then
	mutex_unlock
fi

exit "$exit_code"
