#!/bin/bash
# rescue.sh
# Rescue script for EndeavourOS/Arch from Ubuntu
# Features:
# - Auto-detect root partition
# - Optional EFI and /home partitions
# - Filesystem-aware fsck
# - Repair menu (update, GRUB reinstall, package reinstall, shell)
# - Logging and summary report

set -e

VERSION="1.0"
CREATED="10-03-2025"
MODIFIED="Fri October 03 2025 - 12:21 AM"
START_TIME=$(date '+%Y-%m-%d_%H-%M-%S')
BOOTSTRAP_LOG="$PWD/rescue-bootstrap-$START_TIME.log"
LOGFILE="$BOOTSTRAP_LOG" # temporary until we know the distro
# Default settings & Flags for summary
UPDATED=false
GRUB_REINSTALLED=false
PACKAGES_REINSTALLED=false
SHELL_ENTERED=false
FSCK_RUN=false
HOME_MOUNTED=false
INITRAMFS_REBUILT=false
ORPHAN_DISABLED=false
CHROOT_ACTIONS=()
VERBOSE_LOGGING=true
DEBUGMODE=false
DRYRUN=false

# # Show a help dialog
show_help() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Options:
  --verbose       Enable verbose logging (default)
  --quiet         Disable verbose logging (lean audit trail)
  --no-sudo       Skip sudo auto-elevation (debug mode)
  --help, -h      Show this help message and exit

Examples:
  $(basename "$0") --quiet
  $(basename "$0") --verbose
  $(basename "$0") --no-sudo

EOF
}

# Parse all arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
    --verbose)
        VERBOSE_LOGGING=true
        ;;
    --version)
        echo "$(basename "$0") version $VERSION created $CREATED - updated $MODIFIED"
        exit 0
        ;;
    --quiet)
        VERBOSE_LOGGING=false
        ;;
    --no-sudo)
        DEBUGMODE=true
        ;;
    --dry-run)
        DRYRUN=true
        ;;
    --help | -h)
        show_help
        exit 0
        ;;
    *)
        echo "Unknown option: $1"
        show_help
        exit 1
        ;;
    esac
    shift
done

# Root enforcement (only if not in debug mode and not just showing help)
if [ "$DEBUGMODE" = false ] && [ "$EUID" -ne 0 ]; then
    echo "This script requires root privileges. Re-executing with sudo..."
    exec sudo "$0" "$@"
fi

# Log output to the $LOGFILE
log() {
    local msg="$*"
    if [ "$VERBOSE_LOGGING" = true ]; then
        # Full chatter: timestamp + message
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] $msg" | tee -a "$LOGFILE"
    else
        # Lean mode: only log important markers (skip empty chatter)
        if [ -n "$msg" ]; then
            echo "[$(date '+%Y-%m-%d %H:%M:%S')] $msg" | tee -a "$LOGFILE"
        fi
    fi
}

# In debug mode?
if [ "$DEBUGMODE" == "true" ]; then
    log "=== DEBUG MODE: running without sudo ==="
fi

# Resolve UUID= or LABEL= into a real /dev/... path
resolve_device() {
    local spec="$1"
    if [[ "$spec" =~ ^UUID= ]]; then
        local uuid="${spec#UUID=}"
        blkid -U "$uuid" 2>/dev/null
    elif [[ "$spec" =~ ^LABEL= ]]; then
        local label="${spec#LABEL=}"
        blkid -L "$label" 2>/dev/null
    else
        # Already a /dev/... path
        echo "$spec"
    fi
}

# Filesystem-aware fsck
run_fsck() {
    local part=$1
    local fstype
    fstype=$(lsblk -no FSTYPE "$part")

    log "Detected filesystem on $part: $fstype"

    case "$fstype" in
    ext4 | ext3 | ext2)
        if sudo tune2fs -l "$part" | grep -q 'orphan_file'; then
            log "Partition $part has orphan_file enabled"
            if ! sudo fsck.ext4 -n "$part" 2>&1 | grep -q 'unsupported feature'; then
                sudo fsck.ext4 -f -y "$part" | tee -a "$LOGFILE"
            else
                log "e2fsck too old for orphan_file, disabling temporarily"
                sudo tune2fs -O ^orphan_file "$part"
                ORPHAN_DISABLED=true
                sudo e2fsck -f "$part" | tee -a "$LOGFILE"
            fi
        else
            sudo fsck.ext4 -f -y "$part" | tee -a "$LOGFILE"
        fi
        ;;
    btrfs)
        log "Running btrfs check (non-repair mode for safety)"
        sudo btrfs check "$part" | tee -a "$LOGFILE"
        ;;
    xfs)
        log "Running xfs_repair"
        sudo xfs_repair "$part" | tee -a "$LOGFILE"
        ;;
    vfat | fat32)
        log "Running fsck.vfat"
        sudo fsck.vfat -a "$part" | tee -a "$LOGFILE"
        ;;
    *)
        log "Unknown FS type ($fstype), falling back to generic fsck"
        sudo fsck -f -y "$part" | tee -a "$LOGFILE"
        ;;
    esac
}

update_arch_mirrors() {
    log "Updating Arch Linux mirrorlist using reflector..."

    if ! command -v reflector >/dev/null 2>&1; then
        log "Error: reflector is not installed. Installing it now..."
        pacman -Sy --noconfirm reflector || {
            log "Failed to install reflector. Cannot update mirrors."
            return 1
        }
    fi

    log "Running reflector to fetch the 20 latest HTTPS mirrors..."
    if reflector --latest 20 --protocol https --sort rate \
        --save /etc/pacman.d/mirrorlist >>"$FINAL_LOG" 2>&1; then
        log "Mirrorlist successfully updated."
    else
        log "Error occurred while updating mirrorlist. Check logs for details."
        return 1
    fi
}

global_fsck() {
    CHECKED=0 REPAIRED=0 SKIPPED=0

    echo "Available disks:"
    lsblk -d -o NAME,SIZE,MODEL
    read -rp "Enter disk (e.g. /dev/sda, /dev/nvme0n1): " DISK
    read -rp "Mode: [c]heck only or [r]epair? " MODE

    # Build preview
    PREVIEW=""
    for entry in $(lsblk -ln -o NAME,FSTYPE "$DISK" | awk '$2 != "" {print $1":"$2}'); do
        DEV="/dev/${entry%%:*}"
        FSTYPE="${entry##*:}"
        case "$FSTYPE" in
        ext[234]) TOOL="fsck.ext4" ;;
        xfs) TOOL="xfs_repair" ;;
        btrfs) TOOL="btrfs check" ;;
        vfat | fat*) TOOL="fsck.vfat" ;;
        ntfs) TOOL="ntfsfix" ;;
        *) TOOL="(no handler)" ;;
        esac
        PREVIEW="$PREVIEW\n  $DEV ($FSTYPE) → $TOOL"
    done

    echo -e "Will check:$PREVIEW"
    read -rp "Proceed? [y/N] " CONFIRM
    [ "$CONFIRM" != "y" ] && echo "Aborted." && return

    echo "=== GLOBAL FILESYSTEM CHECK SESSION STARTED for $DISK ===" | tee -a "$LOGFILE"

    # Run checks
    for entry in $(lsblk -ln -o NAME,FSTYPE "$DISK" | awk '$2 != "" {print $1":"$2}'); do
        DEV="/dev/${entry%%:*}"
        FSTYPE="${entry##*:}"

        # Skip unsupported
        case "$FSTYPE" in
        swap | LVM* | zfs)
            log "Skipping $DEV ($FSTYPE)"
            SKIPPED=$((SKIPPED + 1))
            continue
            ;;
        esac

        CHECKED=$((CHECKED + 1))
        log "Checking $DEV ($FSTYPE)"

        case "$FSTYPE" in
        ext[234])
            if [ "$MODE" = "r" ]; then
                fsck.ext4 -f -y "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
            else
                fsck.ext4 -n "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
            fi
            ;;
        xfs)
            if [ "$MODE" = "r" ]; then
                xfs_repair "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
            else
                log "xfs: check-only not supported, skipping $DEV"
                SKIPPED=$((SKIPPED + 1))
            fi
            ;;
        btrfs)
            if [ "$MODE" = "r" ]; then
                btrfs check --repair "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
            else
                btrfs check "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
            fi
            ;;
        vfat | fat*)
            if [ "$MODE" = "r" ]; then
                fsck.vfat -a "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
            else
                fsck.vfat -n "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
            fi
            ;;
        ntfs)
            if [ "$MODE" = "r" ]; then
                ntfsfix "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
            else
                log "ntfs: check-only not supported, skipping $DEV"
                SKIPPED=$((SKIPPED + 1))
            fi
            ;;
        *)
            log "No fsck handler for $FSTYPE on $DEV"
            SKIPPED=$((SKIPPED + 1))
            ;;
        esac
    done

    {
        echo "=== Filesystem Check Summary ==="
        echo "Partitions checked: $CHECKED"
        echo "Partitions repaired: $REPAIRED"
        echo "Partitions skipped/failed: $SKIPPED"
    } | tee -a "$LOGFILE"

    echo "=== GLOBAL FILESYSTEM CHECK SESSION COMPLETE for $DISK ===" | tee -a "$LOGFILE"
}

#log "=== Starting Bill's Rescue Chroot ==="
log "=== Rescue session started at $START_TIME ==="

scan_partitions() {
    log "Scanning for Linux root partitions..."
    local idx=1
    DETECTED_PARTS=()

    # Look at all partitions with Linux-y filesystems
    for part in $(lsblk -lnpo NAME,TYPE | awk '$2=="part"{print $1}'); do
        fstype=$(lsblk -no FSTYPE "$part")
        case "$fstype" in
        ext4 | ext3 | btrfs | xfs)
            sudo mount -o ro "$part" /mnt 2>/dev/null || continue
            if [ -f /mnt/etc/os-release ]; then
                name=$(grep PRETTY_NAME /mnt/etc/os-release | cut -d= -f2- | tr -d '"')
                id=$(grep -i '^ID=' /mnt/etc/os-release | cut -d= -f2 | tr -d '"')
                ver=$(grep -i '^VERSION_ID=' /mnt/etc/os-release | cut -d= -f2 | tr -d '"')
                like=$(grep -i '^ID_LIKE=' /mnt/etc/os-release | cut -d= -f2 | tr -d '"')
                echo "[$idx] $part → $name"
                DETECTED_PARTS[$idx]="$part:$id:$ver:$like"
                idx=$((idx + 1))
            fi
            sudo umount /mnt 2>/dev/null
            ;;
        esac
    done

    # After listing detected partitions
    echo "[$idx] Exit without choosing"
    DETECTED_PARTS[$idx]="EXIT"
    idx=$((idx + 1))
    echo "[$idx] Global filesystem check"
    DETECTED_PARTS[$idx]="GLOBALFSCK"

    if [ ${#DETECTED_PARTS[@]} -eq 0 ]; then
        log "No Linux systems detected automatically."
        return 1
    fi

    read -rp "Select a root partition by number: " choice
    case "${DETECTED_PARTS[$choice]}" in
    EXIT)
        log "User chose to exit without selecting a system"
        echo "Exiting without choosing a root partition."
        return 1 # exit/no partitions
        ;;
    GLOBALFSCK)
        log "=== GLOBAL FILESYSTEM CHECK SESSION STARTED ==="
        log "User chose to perform filesystem checks on all partitions."
        echo "Performing filesystem checks on all partitions (except swap, LVM, etc)."
        global_fsck
        log "=== GLOBAL FILESYSTEM CHECK SESSION COMPLETE ==="

        ROOT_PART=""
        DISTRO_ID="GLOBALFSCK"
        VERSION_ID=""
        ROOT_BASENAME="NOPART"
        FINAL_LOG="$PWD/rescue-${DISTRO_ID}-${ROOT_BASENAME}-${START_TIME}.log"

        return 2 # special code for GLOBALFSCK
        ;;
    *)
        ROOT_PART="${DETECTED_PARTS[$choice]%%:*}"
        DISTRO_ID="$(echo "${DETECTED_PARTS[$choice]}" | cut -d: -f2)"
        VERSION_ID="$(echo "${DETECTED_PARTS[$choice]}" | cut -d: -f3)"

        log "Selected root partition: $ROOT_PART"
        log "Detected distro: $DISTRO_ID $VERSION_ID"
        return 0 # normal partition
        ;;
    esac

    ROOT_PART="${DETECTED_PARTS[$choice]%%:*}"
    DISTRO_ID="$(echo "${DETECTED_PARTS[$choice]}" | cut -d: -f2)"
    VERSION_ID="$(echo "${DETECTED_PARTS[$choice]}" | cut -d: -f3)"

    log "Selected root partition: $ROOT_PART"
    log "Detected distro: $DISTRO_ID $VERSION_ID"
}

ROOT_BASENAME=$(basename "$ROOT_PART")
FINAL_LOG="$PWD/rescue-${DISTRO_ID}-${ROOT_BASENAME}-${START_TIME}.log"

if [ -n "$ROOT_PART" ]; then
    ROOT_BASENAME=$(basename "$ROOT_PART")
    FINAL_LOG="$PWD/rescue-${DISTRO_ID}-${ROOT_BASENAME}-${START_TIME}.log"
else
    log "Error: ROOT_PART is empty, cannot build log filename"
    log "=== WARNING: ROOT_PART was empty, using NOPART placeholder ==="
    FINAL_LOG="$PWD/rescue-${DISTRO_ID}-NOPART-${START_TIME}.log"
    #exit 1
fi

# Merge bootstrap log into final log
{
    echo "=== BOOTSTRAP LOG MERGED INTO FINAL LOG ==="
    cat "$BOOTSTRAP_LOG"
    echo "=== END OF BOOTSTRAP LOG ==="
} >>"$FINAL_LOG"

# Switch logging to the final log
LOGFILE="$FINAL_LOG"
log "Switched logging to $LOGFILE"

# Clean up bootstrap log
rm -f "$BOOTSTRAP_LOG"

log "Logging to $LOGFILE"

# Scan for Linux systems and let user choose
# if ! scan_partitions; then
#     read -rp "Enter root partition manually (e.g. /dev/nvme0n1p2): " ROOT_PART
# fi
scan_partitions
case $? in
0)
    # normal partition selected, continue workflow
    ;;
1)
    # no partitions or user exited
    read -rp "Enter root partition manually (e.g. /dev/nvme0n1p2): " ROOT_PART
    ;;
2)
    # global run complete, stop here
    exit 0
    ;;
esac

# Parse fstab inside the mounted system
if [ -f /mnt/etc/fstab ]; then
    log "Parsing /etc/fstab from $ROOT_PART"

    # EFI
    EFI_SPEC=$(awk '$2=="/boot/efi" {print $1}' /mnt/etc/fstab | head -n1)
    if [ -n "$EFI_SPEC" ]; then
        EFI_PART=$(resolve_device "$EFI_SPEC")
        if [ -n "$EFI_PART" ]; then
            log "Resolved EFI partition: $EFI_SPEC → $EFI_PART"
            sudo mount "$EFI_PART" /mnt/boot/efi
        else
            log "Could not resolve EFI partition: $EFI_SPEC"
        fi
    fi

    # /home
    HOME_SPEC=$(awk '$2=="/home" {print $1}' /mnt/etc/fstab | head -n1)
    if [ -n "$HOME_SPEC" ]; then
        HOME_PART=$(resolve_device "$HOME_SPEC")
        if [ -n "$HOME_PART" ]; then
            log "Resolved /home partition: $HOME_SPEC → $HOME_PART"
            sudo mount "$HOME_PART" /mnt/home
            HOME_MOUNTED=true
        else
            log "Could not resolve /home partition: $HOME_SPEC"
        fi
    fi
fi

# Offer fsck before mounting
read -rp "Run filesystem check/repair on $ROOT_PART before mounting? [y/N]: " FSCK_CHOICE
[[ "$FSCK_CHOICE" =~ ^[Yy]$ ]] && run_fsck "$ROOT_PART" && FSCK_RUN=true

if [ -n "$EFI_PART" ]; then
    read -rp "Run filesystem check/repair on $EFI_PART before mounting? [y/N]: " FSCK_EFI
    [[ "$FSCK_EFI" =~ ^[Yy]$ ]] && run_fsck "$EFI_PART" && FSCK_RUN=true
fi

if [ -n "$HOME_PART" ]; then
    read -rp "Run filesystem check/repair on $HOME_PART before mounting? [y/N]: " FSCK_HOME
    [[ "$FSCK_HOME" =~ ^[Yy]$ ]] && run_fsck "$HOME_PART" && FSCK_RUN=true
fi

# Mount partitions
sudo mount "$ROOT_PART" /mnt
log "Mounted root partition $ROOT_PART at /mnt"

# Detect target distro
if [ -f /mnt/etc/os-release ]; then
    DISTRO=$(grep -i '^ID=' /mnt/etc/os-release | cut -d= -f2 | tr -d '"')
    log "Detected target distro: $DISTRO"
else
    DISTRO="unknown"
    log "Could not detect distro (no /etc/os-release)"
fi

# Set distro-specific commands
case "$DISTRO" in
arch | endeavouros | manjaro)
    PKG_UPDATE="pacman -Syu --noconfirm"
    INITRAMFS_CMD="mkinitcpio -P"
    GRUB_CFG="grub-mkconfig -o /boot/grub/grub.cfg"
    PKG_REINSTALL="pacman -Qnq | xargs -r pacman -S --noconfirm"
    ;;
ubuntu | debian)
    PKG_UPDATE="apt-get update && apt-get upgrade -y"
    INITRAMFS_CMD="update-initramfs -u -k all"
    GRUB_CFG="update-grub"
    PKG_REINSTALL="dpkg --get-selections | awk '{print \$1}' | xargs -r apt-get install --reinstall -y"
    ;;
fedora)
    PKG_UPDATE="dnf upgrade -y"
    INITRAMFS_CMD="dracut --regenerate-all --force"
    GRUB_CFG="grub2-mkconfig -o /boot/grub2/grub.cfg"
    PKG_REINSTALL="rpm -qa --qf '%{NAME}\n' | xargs -r dnf reinstall -y"
    ;;
opensuse* | suse)
    PKG_UPDATE="zypper refresh && zypper update -y"
    INITRAMFS_CMD="dracut --regenerate-all --force"
    GRUB_CFG="grub2-mkconfig -o /boot/grub2/grub.cfg"
    PKG_REINSTALL="rpm -qa --qf '%{NAME}\n' | xargs -r zypper install -y --force"
    ;;
*)
    # Fallback to ID_LIKE if DISTRO_ID not recognized
    FALLBACK_USED=true
    FALLBACK_VALUE="$ID_LIKE"
    log "=== Using fallback family: debian ==="
    if [ -n "$ID_LIKE" ]; then
        for like in $ID_LIKE; do
            case "$like" in
            debian)
                log "Unrecognized distro ($DISTRO_ID), but ID_LIKE includes debian → using Debian/Ubuntu logic"
                FALLBACK_USED=true
                FALLBACK_VALUE="debian"
                PKG_UPDATE="apt-get update && apt-get upgrade -y"
                INITRAMFS_CMD="update-initramfs -u -k all"
                GRUB_CFG="update-grub"
                PKG_REINSTALL="dpkg --get-selections | awk '{print \$1}' | xargs -r apt-get install --reinstall -y"
                break
                ;;
            rhel | fedora)
                log "Unrecognized distro ($DISTRO_ID), but ID_LIKE includes rhel/fedora → using Fedora logic"
                FALLBACK_USED=true
                FALLBACK_VALUE="fedora"
                PKG_UPDATE="dnf upgrade -y"
                INITRAMFS_CMD="dracut --regenerate-all --force"
                GRUB_CFG="grub2-mkconfig -o /boot/grub2/grub.cfg"
                PKG_REINSTALL="rpm -qa --qf '%{NAME}\n' | xargs -r dnf reinstall -y"
                break
                ;;
            suse)
                log "Unrecognized distro ($DISTRO_ID), but ID_LIKE includes suse → using openSUSE logic"
                FALLBACK_USED=true
                FALLBACK_VALUE="suse"
                PKG_UPDATE="zypper refresh && zypper update -y"
                INITRAMFS_CMD="dracut --regenerate-all --force"
                GRUB_CFG="grub2-mkconfig -o /boot/grub2/grub.cfg"
                PKG_REINSTALL="rpm -qa --qf '%{NAME}\n' | xargs -r zypper install -y --force"
                break
                ;;
            arch)
                log "Unrecognized distro ($DISTRO_ID), but ID_LIKE includes arch → using Arch logic"
                FALLBACK_USED=true
                FALLBACK_VALUE="arch"
                PKG_UPDATE="pacman -Syu --noconfirm"
                INITRAMFS_CMD="mkinitcpio -P"
                GRUB_CFG="grub-mkconfig -o /boot/grub/grub.cfg"
                PKG_REINSTALL="pacman -Qnq | xargs -r pacman -S --noconfirm"
                break
                ;;
            esac
        done
    else
        log "No matching handler for DISTRO_ID=$DISTRO_ID and ID_LIKE is empty"
    fi
    ;;
esac

[ -n "$EFI_PART" ] && sudo mount "$EFI_PART" /mnt/boot/efi && log "Mounted EFI partition $EFI_PART at /mnt/boot/efi"
if [ -n "$HOME_PART" ]; then
    sudo mount "$HOME_PART" /mnt/home
    log "Mounted /home partition $HOME_PART at /mnt/home"
    HOME_MOUNTED=true
fi

# Bind system dirs
for dir in /dev /dev/pts /proc /sys /run; do
    sudo mount --bind $dir /mnt$dir
done
log "Bound system directories into chroot"

# pass any combination of functions and commands
# run_in_chroot "update_arch_mirrors && pacman -Syu --noconfirm"
run_in_chroot() {
    local action="$*"

    # Ensure bind mounts
    for d in /dev /proc /sys /run; do
        mountpoint -q "/mnt$d" || sudo mount --bind "$d" "/mnt$d"
    done

    # Export all functions defined in this script so they’re visible in chroot
    for fn in $(declare -F | awk '{print $3}'); do
        export -f "$fn"
    done

    # Track the action for later reporting
    CHROOT_ACTIONS+=("$action")

    if [ "$DRYRUN" = true ]; then
        log "[DRY-RUN] Would execute inside chroot: $action"
        return 0
    fi

    log "=== Executing inside chroot: $action ==="
    if sudo chroot /mnt bash -c "$action" >>"$FINAL_LOG" 2>&1; then
        log "=== Completed successfully: $action ==="
    else
        log "** Error while executing inside chroot: $action **"
    fi
}

# Repair menu
while true; do
    echo
    echo ">>> Rescue Menu (inside $DISTRO)"
    echo "1) Update package database & install updates"
    echo "2) Reinstall GRUB (UEFI)"
    echo "3) Reinstall all currently installed packages"
    echo "4) Drop into full chroot shell"
    echo "5) Run fsck on root/EFI/home partitions (unmounts temporarily)"
    echo "6) Rebuild initramfs and regenerate GRUB config"
    echo "7) Export package list"
    echo "8) Show restore examples"
    echo "9) Exit and cleanup"
    #echo "G) Global filesystem check (c), fix optional (r)"

    read -rp "Choose an option: " choice

    case "$choice" in
    1)
        if [[ "$DISTRO" == "arch" ]] || [[ "$ID_LIKE" == "arch" ]]; then
            log "=== Updating mirrors for arch based systems ==="
            run_in_chroot update_arch_mirrors || log "** Unable to update the mirrors as requested. Continuing with the currently chosen mirrors. **"
        fi
        run_in_chroot "$PKG_UPDATE"
        UPDATED=true
        ;;
    2)
        run_in_chroot "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Rescue && $GRUB_CFG"
        GRUB_REINSTALLED=true
        ;;
    3)
        echo "⚠️  This will reinstall ALL packages for $DISTRO."
        read -rp "Proceed? [y/N]: " CONFIRM_REINSTALL
        if [[ "$CONFIRM_REINSTALL" =~ ^[Yy]$ ]]; then
            log "User confirmed full package reinstall"
            run_in_chroot "$PKG_REINSTALL"
            PACKAGES_REINSTALLED=true
        else
            log "User declined full package reinstall"
            echo "Skipped package reinstall."
        fi
        ;;
        #run_in_chroot "$PKG_REINSTALL"; PACKAGES_REINSTALLED=true ;;
    4)
        log "Dropping into interactive chroot shell"
        sudo chroot /mnt /bin/bash
        SHELL_ENTERED=true
        ;;
    5)
        log "Unmounting for fsck..."
        for dir in /run /sys /proc /dev/pts /dev; do
            sudo umount /mnt$dir || true
        done
        [ -n "$HOME_PART" ] && sudo umount /mnt/home || true
        [ -n "$EFI_PART" ] && sudo umount /mnt/boot/efi || true
        sudo umount /mnt || true

        run_fsck "$ROOT_PART"
        FSCK_RUN=true
        [ -n "$EFI_PART" ] && run_fsck "$EFI_PART" && FSCK_RUN=true
        [ -n "$HOME_PART" ] && run_fsck "$HOME_PART" && FSCK_RUN=true

        # Remount
        sudo mount "$ROOT_PART" /mnt
        [ -n "$EFI_PART" ] && sudo mount "$EFI_PART" /mnt/boot/efi
        if [ -n "$HOME_PART" ]; then
            sudo mount "$HOME_PART" /mnt/home
            HOME_MOUNTED=true
        fi
        for dir in /dev /dev/pts /proc /sys /run; do
            sudo mount --bind $dir /mnt$dir
        done
        log "Remounted after fsck"
        ;;
    6)
        run_in_chroot "$INITRAMFS_CMD"
        INITRAMFS_REBUILT=true
        run_in_chroot "$GRUB_CFG"
        ;;
    7)
        # Copy package lists to the host, after exiting chroot
        copy_out() {
            if [ -f /mnt/root/pkglist.txt ]; then
                sudo cp /mnt/root/pkglist.txt "$PWD/pkglist-${DISTRO}-${START_TIME}.txt"
                log "Exported pkglist.txt to host working directory"
            else
                log "Warning: /mnt/root/pkglist.txt not found"
            fi
            if [ -f /mnt/root/aurlist.txt ]; then
                sudo cp /mnt/root/aurlist.txt "$PWD/aurlist-${DISTRO}-${START_TIME}.txt"
                log "Exported aurlist.txt to host working directory"
            else
                log "Warning: /mnt/root/aurlist.txt not found"
            fi
            if [ -f /mnt/root/pkglist-full.txt ]; then
                sudo cp /mnt/root/pkglist-full.txt "$PWD/pkglist-${DISTRO}-full-${START_TIME}.txt"
                log "Exported pkglist-full.txt to host working directory"
            else
                log "Warning: /mnt/root/pkglist-full.txt not found"
            fi
            if [ -f /mnt/root/pkglist-explicit.txt ]; then
                sudo cp /mnt/root/pkglist-explicit.txt "$PWD/pkglist-${DISTRO}-explicit-${START_TIME}.txt"
                log "Exported pkglist-explicit.txt to host working directory"
            else
                log "Warning: /mnt/root/pkglist-explicit.txt not found"
            fi
            log "Finished exporting package list(s) to host working directory"
        }

        log "Exporting package list for $DISTRO"
        case "$DISTRO" in
        arch | endeavouros | manjaro)
            EXPORT_TIME=$(date '+%Y-%m-%d %H:%M:%S')
            log "=== Package list export started at $EXPORT_TIME ==="
            run_in_chroot "pacman -Qqe > /root/pkglist.txt && pacman -Qqm > /root/aurlist.txt"
            copy_out
            ;;
        ubuntu | debian)
            EXPORT_TIME=$(date '+%Y-%m-%d %H:%M:%S')
            log "=== Package list export started at $EXPORT_TIME ==="
            run_in_chroot "dpkg --get-selections | awk '{print \$1}' > /root/pkglist.txt"
            copy_out
            ;;
        fedora | rhel)
            EXPORT_TIME=$(date '+%Y-%m-%d %H:%M:%S')
            log "=== Package list export started at $EXPORT_TIME ==="
            run_in_chroot "rpm -qa --qf '%{NAME}\n' > /root/pkglist.txt"
            copy_out
            ;;
        opensuse* | suse)
            EXPORT_TIME=$(date '+%Y-%m-%d %H:%M:%S')
            log "=== Package list export started at $EXPORT_TIME ==="
            # Full list of all installed packages
            run_in_chroot "rpm -qa --qf '%{NAME}\n' > /root/pkglist-full.txt"
            # Explicitly installed packages (user-installed, not deps)
            run_in_chroot "zypper se --installed-only --type package --details | awk '{print \$3}' > /root/pkglist-explicit.txt"
            copy_out
            ;;
        *)
            log "Package export not supported for this distro"
            ;;
        esac
        ;;
    8)
        if [ -f restore_examples ]; then
            cat ./restore_examples
        fi || echo "Restore examples not found"
        ;;
    9)
        log "Exiting rescue menu"
        break
        ;;
        # G | g)
        #     CHECKED=0
        #     REPAIRED=0
        #     SKIPPED=0

        #     echo "Available disks:"
        #     lsblk -d -o NAME,SIZE,MODEL
        #     read -rp "Enter disk (e.g. /dev/sda, /dev/nvme0n1): " DISK

        #     read -rp "Mode: [c]heck only or [r]epair? " MODE

        #     # Build preview list
        #     PREVIEW=""
        #     for entry in $(lsblk -ln -o NAME,FSTYPE "$DISK" | awk '$2 != "" {print $1":"$2}'); do
        #         DEV="/dev/${entry%%:*}"
        #         FSTYPE="${entry##*:}"
        #         case "$FSTYPE" in
        #         ext[234]) TOOL="fsck.ext4" ;;
        #         xfs) TOOL="xfs_repair" ;;
        #         btrfs) TOOL="btrfs check" ;;
        #         vfat | fat*) TOOL="fsck.vfat" ;;
        #         ntfs) TOOL="ntfsfix" ;;
        #         *) TOOL="(no handler)" ;;
        #         esac
        #         PREVIEW="$PREVIEW\n  $DEV ($FSTYPE) → $TOOL"
        #     done

        #     echo -e "Will check:$PREVIEW"
        #     read -rp "Proceed? [y/N] " CONFIRM
        #     [ "$CONFIRM" != "y" ] && echo "Aborted." && continue

        #     # Run checks
        #     for entry in $(lsblk -ln -o NAME,FSTYPE "$DISK" | awk '$2 != "" {print $1":"$2}'); do
        #         DEV="/dev/${entry%%:*}"
        #         FSTYPE="${entry##*:}"
        #         CHECKED=$((CHECKED + 1))
        #         log "Checking $DEV ($FSTYPE)"
        #         case "$FSTYPE" in
        #         ext[234])
        #             if [ "$MODE" = "r" ]; then
        #                 if fsck.ext4 -f -y "$DEV" | tee -a "$LOGFILE"; then
        #                     REPAIRED=$((REPAIRED + 1))
        #                 else
        #                     SKIPPED=$((SKIPPED + 1))
        #                 fi
        #             else
        #                 fsck.ext4 -n "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
        #             fi
        #             ;;
        #         xfs)
        #             if [ "$MODE" = "r" ]; then
        #                 xfs_repair "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
        #             else
        #                 log "xfs: check-only not supported, skipping $DEV"
        #                 SKIPPED=$((SKIPPED + 1))
        #             fi
        #             ;;
        #         btrfs)
        #             if [ "$MODE" = "r" ]; then
        #                 btrfs check --repair "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
        #             else
        #                 btrfs check "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
        #             fi
        #             ;;
        #         vfat | fat*)
        #             if [ "$MODE" = "r" ]; then
        #                 fsck.vfat -a "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
        #             else
        #                 fsck.vfat -n "$DEV" | tee -a "$LOGFILE" || SKIPPED=$((SKIPPED + 1))
        #             fi
        #             ;;
        #         ntfs)
        #             if [ "$MODE" = "r" ]; then
        #                 ntfsfix "$DEV" | tee -a "$LOGFILE" && REPAIRED=$((REPAIRED + 1)) || SKIPPED=$((SKIPPED + 1))
        #             else
        #                 log "ntfs: check-only not supported, skipping $DEV"
        #                 SKIPPED=$((SKIPPED + 1))
        #             fi
        #             ;;
        #        *)
        #            log "No fsck handler for $FSTYPE on $DEV"
        #            SKIPPED=$((SKIPPED + 1))
        #            ;;
        #     esac
        # done

        # # Summary block
        # echo "=== Filesystem Check Summary ==="
        # echo "Partitions checked: $CHECKED"
        # echo "Partitions repaired: $REPAIRED"
        # echo "Partitions skipped/failed: $SKIPPED"

        # {
        #     echo "=== Filesystem Check Summary ==="
        #     echo "Partitions checked: $CHECKED"
        #     echo "Partitions repaired: $REPAIRED"
        #     echo "Partitions skipped/failed: $SKIPPED"
        # } >>"$LOGFILE"

        # log "=== Global filesystem check complete for $DISK ==="
        # ;;

    *) echo "Invalid choice" ;;
    esac
done

log "Cleaning up mounts..."

# Unmount bind mounts in reverse order
for dir in /run /sys /proc /dev/pts /dev; do
    if mountpoint -q /mnt$dir; then
        sudo umount /mnt$dir || log "Failed to unmount /mnt$dir"
    fi
done

# Unmount /home if mounted
if [ -n "$HOME_PART" ] && mountpoint -q /mnt/home; then
    sudo umount /mnt/home || log "Failed to unmount /mnt/home"
fi

# Unmount EFI if mounted
if [ -n "$EFI_PART" ] && mountpoint -q /mnt/boot/efi; then
    sudo umount /mnt/boot/efi || log "Failed to unmount /mnt/boot/efi"
fi

# Finally unmount root
if mountpoint -q /mnt; then
    if ! sudo umount /mnt; then
        log "/mnt still busy, forcing lazy unmount"
        sudo umount -l /mnt
    fi
fi

if $ORPHAN_DISABLED; then
    log "Re-enabling orphan_file feature on $ROOT_PART"
    sudo tune2fs -O orphan_file "$ROOT_PART"
    [ -n "$HOME_PART" ] && sudo tune2fs -O orphan_file "$HOME_PART"
fi

END_TIME=$(date '+%Y-%m-%d %H:%M:%S')
log "=== Chroot session finished ==="

# Show actions ran in chroot before the summary
if [ ${#CHROOT_ACTIONS[@]} -gt 0 ]; then
    echo "=== Chroot Actions Executed ===" | tee -a "$FINAL_LOG"
    for act in "${CHROOT_ACTIONS[@]}"; do
        echo " - $act" | tee -a "$FINAL_LOG"
    done
    echo "" | tee -a "$FINAL_LOG"
fi

# Print summary
echo
echo "=== Rescue Session Summary ==="
echo "Start time: $START_TIME"
echo "End time:   $END_TIME"
echo "Root partition: $ROOT_PART"
[ -n "$EFI_PART" ] && echo "EFI partition: $EFI_PART"
[ -n "$HOME_PART" ] && echo "/home partition: $HOME_PART"
echo "Actions performed:"
$UPDATED && echo "  - Package database updated"
$GRUB_REINSTALLED && echo "  - GRUB reinstalled"
$INITRAMFS_REBUILT && echo "  - Initramfs rebuilt and GRUB config regenerated"
$PACKAGES_REINSTALLED && echo "  - All packages reinstalled"
$SHELL_ENTERED && echo "  - Interactive chroot shell entered"
$ORPHAN_DISABLED && echo "  - orphan_file temporarily disabled for fsck and re-enabled after cleanup"
$FSCK_RUN && echo "  - Filesystem check/repair run"
$HOME_MOUNTED && echo "  - /home partition mounted"
echo "Log file: $LOGFILE"

# If package lists were exported, log them in the footer
EXPORTED_FILES=$(ls "$PWD"/pkglist-* "$PWD"/aurlist-* 2>/dev/null | xargs -n1 basename)
if [ -n "$EXPORTED_FILES" ]; then
    log "=== Package list exported: $EXPORTED_FILES ==="
fi

if [ "$FALLBACK_USED" = true ]; then
    log "=== Fallback used: ID_LIKE=$FALLBACK_VALUE ==="
fi
# Append a footer to the log
log "=== SESSION COMPLETE at $END_TIME ==="