#!/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

# Variables
readonly SCRIPT_NAME="limine-snapper-restore"
readonly BASE_CONFIG="/etc/limine-snapper-sync.conf"
readonly DEFAULT_CONFIG="/etc/default/limine"
readonly TMP_CONFIG="/tmp/limine-snapper-sync.conf"
readonly LOCKFILE="/tmp/limine-snapper-restore.lock"
readonly LIMINE_LOCK_FILE="/tmp/limine-global.lock"

TITLE="Restore this snapshot now!"
MESSAGE="You are currently using this snapshot. Please restore it before rebooting to the normal system."
ENABLE_MUTEX="no"

_color_reset=""
_color_yellow=""
_color_red=""
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"
fi

info_msg() {
	echo "$1"
}

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 [[ -n "${SNAPSHOT_ID}" ]]; then
	if [ "$EUID" -ne 0 ]; then
		error_msg "Root privileges are required to restore snapshot ID '$SNAPSHOT_ID'"
		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 "${BASE_CONFIG}" ]]; then
	load_key_value_config "${BASE_CONFIG}"
fi

if [[ -f "${DEFAULT_CONFIG}" ]]; then
	load_key_value_config "${DEFAULT_CONFIG}"
fi

# 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
				warning_msg "${TMP_CONFIG} is not owned by root, this file is ignored."
			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 || {
		warning_msg "Mutex lock timeout on ${name}."
		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
		error_msg "notify-send or dunst is not installed."
		logger -t "${SCRIPT_NAME}" -p err "notify-send or dunst is not installed."
		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.
			cmd_str="${AUTH_METHOD} ${RESTORE_CMD[*]}"
			$terminal_cmd "$cmd_str"
		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
		error_msg "No suitable terminal found."
		logger -t "${SCRIPT_NAME}" -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
		error_msg "Root privileges are required."
		exit 1
	fi
fi

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

if [[ "$RESTORE_KERNELS_ONLY" == "yes" ]]; then
	ENABLE_MUTEX="yes"
	RESTORE_CMD=(/usr/lib/limine/limine-snapper-sync --restore-kernels "${SNAPSHOT_ID}")
else
	ENABLE_MUTEX="no"
	RESTORE_CMD=(/usr/lib/limine/limine-snapper-sync --restore)
fi

if [[ ${ENABLE_MUTEX} == "yes" ]]; then
	mutex_lock "${SCRIPT_NAME}"
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"
