#!/bin/bash
# Handles non-interactive restore modes (1-3) for EFI systems.
# Based on BigLinux GRUB Restore (www.biglinux.com.br)
# Updated and modified for Manjaro by Stefano Capitani, 2026
# License: GPL v3 or later

set -e

# --- Helper Functions ---
export TEXTDOMAINDIR="/usr/share/locale"
export TEXTDOMAIN=grub-restore

log_message() { echo "$(date '+%Y-%m-%d %H:%M:%S'): $1" >&2; }
show_progress() { echo "PROGRESS: $1"; }
handle_error() { 
    log_message "ERROR: $1"
    # Unmount everything before exiting on error
    umount -l /mnt/boot/efi 2>/dev/null || true
    umount -l /mnt 2>/dev/null || true
    # Unmount separate partitions
    for mp in $(findmnt -rn -o TARGET | grep "^/mnt/" | sort -r); do
        umount -l "$mp" 2>/dev/null || true
    done
    echo "ERROR: $1" >&2
    exit "${2:-1}"
}

safe_mount() {
    local dev="$1"
    local target="$2"
    local opts="$3"
    
    log_message "Attempting to mount $dev on $target with opts: $opts"
    mkdir -p "$target"
    
    if [ -n "$opts" ]; then
        mount "$dev" "$target" -o "$opts" || handle_error "Failed to mount $dev on $target"
    else
        mount "$dev" "$target" || handle_error "Failed to mount $dev on $target"
    fi
}

execute_chroot() {
    local cmd="$1"
    local desc="$2"
    log_message "Chroot executing: $desc ($cmd)"
    show_progress "$desc"
    # Using bash via stdin to avoid quoting issues in manjaro-chroot/bash -c
    if ! echo "$cmd" | manjaro-chroot /mnt /bin/bash; then 
        handle_error "Failed during: $desc" 
    fi
}

# Function to mount separate partitions from fstab
mount_separate_partitions() {
    local selected_partition="$1"
    local partition_format="$2"
    local separate_file="/tmp/separate-partitions-$(basename "$selected_partition")"
    
    if [ ! -f "$separate_file" ] || [ ! -s "$separate_file" ]; then
        log_message "No separate partitions file found for $selected_partition"
        return 0
    fi
    
    log_message "Mounting separate partitions from $separate_file"
    
    # Read separate partitions file
    # Format: device:mountpoint:fstype:options:uuid:label
    while IFS=: read -r device mountpoint fstype options uuid label; do
        # Skip if empty fields
        [ -z "$device" ] && continue
        [ -z "$mountpoint" ] && continue
        [ ! -b "$device" ] && continue
        
        # Create mount point under /mnt
        local target="/mnt$mountpoint"
        log_message "Mounting separate partition $device ($fstype) on $target"
        
        # Prepare mount options
        local mount_opts="$options"
        [ "$mount_opts" = "defaults" ] && mount_opts=""
        
        # For Btrfs, ensure subvol is correct
        if [ "$fstype" = "btrfs" ] && [[ ! "$mount_opts" == *"subvol="* ]]; then
            # For separate Btrfs partitions, don't use subvol=@ unless specified
            # User should have configured it in fstab
            mount_opts="${mount_opts:+$mount_opts,}defaults"
        elif [ -z "$mount_opts" ]; then
            mount_opts="defaults"
        fi
        
        mkdir -p "$target"
        if mount -o "$mount_opts" "$device" "$target"; then
            log_message "Successfully mounted $device on $target"
        else
            log_message "WARNING: Failed to mount $device on $target, continuing..."
        fi
    done < "$separate_file"
}

# --- Main Script ---
log_message "Starting EFI GRUB restore script"

# Verify temp files exist
[ ! -f /tmp/efi-selected ] && handle_error "EFI partition selection not found."
[ ! -f /tmp/os-prober-selected ] && handle_error "System partition selection not found."

# Read config
EFI_PARTITION=$(cat /tmp/efi-selected)
SELECTED_PARTITION=$(awk -F: '{ print $1 }' /tmp/os-prober-selected)
PARTITION_FORMAT=$(awk -F: '{ print $5 }' /tmp/os-prober-selected) # Field 5 is filesystem (e.g. btrfs)
RESTORE_MODE=$(cat /tmp/grub-restore-apply-mode 2>/dev/null || echo "1")

# --- Preparation ---
show_progress "Preparing mount points..."

# Cleanup old mounts
umount -l /mnt/boot/efi 2>/dev/null || true
umount -l /mnt 2>/dev/null || true

# Btrfs vs Others Mount Logic
if [ "$PARTITION_FORMAT" = "btrfs" ]; then
    # Mount @ subvolume (BigLinux/Manjaro default)
    safe_mount "$SELECTED_PARTITION" "/mnt" "subvol=@,defaults,compress=zstd"
else
    safe_mount "$SELECTED_PARTITION" "/mnt" "defaults"
fi

# Ensure boot/efi directory exists on target before mounting
mkdir -p /mnt/boot/efi
safe_mount "$EFI_PARTITION" "/mnt/boot/efi" "defaults"

# Mount separate partitions (home, var, boot, etc.)
mount_separate_partitions "$SELECTED_PARTITION" "$PARTITION_FORMAT"

# --- Execution ---
case "$RESTORE_MODE" in
    "1")
        log_message "Executing Simple Restore"
        execute_chroot "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=manjaro --recheck" "Installing GRUB"
        execute_chroot "grub-mkconfig -o /boot/grub/grub.cfg" "Updating configuration"
        ;;
    "2")
        log_message "Executing Intermediate Restore"
        execute_chroot "rm -f /var/lib/pacman/db.lck" "Cleaning package lock"
        execute_chroot "pacman --noconfirm -Sy grub" "Reinstalling GRUB package"
        execute_chroot "pacman --noconfirm -S grub-theme-manjaro" "Installing Manjaro GRUB Theme"
        execute_chroot "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=manjaro --recheck" "Installing GRUB"
        execute_chroot "mkinitcpio -P" "Regenerating initramfs"
        execute_chroot "grub-mkconfig -o /boot/grub/grub.cfg" "Generating GRUB config"
        ;;
    "3")
        log_message "Executing Complete Restore"
        execute_chroot "rm -f /var/lib/pacman/db.lck" "Cleaning package lock"
        # Refresh keys before attempting system update in recovery mode
        execute_chroot "pacman-key --init && pacman-key --populate archlinux manjaro" "Refreshing keys"
        execute_chroot "pacman --noconfirm -Syu" "Updating system"
        execute_chroot "pacman --noconfirm -S grub linux" "Installing core packages"
        execute_chroot "pacman --noconfirm -S grub-theme-manjaro" "Installing Manjaro GRUB Theme"
        execute_chroot "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=manjaro --recheck" "Installing GRUB"
        execute_chroot "mkinitcpio -P" "Regenerating initramfs"
        execute_chroot "grub-mkconfig -o /boot/grub/grub.cfg" "Generating GRUB config"
        ;;
    *)
        handle_error "Invalid mode: $RESTORE_MODE"
        ;;
esac

# --- Cleanup ---
show_progress "Finishing process..."
sync

# Unmount separate partitions first (reverse order)
for mp in $(findmnt -rn -o TARGET | grep "^/mnt/" | sort -r); do
    umount -l "$mp" 2>/dev/null || true
done

umount -l /mnt/boot/efi
umount -l /mnt

log_message "EFI GRUB restore completed successfully"
show_progress "Restore operation completed!"
exit 0
