#!/bin/bash

# Improved LVM Migration Script
# Fixes the boot issues from the previous failed LVM migration
# Properly handles LUKS + LVM combination with robust boot configuration

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 variables
INTERNAL_DRIVE=""
EXTERNAL_DRIVE=""
VG_NAME="migration-vg"  # Changed to avoid conflict with existing system-vg
ROOT_LV="root"
HOME_LV="home"
SWAP_LV="swap"
BOOT_LV="boot"

# Work directory
WORK_DIR="/mnt/lvm_migration"

# Detected partitions and info
declare -A INTERNAL_PARTITIONS
declare -A PARTITION_FILESYSTEMS
declare -A PARTITION_SIZES

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

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

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

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

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
}

detect_drives() {
    log "Detecting available drives..."
    
    local all_drives=($(lsblk -dpno NAME,TYPE | grep "disk" | awk '{print $1}'))
    local drives=()
    
    # Filter out the USB stick we're running from
    for drive in "${all_drives[@]}"; do
        if mount | grep -q "$drive" && mount | grep -q "/lib/live\|overlay\|/media.*live"; then
            log "Excluding live USB drive: $drive"
            continue
        fi
        drives+=("$drive")
    done
    
    if [ ${#drives[@]} -lt 2 ]; then
        error "Need at least 2 drives for migration. Found only ${#drives[@]} suitable drives"
    fi
    
    echo "Available drives:"
    for i in "${!drives[@]}"; do
        local drive="${drives[$i]}"
        local info=$(lsblk -dpno SIZE,MODEL "$drive" | xargs)
        echo "$((i+1)). $drive - $info"
        lsblk "$drive" | tail -n +2 | sed 's/^/   /'
        echo
    done
    
    # Auto-detect with user confirmation
    local suggested_internal=""
    local suggested_external=""
    
    # Prefer NVMe for internal, USB for external
    for drive in "${drives[@]}"; do
        if [[ "$drive" == *"nvme"* ]]; then
            suggested_internal="$drive"
            break
        fi
    done
    
    for drive in "${drives[@]}"; do
        if [ "$drive" != "$suggested_internal" ]; then
            if udevadm info --query=property --name="$drive" 2>/dev/null | grep -q "ID_BUS=usb"; then
                suggested_external="$drive"
                break
            fi
        fi
    done
    
    if [ -n "$suggested_internal" ] && [ -n "$suggested_external" ]; then
        echo "Suggested configuration:"
        echo "  Internal (source): $suggested_internal"
        echo "  External (LVM target): $suggested_external"
        read -p "Use this configuration? [Y/n]: " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            INTERNAL_DRIVE="$suggested_internal"
            EXTERNAL_DRIVE="$suggested_external"
        fi
    fi
    
    # Manual selection if needed
    if [ -z "$INTERNAL_DRIVE" ]; then
        read -p "Select INTERNAL drive number: " choice
        INTERNAL_DRIVE="${drives[$((choice-1))]}"
    fi
    
    if [ -z "$EXTERNAL_DRIVE" ]; then
        read -p "Select EXTERNAL drive number: " choice
        EXTERNAL_DRIVE="${drives[$((choice-1))]}"
    fi
    
    # Safety checks
    if [ "$INTERNAL_DRIVE" = "$EXTERNAL_DRIVE" ]; then
        error "Internal and external drives cannot be the same!"
    fi
    
    local external_size_bytes=$(lsblk -bno SIZE "$EXTERNAL_DRIVE" | head -1 | tr -d ' ')
    local internal_size_bytes=$(lsblk -bno SIZE "$INTERNAL_DRIVE" | head -1 | tr -d ' ')
    
    if [ -n "$external_size_bytes" ] && [ -n "$internal_size_bytes" ] && [ "$external_size_bytes" -lt "$internal_size_bytes" ]; then
        error "External drive is smaller than internal drive"
    fi
    
    echo
    echo "Final configuration:"
    echo "  Internal (source): $INTERNAL_DRIVE"
    echo "  External (LVM target): $EXTERNAL_DRIVE"
    
    # Final safety confirmation
    echo
    echo -e "${RED}⚠️  FINAL SAFETY CHECK ⚠️${NC}"
    echo "This will COMPLETELY WIPE: $EXTERNAL_DRIVE"
    echo "Current partitions that will be DESTROYED:"
    lsblk "$EXTERNAL_DRIVE"
    echo
    read -p "Type 'MIGRATE' to confirm LVM migration: " confirmation
    if [ "$confirmation" != "MIGRATE" ]; then
        error "Migration cancelled by user"
    fi
    
    success "Drive selection completed"
}

analyze_source_system() {
    log "Analyzing source system..."
    
    # Get partitions using a more reliable method
    local partitions=($(lsblk -lpno NAME "$INTERNAL_DRIVE" | tail -n +2))
    
    echo "Source drive partitions:"
    for part in "${partitions[@]}"; do
        # Check if partition actually exists before querying
        if [ ! -b "$part" ]; then
            continue
        fi
        
        local size=$(lsblk -no SIZE "$part" 2>/dev/null || echo "unknown")
        local fstype=$(lsblk -no FSTYPE "$part" 2>/dev/null || echo "")
        local label=$(lsblk -no LABEL "$part" 2>/dev/null || echo "")
        local mountpoint=$(lsblk -no MOUNTPOINT "$part" 2>/dev/null || echo "")
        
        echo "  $part: $size, $fstype, ${label:-'no label'}"
        
        PARTITION_FILESYSTEMS["$part"]="$fstype"
        PARTITION_SIZES["$part"]="$size"
        
        # Identify partitions
        if [[ "$fstype" == "vfat" ]] && [[ "$part" == *"1" ]] || [[ "$part" == *"2" ]]; then
            INTERNAL_PARTITIONS["efi"]="$part"
        elif [[ "$mountpoint" == "/" ]] || [[ "$label" == "root"* ]] || [[ "$fstype" == "ext4" && "$part" == *"1" ]]; then
            INTERNAL_PARTITIONS["root"]="$part"
        elif [[ "$mountpoint" == "/home" ]] || [[ "$label" == "home"* ]]; then
            INTERNAL_PARTITIONS["home"]="$part"
        elif [[ "$fstype" == "crypto_LUKS" ]]; then
            # Encrypted partition - likely home
            INTERNAL_PARTITIONS["encrypted_home"]="$part"
        fi
    done
    
    success "Source system analysis completed"
}

calculate_partition_sizes() {
    log "Calculating required partition sizes..."
    
    # Get actual source partition sizes
    local source_root_gb=0
    local source_home_gb=0
    
    # Check actual partition sizes from the source
    for part_name in "${!INTERNAL_PARTITIONS[@]}"; do
        local part_device="${INTERNAL_PARTITIONS[$part_name]}"
        if [ -b "$part_device" ]; then
            local size_bytes=$(lsblk -bno SIZE "$part_device" 2>/dev/null | head -1 | tr -d ' ')
            if [ -n "$size_bytes" ]; then
                local size_gb=$((size_bytes / 1024 / 1024 / 1024))
                
                case "$part_name" in
                    "root")
                        source_root_gb=$size_gb
                        log "Source root: ${size_gb}GB"
                        ;;
                    "home"|"encrypted_home")
                        source_home_gb=$size_gb
                        log "Source home: ${size_gb}GB"
                        ;;
                esac
            fi
        fi
    done
    
    # Get target drive total space
    local total_space_bytes=$(lsblk -bno SIZE "$EXTERNAL_DRIVE" | head -1 | tr -d ' ')
    local total_space_gb=$((total_space_bytes / 1024 / 1024 / 1024))
    
    log "Target drive total space: ${total_space_gb}GB"
    
    # Fixed sizes for system partitions
    local swap_size="8G"
    local boot_size="2G"
    
    # Calculate available space for data partitions (leave 2GB for overhead/EFI)
    local available_for_data=$((total_space_gb - 8 - 2 - 2))  # 464GB available
    
    # For same-size drives, distribute space proportionally to source
    local total_source_data=$((source_root_gb + source_home_gb))
    
    if [ "$total_source_data" -gt "$available_for_data" ]; then
        # Source is larger than target, scale down proportionally
        local scale_factor_percent=$((available_for_data * 100 / total_source_data))
        local root_size="$((source_root_gb * scale_factor_percent / 100))G"
        local home_size="$((source_home_gb * scale_factor_percent / 100))G"
        
        warning "Scaling down partitions to fit target drive:"
        warning "  Scale factor: ${scale_factor_percent}%"
    else
        # Target has enough space, use source sizes with small buffers
        local root_size="$((source_root_gb + 5))G"  # 5GB buffer for root
        local remaining_space=$((available_for_data - source_root_gb - 5))
        local home_size="${remaining_space}G"  # Use all remaining space for home
    fi
    
    # Export calculated sizes
    CALCULATED_ROOT_SIZE="$root_size"
    CALCULATED_HOME_SIZE="$home_size"
    CALCULATED_SWAP_SIZE="$swap_size"
    CALCULATED_BOOT_SIZE="$boot_size"
    
    log "Final calculated sizes:"
    log "  Root: $CALCULATED_ROOT_SIZE"
    log "  Home: $CALCULATED_HOME_SIZE" 
    log "  Swap: $CALCULATED_SWAP_SIZE"
    log "  Boot: $CALCULATED_BOOT_SIZE"
    
    # Verify total fits
    local total_allocated=$((${CALCULATED_ROOT_SIZE%G} + ${CALCULATED_HOME_SIZE%G} + ${CALCULATED_SWAP_SIZE%G} + ${CALCULATED_BOOT_SIZE%G}))
    log "Total allocated: ${total_allocated}GB of ${total_space_gb}GB"
    
    success "Partition sizes calculated"
}

check_prerequisites() {
    log "Checking prerequisites and installing required tools..."
    
    # Check if running from live system
    if ! df / | grep -q "loop\|overlay\|tmpfs"; then
        warning "Not running from live system - this may cause issues"
        confirm_action "Continue anyway?"
    fi
    
    # Install/update required packages
    log "Installing required packages..."
    apt update >/dev/null 2>&1
    apt install -y lvm2 cryptsetup rsync parted pv grub-efi-amd64 grub-common \
                   e2fsprogs dosfstools bc util-linux initramfs-tools \
                   efibootmgr os-prober >/dev/null 2>&1
    
    # Ensure LVM2 is properly loaded
    modprobe dm-mod
    modprobe dm-crypt
    vgchange -ay 2>/dev/null || true
    
    success "Prerequisites installed"
}

create_lvm_layout() {
    log "Creating LVM layout on external drive..."
    
    # Use the calculated sizes
    local root_size="$CALCULATED_ROOT_SIZE"
    local home_size="$CALCULATED_HOME_SIZE"
    local swap_size="$CALCULATED_SWAP_SIZE"
    local boot_size="$CALCULATED_BOOT_SIZE"
    
    log "Using calculated sizes: Root=$root_size, Home=$home_size, Swap=$swap_size, Boot=$boot_size"
    
    # Properly unmount and deactivate any existing LVM on the target drive
    log "Cleaning up existing LVM on target drive..."
    
    # Check if target drive partitions are currently mounted or in use
    local partitions_in_use=false
    for part in "${EXTERNAL_DRIVE}"*; do
        if [ -b "$part" ]; then
            if mount | grep -q "$part"; then
                log "Unmounting $part..."
                umount "$part" 2>/dev/null || {
                    warning "Could not unmount $part - it may be in use"
                    partitions_in_use=true
                }
            fi
            
            # Check if this partition has a VG on it
            local vg_on_part=$(pvs --noheadings -o vg_name "$part" 2>/dev/null | tr -d ' ')
            if [ -n "$vg_on_part" ]; then
                log "Deactivating VG '$vg_on_part' on $part"
                vgchange -an "$vg_on_part" 2>/dev/null || true
                vgremove -f "$vg_on_part" 2>/dev/null || true
            fi
        fi
    done
    
    if [ "$partitions_in_use" = true ]; then
        warning "Some partitions are in use. Continuing anyway..."
    fi
    
    # Remove any existing LVM structures with force (only on target drive)
    log "Removing existing PV structures on target drive..."
    for part in "${EXTERNAL_DRIVE}"*; do
        if [ -b "$part" ]; then
            pvremove -ff "$part" 2>/dev/null || true
        fi
    done
    
    # Wipe filesystem signatures and partition table
    log "Wiping drive signatures..."
    wipefs -af "$EXTERNAL_DRIVE" 2>/dev/null || true
    dd if=/dev/zero of="$EXTERNAL_DRIVE" bs=1M count=100 2>/dev/null || true
    
    # Force kernel to re-read partition table
    partprobe "$EXTERNAL_DRIVE" 2>/dev/null || true
    sleep 2
    
    # Create new partition table
    log "Creating new partition table..."
    parted -s "$EXTERNAL_DRIVE" mklabel gpt
    
    # Create EFI partition (512MB)
    log "Creating EFI partition..."
    parted -s "$EXTERNAL_DRIVE" mkpart primary fat32 1MiB 513MiB
    parted -s "$EXTERNAL_DRIVE" set 1 boot on
    parted -s "$EXTERNAL_DRIVE" set 1 esp on
    
    # Create LVM partition (rest of disk)
    log "Creating LVM partition..."
    parted -s "$EXTERNAL_DRIVE" mkpart primary 513MiB 100%
    parted -s "$EXTERNAL_DRIVE" set 2 lvm on
    
    # Force kernel to re-read the new partition table
    log "Refreshing partition table..."
    partprobe "$EXTERNAL_DRIVE" || {
        warning "partprobe failed, trying alternative methods..."
        echo 1 > /sys/block/$(basename "$EXTERNAL_DRIVE")/device/rescan 2>/dev/null || true
        hdparm -z "$EXTERNAL_DRIVE" 2>/dev/null || true
    }
    
    # Wait for partitions to appear
    local retry_count=0
    while [ ! -b "${EXTERNAL_DRIVE}1" ] || [ ! -b "${EXTERNAL_DRIVE}2" ]; do
        sleep 2
        retry_count=$((retry_count + 1))
        if [ $retry_count -gt 10 ]; then
            error "Partitions not appearing after 20 seconds. Please reboot and try again."
        fi
        log "Waiting for partitions to appear... ($retry_count/10)"
    done
    
    log "Partitions created successfully"
    
    # Create filesystems
    mkfs.fat -F32 "${EXTERNAL_DRIVE}1" || error "Failed to create EFI filesystem"
    
    # Setup LVM with force flag to handle existing signatures
    log "Creating physical volume..."
    pvcreate -ff "${EXTERNAL_DRIVE}2" || error "Failed to create physical volume"
    
    # Check if VG name already exists and handle it
    if vgs "$VG_NAME" >/dev/null 2>&1; then
        log "Volume group $VG_NAME already exists, removing it first..."
        vgremove -f "$VG_NAME" 2>/dev/null || true
        # Wait a moment for cleanup
        sleep 2
    fi
    
    log "Creating volume group..."
    vgcreate "$VG_NAME" "${EXTERNAL_DRIVE}2" || error "Failed to create volume group"
    
    # Verify VG creation
    if ! vgs "$VG_NAME" >/dev/null 2>&1; then
        error "Volume group $VG_NAME was not created successfully"
    fi
    
    # Create logical volumes
    lvcreate -L "$root_size" -n "$ROOT_LV" "$VG_NAME" || error "Failed to create root LV"
    lvcreate -L "$home_size" -n "$HOME_LV" "$VG_NAME" || error "Failed to create home LV"
    lvcreate -L "$swap_size" -n "$SWAP_LV" "$VG_NAME" || error "Failed to create swap LV"
    lvcreate -L "$boot_size" -n "$BOOT_LV" "$VG_NAME" || error "Failed to create boot LV"
    
    # Create filesystems on LVM volumes
    mkfs.ext4 -L "root" "/dev/$VG_NAME/$ROOT_LV" || error "Failed to create root filesystem"
    mkfs.ext4 -L "home" "/dev/$VG_NAME/$HOME_LV" || error "Failed to create home filesystem"
    mkfs.ext4 -L "boot" "/dev/$VG_NAME/$BOOT_LV" || error "Failed to create boot filesystem"
    mkswap -L "swap" "/dev/$VG_NAME/$SWAP_LV" || error "Failed to create swap"
    
    success "LVM layout created successfully"
}

handle_encrypted_partitions() {
    log "Handling encrypted partitions..."
    
    for part_name in "${!INTERNAL_PARTITIONS[@]}"; do
        local part_device="${INTERNAL_PARTITIONS[$part_name]}"
        local fstype="${PARTITION_FILESYSTEMS[$part_device]}"
        
        if [[ "$fstype" == "crypto_LUKS" ]]; then
            log "Found encrypted partition: $part_device"
            local crypt_name="migration_${part_name}"
            
            echo "Please enter password for encrypted partition ($part_device):"
            if cryptsetup open "$part_device" "$crypt_name"; then
                success "Unlocked $part_device as /dev/mapper/$crypt_name"
                INTERNAL_PARTITIONS["$part_name"]="/dev/mapper/$crypt_name"
                
                # Update filesystem type
                local decrypted_fs=$(lsblk -no FSTYPE "/dev/mapper/$crypt_name")
                PARTITION_FILESYSTEMS["/dev/mapper/$crypt_name"]="$decrypted_fs"
            else
                error "Failed to unlock encrypted partition"
            fi
        fi
    done
    
    success "Encrypted partitions handled"
}

mount_filesystems() {
    log "Mounting filesystems..."
    
    mkdir -p "$WORK_DIR"/{internal_root,internal_home,external_root,external_home,external_boot}
    
    # Mount internal filesystems
    if [ -n "${INTERNAL_PARTITIONS[root]}" ]; then
        mount "${INTERNAL_PARTITIONS[root]}" "$WORK_DIR/internal_root"
    fi
    
    if [ -n "${INTERNAL_PARTITIONS[home]}" ]; then
        mount "${INTERNAL_PARTITIONS[home]}" "$WORK_DIR/internal_home"
    elif [ -n "${INTERNAL_PARTITIONS[encrypted_home]}" ]; then
        mount "${INTERNAL_PARTITIONS[encrypted_home]}" "$WORK_DIR/internal_home"
    fi
    
    # Mount external LVM filesystems
    mount "/dev/$VG_NAME/$ROOT_LV" "$WORK_DIR/external_root"
    mount "/dev/$VG_NAME/$HOME_LV" "$WORK_DIR/external_home"
    mount "/dev/$VG_NAME/$BOOT_LV" "$WORK_DIR/external_boot"
    
    # Mount EFI
    mkdir -p "$WORK_DIR/external_root/boot/efi"
    mount "${EXTERNAL_DRIVE}1" "$WORK_DIR/external_root/boot/efi"
    
    success "Filesystems mounted"
}

copy_system_data() {
    log "Copying system data..."
    
    # Copy root filesystem
    if [ -d "$WORK_DIR/internal_root" ]; then
        log "Copying root filesystem..."
        rsync -avxHAX --progress \
            --exclude=/home/* --exclude=/proc/* --exclude=/sys/* \
            --exclude=/dev/* --exclude=/run/* --exclude=/tmp/* \
            --exclude=/var/tmp/* --exclude=/mnt/* --exclude=/media/* \
            "$WORK_DIR/internal_root/" "$WORK_DIR/external_root/"
    fi
    
    # Copy home filesystem
    if [ -d "$WORK_DIR/internal_home" ]; then
        log "Copying home filesystem..."
        rsync -avxHAX --progress "$WORK_DIR/internal_home/" "$WORK_DIR/external_home/"
    elif [ -d "$WORK_DIR/internal_root/home" ]; then
        log "Copying /home from root filesystem..."
        rsync -avxHAX --progress "$WORK_DIR/internal_root/home/" "$WORK_DIR/external_home/"
    fi
    
    # Copy boot files
    if [ -d "$WORK_DIR/internal_root/boot" ]; then
        log "Copying boot files..."
        rsync -avxHAX --progress \
            --exclude=/boot/efi/* \
            "$WORK_DIR/internal_root/boot/" "$WORK_DIR/external_boot/"
    fi
    
    success "System data copied"
}

configure_lvm_system() {
    log "Configuring LVM system..."
    
    # Get UUIDs
    local root_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/$ROOT_LV")
    local home_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/$HOME_LV")
    local boot_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/$BOOT_LV")
    local efi_uuid=$(blkid -s UUID -o value "${EXTERNAL_DRIVE}1")
    local swap_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/$SWAP_LV")
    
    # Create new fstab
    cat > "$WORK_DIR/external_root/etc/fstab" << EOF
# /etc/fstab: static file system information for LVM system
UUID=$root_uuid             /              ext4    defaults   0 1
UUID=$efi_uuid              /boot/efi      vfat    defaults   0 2
UUID=$boot_uuid             /boot          ext4    defaults   0 2
UUID=$home_uuid             /home          ext4    defaults   0 2
UUID=$swap_uuid             none           swap    sw         0 0
tmpfs                       /tmp           tmpfs   defaults,noatime,mode=1777 0 0
EOF

    # Configure LVM in initramfs
    echo "$VG_NAME" >> "$WORK_DIR/external_root/etc/initramfs-tools/conf.d/lvm"
    echo "BOOT=local" >> "$WORK_DIR/external_root/etc/initramfs-tools/conf.d/resume"
    
    # Ensure LVM modules are included
    cat > "$WORK_DIR/external_root/etc/initramfs-tools/modules" << EOF
# LVM modules
dm-mod
dm-crypt
dm-snapshot
EOF

    # Update GRUB configuration for LVM
    sed -i 's/#GRUB_ENABLE_CRYPTODISK=y/GRUB_ENABLE_CRYPTODISK=y/' "$WORK_DIR/external_root/etc/default/grub"
    echo 'GRUB_PRELOAD_MODULES="lvm"' >> "$WORK_DIR/external_root/etc/default/grub"
    
    success "LVM system configured"
}

install_bootloader() {
    log "Installing bootloader with LVM support..."
    
    # Bind mount necessary filesystems
    mount --bind /dev "$WORK_DIR/external_root/dev"
    mount --bind /proc "$WORK_DIR/external_root/proc"
    mount --bind /sys "$WORK_DIR/external_root/sys"
    mount --bind /run "$WORK_DIR/external_root/run"
    
    # Install and configure bootloader in chroot
    chroot "$WORK_DIR/external_root" /bin/bash -c "
        # Ensure LVM is available
        vgscan
        vgchange -ay
        
        # Update initramfs with LVM support
        echo 'MODULES=dep' > /etc/initramfs-tools/initramfs.conf
        echo 'BOOT=local' >> /etc/initramfs-tools/initramfs.conf
        update-initramfs -u -k all
        
        # Install GRUB with LVM support
        grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian --recheck $EXTERNAL_DRIVE
        
        # Update GRUB configuration
        update-grub
        
        # Verify LVM tools are available
        which lvm && echo 'LVM tools available'
        ls -la /boot/initrd.img-* | head -1
    "
    
    # Unmount bind mounts
    umount "$WORK_DIR/external_root/dev" 2>/dev/null || true
    umount "$WORK_DIR/external_root/proc" 2>/dev/null || true
    umount "$WORK_DIR/external_root/sys" 2>/dev/null || true
    umount "$WORK_DIR/external_root/run" 2>/dev/null || true
    
    success "Bootloader installed with LVM support"
}

verify_lvm_boot() {
    log "Verifying LVM boot configuration..."
    
    # Check if initramfs contains LVM modules
    local initrd_file=$(ls "$WORK_DIR/external_boot/initrd.img-"* 2>/dev/null | head -1)
    if [ -n "$initrd_file" ]; then
        if lsinitramfs "$initrd_file" | grep -q "dm-mod\|lvm"; then
            success "Initramfs contains LVM modules"
        else
            warning "Initramfs may be missing LVM modules"
        fi
    fi
    
    # Check GRUB configuration
    if grep -q "lvm" "$WORK_DIR/external_boot/grub/grub.cfg"; then
        success "GRUB configuration includes LVM support"
    else
        warning "GRUB configuration may not have proper LVM support"
    fi
    
    # Check fstab
    if grep -q "/dev/$VG_NAME" "$WORK_DIR/external_root/etc/fstab"; then
        success "fstab configured for LVM"
    else
        warning "fstab configuration issue"
    fi
    
    success "LVM boot verification completed"
}

cleanup() {
    log "Cleaning up..."
    
    # Unmount filesystems in reverse order
    umount "$WORK_DIR/external_root/boot/efi" 2>/dev/null || true
    umount "$WORK_DIR/external_root/dev" 2>/dev/null || true
    umount "$WORK_DIR/external_root/proc" 2>/dev/null || true
    umount "$WORK_DIR/external_root/sys" 2>/dev/null || true
    umount "$WORK_DIR/external_root/run" 2>/dev/null || true
    umount "$WORK_DIR/external_root" 2>/dev/null || true
    umount "$WORK_DIR/external_home" 2>/dev/null || true
    umount "$WORK_DIR/external_boot" 2>/dev/null || true
    umount "$WORK_DIR/internal_root" 2>/dev/null || true
    umount "$WORK_DIR/internal_home" 2>/dev/null || true
    
    # Deactivate LVM volumes
    if [ -n "$VG_NAME" ]; then
        vgchange -an "$VG_NAME" 2>/dev/null || true
    fi
    
    # Close encrypted partitions
    for mapper in /dev/mapper/migration_*; do
        if [ -b "$mapper" ]; then
            cryptsetup close "$(basename "$mapper")" 2>/dev/null || true
        fi
    done
    
    # Remove work directory
    rm -rf "$WORK_DIR" 2>/dev/null || true
    
    success "Cleanup completed"
}

main() {
    echo -e "${GREEN}=== Improved LVM Migration Script ===${NC}"
    echo "This script migrates your system to LVM with proper boot configuration"
    echo "Fixes the issues from the previous failed migration"
    echo
    
    check_prerequisites
    detect_drives
    analyze_source_system
    calculate_partition_sizes
    
    echo
    echo "Migration Summary:"
    echo "  Source: $INTERNAL_DRIVE (current system)"
    echo "  Target: $EXTERNAL_DRIVE (new LVM system)"
    echo "  VG Name: $VG_NAME"
    echo "  Planned sizes:"
    echo "    Root: $CALCULATED_ROOT_SIZE"
    echo "    Home: $CALCULATED_HOME_SIZE"
    echo "    Swap: $CALCULATED_SWAP_SIZE"
    echo "    Boot: $CALCULATED_BOOT_SIZE"
    echo
    
    confirm_action "Start LVM migration?"
    
    create_lvm_layout
    handle_encrypted_partitions
    mount_filesystems
    copy_system_data
    configure_lvm_system
    install_bootloader
    verify_lvm_boot
    cleanup
    
    success "LVM migration completed successfully!"
    echo
    echo -e "${GREEN}=== MIGRATION COMPLETE ===${NC}"
    echo "✅ System migrated to LVM with proper boot support"
    echo "✅ Initramfs configured with LVM modules"
    echo "✅ GRUB installed with LVM support"
    echo "✅ Boot configuration verified"
    echo
    echo -e "${BLUE}Next steps:${NC}"
    echo "1. Reboot system"
    echo "2. Set external drive as first boot device in BIOS"
    echo "3. Boot should work without reset loops"
    echo "4. System will ask for any encryption passwords as normal"
    echo
    echo -e "${YELLOW}🎉 Improved LVM migration completed!${NC}"
    echo "This version fixes the boot issues from the previous attempt."
}

# Trap for cleanup
trap cleanup EXIT

main "$@"