#!/bin/bash
# 2026-09-21
# Description: Automated Btrfs rollback utility for TUXEDO open beta
# Purpose: Replaces the active root subvolume (@) with a stable snapshot.

# 0. Root Check
if [ "$EUID" -ne 0 ]; then
    echo "[-] Error: This script must be run as root (use sudo)!"
    exit 1
fi

echo "[+] Starting Btrfs subvolume rollback tool..."

# 1. Fetch root partition UUID
UUID=$(findmnt -n -o UUID /)
if [ -z "$UUID" ]; then
    echo "[-] Error: Could not determine root partition UUID."
    exit 1
fi

# 2. Auto-detect current snapshot if booted into one
DETECTED_SNAP=""
MOUNT_SOURCE=$(findmnt -n -o SOURCE /)
if [[ "$MOUNT_SOURCE" =~ @\.snapshots/([0-9]+)/snapshot ]]; then
    DETECTED_SNAP="${BASH_REMATCH[1]}"
fi

# 3. Prompt user with the auto-detected default (if available)
if [ -n "$DETECTED_SNAP" ]; then
    SNAP_NUM="${SNAP_NUM:-$DETECTED_SNAP}"
    echo "[+] Booted snapshot detected. Auto-selecting snapshot: $SNAP_NUM"
else
    echo "[?] Could not auto-detect snapshot from mount path."
    
    # Display snapper list if available
    if command -v snapper &> /dev/null; then
        echo "[+] Available snapper snapshots:"
        snapper list
        echo "--------------------------------------------------"
    fi
    
    read -p "[?] Enter target snapshot number: (q for quit) " SNAP_NUM

    if [ -z "$SNAP_NUM" ] || [[ "$SNAP_NUM" =~ ^[Qq]$ ]]; then
        echo "[+] Rollback aborted by user."
        exit 0
    fi
    
    if [ "$SNAP_NUM" -eq 0 ] 2>/dev/null; then
        echo "[-] Error: Snapshot 0 represents the active state and cannot be used for a rollback."
        exit 1
    fi
fi

# 4. Mount top-level Btrfs root
mkdir -p /tmp/btrfs-root
if ! mountpoint -q /tmp/btrfs-root; then
    mount -o subvolid=5 "/dev/disk/by-uuid/$UUID" /tmp/btrfs-root
fi

# 5. Verify snapshot exists
SNAP_PATH="/tmp/btrfs-root/@.snapshots/$SNAP_NUM/snapshot"
if [ ! -d "$SNAP_PATH" ]; then
    echo "[-] Error: Snapshot $SNAP_NUM does not exist at $SNAP_PATH"
    umount /tmp/btrfs-root 2>/dev/null
    exit 1
fi

# 6. Perform subvolume swap
cd /tmp/btrfs-root || exit 1

if [ -d "@_broken" ]; then
    TIMESTAMP=$(date +%s)
    echo "[+] Existing @_broken found. Archiving to @_broken_$TIMESTAMP..."
    mv @_broken "@_broken_$TIMESTAMP"
fi

echo "[+] Replacing broken root subvolume with snapshot $SNAP_NUM..."
mv @ @_broken
btrfs subvolume snapshot "@.snapshots/$SNAP_NUM/snapshot" @

# --- 6b. UKI GENERATION IN TARGET SUBVOLUME (@) [FDE only] ---
# Check if LUKS encryption (FDE) is active
IS_ENCRYPTED=false
if [ -s /etc/crypttab ] && grep -v '^#' /etc/crypttab | grep -q '.'; then
    IS_ENCRYPTED=true
fi

if [ "$IS_ENCRYPTED" = true ]; then
    echo "[+] Encrypted system detected. Preparing chroot & regenerating UKI..."

    TARGET_ROOT="/tmp/btrfs-root/@"
    mount --bind "$TARGET_ROOT" "$TARGET_ROOT"

    # Provide virtual filesystems for chroot
    for i in /dev /dev/pts /proc /sys /sys/firmware/efi/efivars /run; do
        if [ -d "$i" ]; then
            mkdir -p "$TARGET_ROOT$i"
            mount -B "$i" "$TARGET_ROOT$i"
        fi
    done

    if [ -d /home ]; then
        mount -B /home "$TARGET_ROOT/home"
    fi

    ESP_DEV=$(findmnt -n -o SOURCE /boot/efi 2>/dev/null || findmnt -n -o SOURCE /boot 2>/dev/null || true)
    ESP_DEV=$(echo "$ESP_DEV" | sed 's/\[.*\]//')

    if [ -n "$ESP_DEV" ] && [ -e "$ESP_DEV" ]; then
        mkdir -p "$TARGET_ROOT/boot/efi"
        mount "$ESP_DEV" "$TARGET_ROOT/boot/efi"
    fi

    HOST_ROOT_UUID=$(findmnt -n -o UUID / 2>/dev/null || blkid -s UUID -o value $(findmnt -n -o SOURCE /) 2>/dev/null || true)

    BASE_SCRIPT=""
    if [ -f "$TARGET_ROOT/usr/share/tuxedo-btrfs/create_boot_uki_base.sh" ]; then
        BASE_SCRIPT="/usr/share/tuxedo-btrfs/create_boot_uki_base.sh"
    elif [ -f "$TARGET_ROOT/usr/share/tuxedo-base-files/create_boot_uki_base.sh" ]; then
        BASE_SCRIPT="/usr/share/tuxedo-base-files/create_boot_uki_base.sh"
    elif [ -f "$TARGET_ROOT/usr/sbin/create_boot_uki_base.sh" ]; then
        BASE_SCRIPT="/usr/sbin/create_boot_uki_base.sh"
    fi

    if [ -n "$BASE_SCRIPT" ]; then
        if chroot "$TARGET_ROOT" env TARGET_SUBVOL="@" HOST_ROOT_UUID="$HOST_ROOT_UUID" "$BASE_SCRIPT"; then
            echo "[+] UKI generation completed successfully for subvol=@!"
        else
            echo "[!] WARNING: UKI generation failed inside chroot!"
        fi
    else
        echo "[-] Error: UKI generator script not found in target subvolume."
    fi

    # Cleanup chroot mounts
    umount -R "$TARGET_ROOT/boot/efi" 2>/dev/null || umount -l "$TARGET_ROOT/boot/efi" 2>/dev/null || true
    umount "$TARGET_ROOT/home" 2>/dev/null || true
    for i in /run /sys/firmware/efi/efivars /sys /proc /dev/pts /dev; do
        umount -R "$TARGET_ROOT$i" 2>/dev/null || umount -l "$TARGET_ROOT$i" 2>/dev/null || true
    done

    umount "$TARGET_ROOT"
else
    echo "[+] Unencrypted system detected. Skipping UKI regeneration."
fi

# --- 7. FIX SDDM PERMISSIONS ---
echo "[+] Fixing display manager permissions..."
if [ -d "/tmp/btrfs-root/@/var/lib/sddm" ]; then
    chown -R sddm:sddm /tmp/btrfs-root/@/var/lib/sddm
    chmod -R 750 /tmp/btrfs-root/@/var/lib/sddm
else
    echo "[!] Warning: /var/lib/sddm not found in target subvolume, skipping permission fix."
fi

# --- 8. CLEANUP AND REBOOT PROMPT ---
cd /root || cd /

# Safely unmount any remaining mounts under /tmp/btrfs-root
umount -R /tmp/btrfs-root 2>/dev/null || umount -l /tmp/btrfs-root 2>/dev/null || true

echo "[+] Rollback complete successfully!"

cp -af /usr/share/tuxedo-btrfs/fix-menu.desktop /home/*/.config/autostart/
chmod 777 /home/*/.config/autostart/fix-menu.desktop

read -p "[?] Do you want to reboot now? (y/n): " REBOOT_CHOICE
if [[ "$REBOOT_CHOICE" =~ ^[Yy]$ ]]; then
    reboot
else
    echo "[+] Please remember to reboot before using the system normally."
fi
