#!/bin/bash

# GRUB Boot Repair Script for LVM Migration
# Fixes GRUB bootloader issues after LVM migration to external M.2 drive

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
VG_NAME="system-vg"
EXTERNAL_DRIVE="/dev/sda"  # Adjust if needed

log() {
    local message="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "${BLUE}$message${NC}"
}

error() {
    local message="[ERROR] $1"
    echo -e "${RED}$message${NC}" >&2
    exit 1
}

warning() {
    local message="[WARNING] $1"
    echo -e "${YELLOW}$message${NC}"
}

success() {
    local message="[SUCCESS] $1"
    echo -e "${GREEN}$message${NC}"
}

confirm_action() {
    echo -e "${YELLOW}$1${NC}"
    read -p "Do you want to continue? [y/N] " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        error "Operation aborted by user"
    fi
}

check_prerequisites() {
    log "Checking prerequisites for GRUB repair..."
    
    # Check if running as root
    if [ "$EUID" -ne 0 ]; then
        error "This script must be run as root. Use: sudo $0"
    fi
    
    # Check if running from live system
    local root_device=$(df / | tail -1 | awk '{print $1}')
    if [[ "$root_device" != *"loop"* ]] && [[ "$root_device" != *"overlay"* ]] && [[ "$root_device" != *"tmpfs"* ]]; then
        warning "This should be run from a live USB system for safety"
        confirm_action "Continue anyway? (Not recommended if booting from the problematic system)"
    fi
    
    # Check required tools
    for tool in grub-install update-grub mount umount lvm; do
        if ! command -v $tool >/dev/null 2>&1; then
            error "$tool not found. Please install required packages."
        fi
    done
    
    success "Prerequisites check passed"
}

detect_external_drive() {
    log "Detecting external M.2 drive with LVM..."
    
    # Check if LVM volume group exists
    if ! vgs "$VG_NAME" >/dev/null 2>&1; then
        # Try to activate volume groups
        vgchange -ay || warning "Could not activate volume groups"
        
        if ! vgs "$VG_NAME" >/dev/null 2>&1; then
            error "Volume group '$VG_NAME' not found. Is the external drive connected?"
        fi
    fi
    
    # Display volume group information
    log "Found volume group '$VG_NAME':"
    vgs "$VG_NAME"
    
    # Check logical volumes
    local required_lvs=("root" "home" "boot")
    for lv in "${required_lvs[@]}"; do
        if ! lvs "$VG_NAME/$lv" >/dev/null 2>&1; then
            error "Logical volume '$VG_NAME/$lv' not found"
        fi
    done
    
    success "External LVM drive detected successfully"
}

mount_external_system() {
    log "Mounting external LVM system..."
    
    local mount_base="/mnt/ext"
    mkdir -p "$mount_base"
    
    # Mount in correct order
    log "Mounting root filesystem..."
    mount "/dev/$VG_NAME/root" "$mount_base" || error "Failed to mount root"
    
    log "Mounting boot filesystem..."
    mkdir -p "$mount_base/boot"
    mount "/dev/$VG_NAME/boot" "$mount_base/boot" || error "Failed to mount boot"
    
    log "Mounting EFI partition..."
    mkdir -p "$mount_base/boot/efi"
    # Find EFI partition (usually first partition on external drive)
    local efi_partition="${EXTERNAL_DRIVE}1"
    if [ -b "$efi_partition" ]; then
        mount "$efi_partition" "$mount_base/boot/efi" || error "Failed to mount EFI partition"
    else
        error "EFI partition $efi_partition not found"
    fi
    
    log "Mounting home filesystem..."
    mkdir -p "$mount_base/home"
    mount "/dev/$VG_NAME/home" "$mount_base/home" || error "Failed to mount home"
    
    success "External system mounted at $mount_base"
    echo "$mount_base"
}

repair_grub() {
    local mount_base="$1"
    log "Repairing GRUB bootloader..."
    
    # Bind mount necessary filesystems for chroot
    log "Setting up chroot environment..."
    mount --bind /dev "$mount_base/dev" || error "Failed to bind /dev"
    mount --bind /proc "$mount_base/proc" || error "Failed to bind /proc"
    mount --bind /sys "$mount_base/sys" || error "Failed to bind /sys"
    mount --bind /run "$mount_base/run" || error "Failed to bind /run"
    
    # Copy DNS resolution
    cp /etc/resolv.conf "$mount_base/etc/resolv.conf" 2>/dev/null || warning "Could not copy DNS settings"
    
    log "Updating initramfs to include LVM support..."
    chroot "$mount_base" /bin/bash -c "
        # Ensure LVM is in initramfs
        echo 'GRUB_ENABLE_CRYPTODISK=y' >> /etc/default/grub 2>/dev/null || true
        
        # Update initramfs with LVM support
        update-initramfs -u -k all
        
        echo 'Initramfs updated successfully'
    " || warning "Initramfs update had issues"
    
    log "Reinstalling GRUB bootloader..."
    chroot "$mount_base" /bin/bash -c "
        # Install GRUB to the external drive
        grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian --recheck '$EXTERNAL_DRIVE'
        
        echo 'GRUB installed successfully'
    " || error "GRUB installation failed"
    
    log "Updating GRUB configuration..."
    chroot "$mount_base" /bin/bash -c "
        # Update GRUB configuration
        update-grub
        
        echo 'GRUB configuration updated successfully'
    " || error "GRUB configuration update failed"
    
    success "GRUB repair completed"
}

check_grub_files() {
    local mount_base="$1"
    log "Checking GRUB installation..."
    
    # Check EFI bootloader
    if [ -f "$mount_base/boot/efi/EFI/debian/grubx64.efi" ]; then
        success "GRUB EFI bootloader found"
    else
        warning "GRUB EFI bootloader missing"
    fi
    
    # Check GRUB configuration
    if [ -f "$mount_base/boot/grub/grub.cfg" ]; then
        success "GRUB configuration found"
        
        # Check if configuration mentions LVM
        if grep -q "$VG_NAME" "$mount_base/boot/grub/grub.cfg"; then
            success "GRUB configuration includes LVM volumes"
        else
            warning "GRUB configuration may not include LVM setup"
        fi
    else
        error "GRUB configuration missing"
    fi
}

cleanup_mounts() {
    local mount_base="$1"
    log "Cleaning up mount points..."
    
    # Unmount in reverse order
    umount "$mount_base/run" 2>/dev/null || true
    umount "$mount_base/sys" 2>/dev/null || true
    umount "$mount_base/proc" 2>/dev/null || true
    umount "$mount_base/dev" 2>/dev/null || true
    umount "$mount_base/home" 2>/dev/null || true
    umount "$mount_base/boot/efi" 2>/dev/null || true
    umount "$mount_base/boot" 2>/dev/null || true
    umount "$mount_base" 2>/dev/null || true
    
    # Remove mount directory
    rmdir "$mount_base" 2>/dev/null || true
    
    success "Mount points cleaned up"
}

show_next_steps() {
    echo
    echo -e "${GREEN}=== GRUB Repair Completed ===${NC}"
    echo
    echo -e "${BLUE}Next steps:${NC}"
    echo "1. Reboot your system"
    echo "2. Enter BIOS/UEFI settings"
    echo "3. Ensure boot order prioritizes the external M.2 SSD"
    echo "4. Look for 'debian' entry in the boot menu"
    echo "5. Boot from the external drive"
    echo
    echo -e "${YELLOW}If boot still fails:${NC}"
    echo "• Check BIOS/UEFI settings for Secure Boot (disable if necessary)"
    echo "• Verify UEFI mode is enabled (not Legacy/CSM)"
    echo "• Try different USB ports if using external enclosure"
    echo "• Run this script again from live USB if needed"
    echo
    echo -e "${GREEN}The system should now boot properly from the external M.2!${NC}"
}

main() {
    echo -e "${GREEN}=== GRUB Boot Repair for LVM Migration ===${NC}"
    echo "This script repairs GRUB bootloader issues after LVM migration"
    echo
    
    check_prerequisites
    detect_external_drive
    
    echo
    echo -e "${YELLOW}This will repair GRUB on the external M.2 drive${NC}"
    echo "Drive: $EXTERNAL_DRIVE"
    echo "Volume Group: $VG_NAME"
    echo
    confirm_action "Proceed with GRUB repair?"
    
    local mount_base=$(mount_external_system)
    
    # Set trap to ensure cleanup
    trap "cleanup_mounts '$mount_base'" EXIT
    
    repair_grub "$mount_base"
    check_grub_files "$mount_base"
    cleanup_mounts "$mount_base"
    
    # Clear trap since we cleaned up successfully
    trap - EXIT
    
    show_next_steps
}

main "$@"