#!/bin/bash

# Fix GRUB Boot for LVM System - External M.2
# This script fixes GRUB configuration to boot from the external LVM system instead of internal drive

set -e

LOG_FILE="/tmp/grub_fix_$(date +%Y%m%d_%H%M%S).log"
EXTERNAL_ROOT="/tmp/external-root"
EXTERNAL_BOOT="/tmp/external-boot"

echo "🔧 GRUB LVM Boot Fix - Starting at $(date)" | tee "$LOG_FILE"
echo "==========================================" | tee -a "$LOG_FILE"

# Function to cleanup mounts on exit
cleanup() {
    echo "🧹 Cleaning up mounts..." | tee -a "$LOG_FILE"
    sudo umount "$EXTERNAL_ROOT/proc" 2>/dev/null || true
    sudo umount "$EXTERNAL_ROOT/sys" 2>/dev/null || true
    sudo umount "$EXTERNAL_ROOT/dev/pts" 2>/dev/null || true
    sudo umount "$EXTERNAL_ROOT/dev" 2>/dev/null || true
    sudo umount "$EXTERNAL_ROOT/boot" 2>/dev/null || true
    sudo umount "$EXTERNAL_ROOT" 2>/dev/null || true
    sudo umount "$EXTERNAL_BOOT" 2>/dev/null || true
    sudo rmdir "$EXTERNAL_ROOT" "$EXTERNAL_BOOT" 2>/dev/null || true
}
trap cleanup EXIT

# Check if LVM is active
echo "📋 Checking LVM status..." | tee -a "$LOG_FILE"
if ! sudo lvs system-vg/root &>/dev/null; then
    echo "❌ LVM system-vg/root not found. Please ensure the external M.2 is connected." | tee -a "$LOG_FILE"
    exit 1
fi

# Create mount points
sudo mkdir -p "$EXTERNAL_ROOT" "$EXTERNAL_BOOT"

# Mount the external LVM system
echo "💾 Mounting external LVM system..." | tee -a "$LOG_FILE"
sudo mount /dev/system-vg/root "$EXTERNAL_ROOT" | tee -a "$LOG_FILE"
sudo mount /dev/system-vg/boot "$EXTERNAL_ROOT/boot" | tee -a "$LOG_FILE"

# Bind mount system directories for chroot
echo "🔗 Setting up chroot environment..." | tee -a "$LOG_FILE"
sudo mount --bind /proc "$EXTERNAL_ROOT/proc"
sudo mount --bind /sys "$EXTERNAL_ROOT/sys"  
sudo mount --bind /dev "$EXTERNAL_ROOT/dev"
sudo mount --bind /dev/pts "$EXTERNAL_ROOT/dev/pts"

# Show current GRUB configuration issue
echo "🔍 Current GRUB issue:" | tee -a "$LOG_FILE"
if sudo mount /dev/system-vg/boot "$EXTERNAL_BOOT" 2>/dev/null; then
    WRONG_ENTRIES=$(sudo grep -c "nvme0n1p1" "$EXTERNAL_BOOT/grub/grub.cfg" 2>/dev/null || echo "0")
    echo "   - Found $WRONG_ENTRIES entries pointing to internal drive (nvme0n1p1)" | tee -a "$LOG_FILE"
    sudo umount "$EXTERNAL_BOOT"
fi

# Get the correct UUID for the LVM root
ROOT_UUID=$(sudo blkid /dev/system-vg/root | grep -o 'UUID="[^"]*"' | cut -d'"' -f2)
echo "   - External LVM root UUID: $ROOT_UUID" | tee -a "$LOG_FILE"

# Update GRUB configuration from within the external system
echo "🔄 Regenerating GRUB configuration..." | tee -a "$LOG_FILE"

# First, ensure EFI partition is mounted in the chroot
EFI_PARTITION=$(lsblk -rno MOUNTPOINT,FSTYPE | grep vfat | head -1 | cut -d' ' -f1)
if [ -n "$EFI_PARTITION" ]; then
    echo "📁 Mounting EFI partition: $EFI_PARTITION" | tee -a "$LOG_FILE"
    sudo mkdir -p "$EXTERNAL_ROOT/boot/efi"
    sudo mount "$EFI_PARTITION" "$EXTERNAL_ROOT/boot/efi" || true
fi

# Update fstab to reflect LVM UUIDs
echo "📝 Updating /etc/fstab with correct UUIDs..." | tee -a "$LOG_FILE"
sudo chroot "$EXTERNAL_ROOT" bash -c "
    # Get UUIDs for all LVM volumes
    ROOT_UUID=\$(blkid /dev/system-vg/root | grep -o 'UUID=\"[^\"]*\"' | cut -d'\"' -f2)
    BOOT_UUID=\$(blkid /dev/system-vg/boot | grep -o 'UUID=\"[^\"]*\"' | cut -d'\"' -f2)
    HOME_UUID=\$(blkid /dev/system-vg/home | grep -o 'UUID=\"[^\"]*\"' | cut -d'\"' -f2)
    SWAP_UUID=\$(blkid /dev/system-vg/swap | grep -o 'UUID=\"[^\"]*\"' | cut -d'\"' -f2)

    # Backup original fstab
    cp /etc/fstab /etc/fstab.backup

    # Create new fstab with LVM UUIDs
    cat > /etc/fstab << EOF
# LVM-based fstab generated by fix_grub_lvm_boot.sh
UUID=\${ROOT_UUID} / ext4 errors=remount-ro 0 1
UUID=\${BOOT_UUID} /boot ext4 defaults 0 2
UUID=\${HOME_UUID} /home ext4 defaults 0 2
UUID=\${SWAP_UUID} none swap sw 0 0
# EFI System Partition
/dev/disk/by-label/SYSTEM-EFI /boot/efi vfat umask=0077 0 1
EOF

    echo '✅ Updated /etc/fstab with LVM UUIDs'
"

# Install GRUB to EFI
echo "🛠️  Installing GRUB to EFI..." | tee -a "$LOG_FILE"
sudo chroot "$EXTERNAL_ROOT" bash -c "
    grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=ubuntu-lvm --recheck 2>&1
    echo '✅ GRUB installed to EFI'
"

# Update GRUB configuration
echo "📝 Updating GRUB configuration..." | tee -a "$LOG_FILE"
sudo chroot "$EXTERNAL_ROOT" bash -c "
    update-grub 2>&1
    echo '✅ GRUB configuration updated'
"

# Verify the fix
echo "✅ Verifying GRUB fix..." | tee -a "$LOG_FILE"
sudo mount /dev/system-vg/boot "$EXTERNAL_BOOT"
NEW_ENTRIES=$(sudo grep -c "system-vg-root" "$EXTERNAL_BOOT/grub/grub.cfg" 2>/dev/null || echo "0")
OLD_ENTRIES=$(sudo grep -c "nvme0n1p1" "$EXTERNAL_BOOT/grub/grub.cfg" 2>/dev/null || echo "0")

echo "   - LVM entries found: $NEW_ENTRIES" | tee -a "$LOG_FILE"
echo "   - Internal drive entries: $OLD_ENTRIES" | tee -a "$LOG_FILE"

if [ "$NEW_ENTRIES" -gt 0 ]; then
    echo "✅ SUCCESS: GRUB now configured to boot from external LVM system!" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "🚀 Next steps:" | tee -a "$LOG_FILE"
    echo "   1. Reboot your system" | tee -a "$LOG_FILE"
    echo "   2. In GRUB menu, select 'ubuntu-lvm' or Ubuntu entries" | tee -a "$LOG_FILE"
    echo "   3. System should now boot from external M.2 LVM" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "📋 Log saved to: $LOG_FILE"
else
    echo "⚠️  Warning: No LVM entries found in new GRUB config. Manual check required." | tee -a "$LOG_FILE"
    echo "   Check $LOG_FILE for details."
fi

echo "" | tee -a "$LOG_FILE"
echo "🔧 GRUB LVM Boot Fix completed at $(date)" | tee -a "$LOG_FILE"