#!/usr/bin/env bash

### 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"
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 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; 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;31m Warning: ${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
}

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

trap cleanup EXIT

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

# 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
# Create lock file
: >>"$LOCKFILE"

# 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 [[ $1 == "--notify" ]]; 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

RESTORE_CMD="${JAVA_CMD} -jar /usr/share/java/limine-snapper-sync.jar --restore"

# Notify and execute the restore command
if is_graphical && [[ $1 == "--notify" ]] && ! 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
    bash -c "$RESTORE_CMD"
else
    # Fallback for TTY or non-graphical environments
    bash -c "$AUTH_METHOD $RESTORE_CMD"
fi
