#!/bin/bash
# Handles non-interactive restore modes (1-3) for Legacy 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

# ... (log, mount, unmount, execute_chroot functions remain the same) ...
#<editor-fold desc="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 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() {
    umount -l "$2" 2>/dev/null || true
    mkdir -p "$2"
    if mount "$1" -o "$3" "$2"; then log_message "Mounted $1 to $2"; else handle_error "Failed to mount $1"; fi
}
safe_unmount() {
    if umount "$1"; then log_message "Unmounted $1"; else log_message "Warning: Could not unmount $1"; fi
}
execute_chroot() {
    log_message "Chroot: $2"
    show_progress "$2"
    if ! manjaro-chroot /mnt $1; then handle_error "Failed: $2"; 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"
}
#</editor-fold>

# Main script
log_message "Starting Legacy GRUB restore script"

# Read config
DISK_SELECTED=$(cat /tmp/grub-disks-selected)
SELECTED_PARTITION=$(awk -F: '{ print $1}' /tmp/os-prober-selected)
PARTITION_FORMAT=$(awk -F: '{ print $5 }' /tmp/os-prober-selected)
RESTORE_MODE=$(cat /tmp/grub-restore-apply-mode)

# Mount
show_progress "Preparing mount points..."
mount_opts=""
[ "$PARTITION_FORMAT" = "btrfs" ] && mount_opts="subvol=@"
safe_mount "$SELECTED_PARTITION" "/mnt" "$mount_opts"

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

# Execute
case "$RESTORE_MODE" in
    "1")
        log_message "Executing Simple Restore"
        if ! grub-install --root-directory=/mnt "$DISK_SELECTED"; then handle_error "GRUB installation to MBR failed"; fi
        ;;
    "2")
        log_message "Executing Intermediate Restore"
        execute_chroot "rm -f /var/lib/pacman/db.lck" "Removing package lock"
        execute_chroot "pacman --noconfirm -Sy grub" "Updating GRUB package"
        execute_chroot "pacman --noconfirm -S grub-theme-manjaro" "Installing Manjaro GRUB Theme"
        execute_chroot "grub-install --force --target=i386-pc --recheck --boot-directory=/boot $DISK_SELECTED" "Installing GRUB"
        execute_chroot "mkinitcpio -P" "Regenerating initramfs"
        execute_chroot "grub-mkconfig -o /boot/grub/grub.cfg" "Generating GRUB configuration"
        ;;
    "3")
        log_message "Executing Complete Restore"
        execute_chroot "rm -f /var/lib/pacman/db.lck" "Removing package lock"
        execute_chroot "pacman --noconfirm -Syyu" "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 --force --target=i386-pc --recheck --boot-directory=/boot $DISK_SELECTED" "Installing GRUB"
        execute_chroot "mkinitcpio -P" "Regenerating initramfs"
        execute_chroot "grub-mkconfig -o /boot/grub/grub.cfg" "Generating GRUB configuration"
        ;;
    *)
        handle_error "Invalid non-interactive mode: $RESTORE_MODE"
        ;;
esac

# Cleanup
show_progress "Syncing filesystems..."
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

show_progress "Unmounting filesystems..."
safe_unmount "/mnt"

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