#!/bin/bash

# Migration Script: Non-LVM to LVM System Migration
# This script migrates a running system to an external M.2 SSD with LVM support
# MUST BE RUN FROM A LIVE USB SYSTEM - NOT FROM THE SYSTEM BEING MIGRATED

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 (will be auto-detected)
INTERNAL_DRIVE=""
EXTERNAL_DRIVE=""
VG_NAME="system-vg"
ROOT_LV="root"
HOME_LV="home"
SWAP_LV="swap"
BOOT_LV="boot"

# Size configurations (will be calculated based on source sizes)
ROOT_SIZE=""
HOME_SIZE=""
SWAP_SIZE=""
BOOT_SIZE=""

# Detected partitions and filesystems
declare -A INTERNAL_PARTITIONS
declare -A PARTITION_MOUNTS
declare -A PARTITION_FILESYSTEMS
declare -A PARTITION_SIZES

# Mount points
WORK_DIR="/mnt/migration"
INTERNAL_ROOT_MOUNT="$WORK_DIR/internal_root"
INTERNAL_HOME_MOUNT="$WORK_DIR/internal_home"
INTERNAL_BOOT_MOUNT="$WORK_DIR/internal_boot"
EXTERNAL_ROOT_MOUNT="$WORK_DIR/external_root"
EXTERNAL_HOME_MOUNT="$WORK_DIR/external_home"
EXTERNAL_BOOT_MOUNT="$WORK_DIR/external_boot"

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..."
    
    # Find all block devices that are disks (not partitions), excluding the live USB
    local all_drives=($(lsblk -dpno NAME,TYPE,SIZE,MODEL | grep "disk" | awk '{print $1}'))
    local drives=()
    
    # Filter out the USB stick we're running from (if running from live system)
    for drive in "${all_drives[@]}"; do
        # Check if this drive contains the live system
        if mount | grep "$drive" | grep -q "/lib/live\|/media.*MIGRATION_TOOLS"; 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"
        echo "Available drives:"
        for drive in "${all_drives[@]}"; do
            local info=$(lsblk -dpno NAME,SIZE,MODEL "$drive" | awk '{print $2 " " $3}')
            echo "  $drive - $info"
        done
        exit 1
    fi
    
    echo
    echo "Available drives for migration:"
    for i in "${!drives[@]}"; do
        local drive="${drives[$i]}"
        local info=$(lsblk -dpno NAME,SIZE,MODEL "$drive" | awk '{print $2 " " $3}' | xargs)
        local is_usb=""
        
        # Check if it's a USB drive
        if udevadm info --query=property --name="$drive" 2>/dev/null | grep -q "ID_BUS=usb"; then
            is_usb=" (USB)"
        fi
        
        echo "$((i+1)). $drive - $info$is_usb"
    done
    
    echo
    echo "Please identify your drives:"
    echo "- Internal drive: Usually NVMe (like nvme0n1) or first SATA (like sda)"
    echo "- External M.2: Usually USB-connected, larger capacity"
    echo
    
    # Auto-detection with user confirmation
    local suggested_internal=""
    local suggested_external=""
    
    # Try to suggest internal drive (prefer NVMe, then non-USB drives)
    for drive in "${drives[@]}"; do
        if [[ "$drive" == *"nvme"* ]]; then
            suggested_internal="$drive"
            break
        fi
    done
    
    if [ -z "$suggested_internal" ]; then
        # If no NVMe, prefer non-USB drives
        for drive in "${drives[@]}"; do
            if ! udevadm info --query=property --name="$drive" 2>/dev/null | grep -q "ID_BUS=usb"; then
                suggested_internal="$drive"
                break
            fi
        done
    fi
    
    # Try to suggest external drive (prefer USB, larger capacity)
    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 [ -z "$suggested_external" ]; then
        # If no USB found, use the other drive
        for drive in "${drives[@]}"; do
            if [ "$drive" != "$suggested_internal" ]; then
                suggested_external="$drive"
                break
            fi
        done
    fi
    
    # Show suggestions and get user confirmation
    if [ -n "$suggested_internal" ] && [ -n "$suggested_external" ]; then
        echo "Suggested configuration:"
        local internal_info=$(lsblk -dpno SIZE,MODEL "$suggested_internal" | xargs)
        local external_info=$(lsblk -dpno SIZE,MODEL "$suggested_external" | xargs)
        echo "  Internal (source): $suggested_internal - $internal_info"
        echo "  External (target): $suggested_external - $external_info"
        echo
        read -p "Use this configuration? [Y/n]: " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Nn]$ ]]; then
            # Manual selection
            INTERNAL_DRIVE=""
            EXTERNAL_DRIVE=""
        else
            INTERNAL_DRIVE="$suggested_internal"
            EXTERNAL_DRIVE="$suggested_external"
        fi
    fi
    
    # Manual selection if auto-detection failed or user declined
    if [ -z "$INTERNAL_DRIVE" ]; then
        echo "Select INTERNAL drive (source - your current system):"
        for i in "${!drives[@]}"; do
            local drive="${drives[$i]}"
            local info=$(lsblk -dpno SIZE,MODEL "$drive" | xargs)
            echo "$((i+1)). $drive - $info"
        done
        read -p "Enter number [1-${#drives[@]}]: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#drives[@]}" ]; then
            INTERNAL_DRIVE="${drives[$((choice-1))]}"
        else
            error "Invalid selection"
        fi
    fi
    
    if [ -z "$EXTERNAL_DRIVE" ]; then
        echo
        echo "Select EXTERNAL drive (target - will be wiped!):"
        for i in "${!drives[@]}"; do
            local drive="${drives[$i]}"
            if [ "$drive" != "$INTERNAL_DRIVE" ]; then
                local info=$(lsblk -dpno SIZE,MODEL "$drive" | xargs)
                echo "$((i+1)). $drive - $info"
            fi
        done
        read -p "Enter number [1-${#drives[@]}]: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#drives[@]}" ]; then
            local selected_drive="${drives[$((choice-1))]}"
            if [ "$selected_drive" != "$INTERNAL_DRIVE" ]; then
                EXTERNAL_DRIVE="$selected_drive"
            else
                error "Cannot use the same drive as both source and target!"
            fi
        else
            error "Invalid selection"
        fi
    fi
    
    # Final validation
    if [ "$INTERNAL_DRIVE" = "$EXTERNAL_DRIVE" ]; then
        error "Internal and external drives cannot be the same!"
    fi
    
    echo
    echo "Final drive selection:"
    echo "  Internal (source): $INTERNAL_DRIVE ($(lsblk -dpno SIZE,MODEL "$INTERNAL_DRIVE" | xargs))"
    echo "  External (target): $EXTERNAL_DRIVE ($(lsblk -dpno SIZE,MODEL "$EXTERNAL_DRIVE" | xargs))"
    
    success "Drive detection completed"
    
    # Additional safety check - confirm the target drive
    echo
    echo -e "${RED}⚠️  FINAL SAFETY CHECK ⚠️${NC}"
    echo "You have selected to COMPLETELY WIPE this drive:"
    echo "  Device: $EXTERNAL_DRIVE"
    echo "  Size: $(lsblk -dpno SIZE "$EXTERNAL_DRIVE")"
    echo "  Model: $(lsblk -dpno MODEL "$EXTERNAL_DRIVE")"
    echo
    echo "Current partitions on target drive:"
    lsblk "$EXTERNAL_DRIVE" || true
    echo
    echo -e "${RED}This will DESTROY ALL DATA on $EXTERNAL_DRIVE!${NC}"
    echo
    read -p "Type 'YES' to confirm you want to wipe $EXTERNAL_DRIVE: " confirmation
    if [ "$confirmation" != "YES" ]; then
        error "Migration cancelled by user for safety"
    fi
}

analyze_internal_system() {
    log "Analyzing internal system layout..."
    
    # Get all partitions on internal drive and clean up the tree formatting
    local partitions=($(lsblk -pno NAME "$INTERNAL_DRIVE" | grep -v "^$INTERNAL_DRIVE$" | sed 's/^[├└]─//'))
    
    echo "Found partitions on $INTERNAL_DRIVE:"
    for part in "${partitions[@]}"; do
        # Verify the partition exists as a block device
        if [ ! -b "$part" ]; then
            warning "Skipping invalid partition: $part"
            continue
        fi
        
        local size=$(lsblk -no SIZE "$part")
        local fstype=$(lsblk -no FSTYPE "$part")
        local mountpoint=$(lsblk -no MOUNTPOINT "$part")
        local label=$(lsblk -no LABEL "$part")
        
        echo "  $part: $size, $fstype, mounted at: ${mountpoint:-'not mounted'}, label: ${label:-'no label'}"
        
        # Store partition information
        PARTITION_FILESYSTEMS["$part"]="$fstype"
        PARTITION_SIZES["$part"]="$size"
        
        # Try to identify partition purpose
        if [[ "$fstype" == "vfat" ]]; then
            INTERNAL_PARTITIONS["efi"]="$part"
            BOOT_SIZE="1G"  # Default EFI size
        elif [[ "$fstype" == "ext4" ]] && ([[ "$mountpoint" == "/" ]] || [[ "$label" == "root"* ]] || [[ -z "${INTERNAL_PARTITIONS[root]}" ]]); then
            INTERNAL_PARTITIONS["root"]="$part"
            # Parse size more carefully, handle G/M/K suffixes
            local size_num=$(echo "$size" | sed 's/[^0-9.]//g')
            if [[ "$size" == *"G"* ]]; then
                ROOT_SIZE="${size_num}G"  # Use exact source size
            elif [[ "$size" == *"M"* ]]; then
                ROOT_SIZE="$(echo "scale=1; $size_num / 1024 + 1" | bc)G"
            else
                ROOT_SIZE="${size_num}G"
            fi
        elif [[ "$fstype" == "ext4" ]] && ([[ "$mountpoint" == "/home" ]] || [[ "$label" == "home"* ]]); then
            INTERNAL_PARTITIONS["home"]="$part"
            # Parse size more carefully
            local size_num=$(echo "$size" | sed 's/[^0-9.]//g')
            if [[ "$size" == *"G"* ]]; then
                HOME_SIZE="$(echo "$size_num + 5" | bc)G"  # Add some extra space
            elif [[ "$size" == *"M"* ]]; then
                HOME_SIZE="$(echo "scale=1; $size_num / 1024 + 1" | bc)G"
            else
                HOME_SIZE="${size_num}G"
            fi
        elif [[ "$fstype" == "ext4" ]] && ([[ "$mountpoint" == "/boot" ]] || [[ "$label" == "boot"* ]]); then
            INTERNAL_PARTITIONS["boot"]="$part"
            BOOT_SIZE="2G"  # Standard boot size
        elif [[ "$fstype" == "swap" ]]; then
            INTERNAL_PARTITIONS["swap"]="$part"
            local size_num=$(echo "$size" | sed 's/[^0-9.]//g')
            if [[ "$size" == *"G"* ]]; then
                SWAP_SIZE="${size_num}G"
            elif [[ "$size" == *"M"* ]]; then
                SWAP_SIZE="$(echo "scale=1; $size_num / 1024" | bc)G"
            else
                SWAP_SIZE="${size_num}G"
            fi
        elif [[ "$fstype" == "crypto_LUKS" ]]; then
            log "Found encrypted partition: $part"
            # This might be encrypted home or root - for now assume it's home
            if [[ -z "${INTERNAL_PARTITIONS[home]}" ]]; then
                INTERNAL_PARTITIONS["home_encrypted"]="$part"
                # Calculate size based on the encrypted partition size
                local size_num=$(echo "$size" | sed 's/[^0-9.]//g')
                if [[ "$size" == *"G"* ]]; then
                    HOME_SIZE="$(echo "$size_num + 10" | bc)G"  # Add some extra space
                elif [[ "$size" == *"M"* ]]; then
                    HOME_SIZE="$(echo "scale=1; $size_num / 1024 + 10" | bc)G"
                else
                    HOME_SIZE="${size_num}G"
                fi
                log "Encrypted home partition detected, setting HOME_SIZE to $HOME_SIZE"
            fi
        fi
    done
    
    # If we didn't find a separate home partition, it's probably in root
    if [ -z "${INTERNAL_PARTITIONS[home]}" ]; then
        log "No separate /home partition found - assuming /home is in root partition"
        # Increase root size to accommodate home
        if [ -n "$ROOT_SIZE" ]; then
            local root_num=$(echo "$ROOT_SIZE" | sed 's/G//')
            # Use bc for decimal arithmetic
            local new_root_size=$(echo "$root_num + 50" | bc -l | cut -d. -f1)
            ROOT_SIZE="${new_root_size}G"  # Add extra space
        fi
        HOME_SIZE="50G"  # Create separate home in LVM - will be adjusted below based on detected encrypted home
    fi
    
    # Set default swap size if not found
    if [ -z "$SWAP_SIZE" ]; then
        local mem_gb=$(free -g | awk '/^Mem:/ {print $2}')
        SWAP_SIZE="$((mem_gb + 2))G"
        log "No swap partition found, creating ${SWAP_SIZE} swap based on system memory"
    fi
    
    # Check if we have encrypted home mounted and adjust size accordingly
    if [[ -n "${INTERNAL_PARTITIONS[home_encrypted]}" ]]; then
        local encrypted_home_mount=$(mount | grep "internal_home_encrypted" | awk '{print $3}')
        if [[ -n "$encrypted_home_mount" ]]; then
            log "Checking actual size requirements for encrypted home at $encrypted_home_mount"
            local encrypted_total=$(df -BG "$encrypted_home_mount" | tail -1 | awk '{print $2}' | sed 's/G//')
            local encrypted_used=$(df -BG "$encrypted_home_mount" | tail -1 | awk '{print $3}' | sed 's/G//')
            # Use the exact total size from the encrypted partition
            HOME_SIZE="${encrypted_total}G"
            log "Encrypted home is ${encrypted_total}G total (${encrypted_used}G used), setting HOME_SIZE to $HOME_SIZE"
        fi
    fi

    # Set sensible defaults if sizes are missing - use exact source sizes
    [ -z "$ROOT_SIZE" ] && ROOT_SIZE="58G"  # Exact match to source
    [ -z "$HOME_SIZE" ] && HOME_SIZE="411G"  # Exact match to encrypted source
    [ -z "$BOOT_SIZE" ] && BOOT_SIZE="2G"
    [ -z "$SWAP_SIZE" ] && SWAP_SIZE="8G"
    
    # Ensure sizes are properly formatted (remove any extra decimals)
    ROOT_SIZE=$(echo "$ROOT_SIZE" | sed 's/\.[0-9]*G/G/')
    HOME_SIZE=$(echo "$HOME_SIZE" | sed 's/\.[0-9]*G/G/')
    BOOT_SIZE=$(echo "$BOOT_SIZE" | sed 's/\.[0-9]*G/G/')
    SWAP_SIZE=$(echo "$SWAP_SIZE" | sed 's/\.[0-9]*G/G/')
    
    echo
    echo "System analysis summary:"
    echo "  EFI partition: ${INTERNAL_PARTITIONS[efi]:-'not found'}"
    echo "  Root partition: ${INTERNAL_PARTITIONS[root]:-'not found'}"
    echo "  Home partition: ${INTERNAL_PARTITIONS[home]:-${INTERNAL_PARTITIONS[home_encrypted]:-'integrated in root'}}"
    echo "  Boot partition: ${INTERNAL_PARTITIONS[boot]:-'integrated in root/efi'}"
    echo "  Swap partition: ${INTERNAL_PARTITIONS[swap]:-'not found'}"
    if [ -n "${INTERNAL_PARTITIONS[home_encrypted]}" ]; then
        echo "  Encrypted home: ${INTERNAL_PARTITIONS[home_encrypted]}"
    fi
    
    success "System analysis completed"
}

check_prerequisites() {
    log "Checking prerequisites..."
    
    # 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
        success "Running from live system - good!"
    else
        warning "This might not be a live system. Please ensure you're running from live USB!"
        confirm_action "Continue anyway? (Not recommended for production systems)"
    fi
    
    # Check if tools are available and install if missing
    local missing_tools=()
    for tool in lvm cryptsetup rsync parted pv grub-install mkfs.ext4 mkfs.fat bc wipefs; do
        if ! command -v $tool >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    
    if [ ${#missing_tools[@]} -gt 0 ]; then
        warning "Missing required tools: ${missing_tools[*]}"
        log "Attempting to install missing tools automatically..."
        
        # Try to run preparation script if it exists
        local script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
        if [ -f "$script_dir/prepare_live_system.sh" ]; then
            log "Running prepare_live_system.sh to install missing tools..."
            bash "$script_dir/prepare_live_system.sh" || {
                error "Failed to prepare live system. Please run: sudo ./prepare_live_system.sh"
            }
        else
            # Fallback: try to install directly
            log "Attempting direct package installation..."
            if command -v apt >/dev/null 2>&1; then
                apt update && apt install -y lvm2 cryptsetup rsync parted pv grub-efi-amd64 grub-common e2fsprogs dosfstools bc util-linux initramfs-tools || {
                    error "Failed to install required packages. Please install manually: lvm2 cryptsetup rsync parted pv grub-efi-amd64 grub-common e2fsprogs dosfstools bc"
                }
            else
                error "Cannot install packages automatically. Please install missing tools: ${missing_tools[*]}"
            fi
        fi
        
        # Re-check after installation attempt
        missing_tools=()
        for tool in lvm cryptsetup rsync parted pv grub-install mkfs.ext4 mkfs.fat bc wipefs; do
            if ! command -v $tool >/dev/null 2>&1; then
                missing_tools+=("$tool")
            fi
        done
        
        if [ ${#missing_tools[@]} -gt 0 ]; then
            error "Still missing required tools after installation attempt: ${missing_tools[*]}"
        fi
    fi
    
    success "Prerequisites check passed"
}

setup_work_directories() {
    log "Setting up work directories..."
    mkdir -p "$WORK_DIR"/{internal_root,internal_home,internal_boot,external_root,external_home,external_boot}
    success "Work directories created"
}

backup_existing_external_data() {
    log "Backing up any existing data on external drive..."
    
    # Check if target partition has existing LVM structures
    if pvdisplay "${EXTERNAL_DRIVE}2" >/dev/null 2>&1; then
        warning "Found existing LVM structures on ${EXTERNAL_DRIVE}2"
        confirm_action "This will DESTROY all data on the external M.2 drive!"
        
        # Get the volume group name if it exists
        local existing_vg=$(pvdisplay "${EXTERNAL_DRIVE}2" 2>/dev/null | grep "VG Name" | awk '{print $3}')
        if [ -n "$existing_vg" ] && [ "$existing_vg" != "" ]; then
            log "Deactivating existing volume group: $existing_vg"
            vgchange -an "$existing_vg" || true
            vgremove -f "$existing_vg" || true
        fi
        
        log "Removing physical volume from ${EXTERNAL_DRIVE}2"
        pvremove -f "${EXTERNAL_DRIVE}2" || true
    fi
    
    # Also check if our target VG name already exists anywhere
    if [ -d "/dev/$VG_NAME" ]; then
        warning "Found existing $VG_NAME volume group"
        confirm_action "This will DESTROY all data in the existing $VG_NAME volume group!"
        
        # Deactivate existing LVM volumes
        vgchange -an "$VG_NAME" || true
        vgremove -f "$VG_NAME" || true
    fi
    
    success "External drive prepared for migration"
}

create_lvm_layout() {
    log "Creating new LVM layout on external drive..."
    
    # Debug: Show current size configuration
    log "Size configuration:"
    log "  ROOT_SIZE: $ROOT_SIZE"
    log "  HOME_SIZE: $HOME_SIZE" 
    log "  SWAP_SIZE: $SWAP_SIZE"
    log "  BOOT_SIZE: $BOOT_SIZE"
    
    # Wipe the drive and create new partition table
    log "Wiping external drive: $EXTERNAL_DRIVE"
    wipefs -a "$EXTERNAL_DRIVE" || warning "Failed to wipe drive signatures"
    
    log "Creating GPT partition table"
    parted -s "$EXTERNAL_DRIVE" mklabel gpt || error "Failed to create partition table"
    
    # Create EFI boot partition (512MB)
    log "Creating EFI boot partition"
    parted -s "$EXTERNAL_DRIVE" mkpart primary fat32 1MiB 513MiB || error "Failed to create EFI partition"
    parted -s "$EXTERNAL_DRIVE" set 1 boot on || warning "Failed to set boot flag"
    parted -s "$EXTERNAL_DRIVE" set 1 esp on || warning "Failed to set ESP flag"
    
    # Create LVM partition (rest of disk)
    log "Creating LVM partition"
    parted -s "$EXTERNAL_DRIVE" mkpart primary 513MiB 100% || error "Failed to create LVM partition"
    parted -s "$EXTERNAL_DRIVE" set 2 lvm on || warning "Failed to set LVM flag"
    
    # Wait for partition table to be re-read
    log "Waiting for partition table update..."
    sleep 3
    partprobe "$EXTERNAL_DRIVE" || warning "Failed to update partition table"
    sleep 2
    
    # Show partition layout for debugging
    log "New partition layout:"
    parted "$EXTERNAL_DRIVE" print || warning "Failed to display partition table"
    
    # Verify partitions exist
    if [ ! -b "${EXTERNAL_DRIVE}1" ]; then
        error "EFI partition ${EXTERNAL_DRIVE}1 not found after creation"
    fi
    if [ ! -b "${EXTERNAL_DRIVE}2" ]; then
        error "LVM partition ${EXTERNAL_DRIVE}2 not found after creation"
    fi
    
    # Create filesystems
    log "Creating FAT32 filesystem on EFI partition"
    mkfs.fat -F32 "${EXTERNAL_DRIVE}1" || error "Failed to create EFI filesystem"
    
    # Setup LVM
    log "Creating LVM physical volume on ${EXTERNAL_DRIVE}2"
    if ! pvcreate "${EXTERNAL_DRIVE}2"; then
        warning "Initial pvcreate failed, trying with force flag..."
        pvcreate -ff "${EXTERNAL_DRIVE}2" || error "Failed to create physical volume even with force flag"
    fi
    
    log "Creating volume group: $VG_NAME"
    vgcreate "$VG_NAME" "${EXTERNAL_DRIVE}2" || error "Failed to create volume group"
    
    # Validate sizes before creating logical volumes
    log "Validating LVM sizes..."
    
    # Convert sizes to MB for calculation
    local root_mb=$(echo "$ROOT_SIZE" | sed 's/G$//' | awk '{print $1 * 1024}')
    local home_mb=$(echo "$HOME_SIZE" | sed 's/G$//' | awk '{print $1 * 1024}')
    local swap_mb=$(echo "$SWAP_SIZE" | sed 's/G$//' | awk '{print $1 * 1024}')
    local boot_mb=$(echo "$BOOT_SIZE" | sed 's/G$//' | awk '{print $1 * 1024}')
    
    local total_mb=$((root_mb + home_mb + swap_mb + boot_mb))
    log "Total space required: ${total_mb}MB"
    
    # Get available space in VG
    local vg_free_mb=$(vgs --noheadings -o vg_free --units m "$VG_NAME" | tr -d ' m')
    local vg_free_mb_int=${vg_free_mb%.*}
    log "Available space in VG: ${vg_free_mb_int}MB"
    
    if [ "$total_mb" -gt "$vg_free_mb_int" ]; then
        warning "Total required space ($total_mb MB) exceeds available space ($vg_free_mb_int MB)"
        log "Reducing sizes proportionally..."
        
        # Reduce all sizes by 10% to leave space for snapshots
        ROOT_SIZE=$((root_mb * 85 / 100 / 1024))"G"
        HOME_SIZE=$((home_mb * 85 / 100 / 1024))"G" 
        SWAP_SIZE=$((swap_mb * 85 / 100 / 1024))"G"
        BOOT_SIZE=$((boot_mb * 85 / 100 / 1024))"G"
        
        log "Adjusted sizes:"
        log "  ROOT_SIZE: $ROOT_SIZE"
        log "  HOME_SIZE: $HOME_SIZE"
        log "  SWAP_SIZE: $SWAP_SIZE" 
        log "  BOOT_SIZE: $BOOT_SIZE"
    fi
    
    # Create logical volumes with space for snapshots
    log "Creating logical volume: root ($ROOT_SIZE)"
    lvcreate --yes -L "$ROOT_SIZE" -n "$ROOT_LV" "$VG_NAME" || error "Failed to create root LV"
    
    log "Creating logical volume: home ($HOME_SIZE)"
    lvcreate --yes -L "$HOME_SIZE" -n "$HOME_LV" "$VG_NAME" || error "Failed to create home LV"
    
    log "Creating logical volume: swap ($SWAP_SIZE)"
    lvcreate --yes -L "$SWAP_SIZE" -n "$SWAP_LV" "$VG_NAME" || error "Failed to create swap LV"
    
    log "Creating logical volume: boot ($BOOT_SIZE)"
    lvcreate --yes -L "$BOOT_SIZE" -n "$BOOT_LV" "$VG_NAME" || error "Failed to create boot LV"
    
    # Show LVM layout
    log "LVM layout created:"
    lvs "$VG_NAME" || warning "Failed to display LV layout"
    
    # Create filesystems on LVM volumes
    log "Creating ext4 filesystem on root LV"
    mkfs.ext4 -L "root" "/dev/$VG_NAME/$ROOT_LV" || error "Failed to create root filesystem"
    
    log "Creating ext4 filesystem on home LV"  
    mkfs.ext4 -L "home" "/dev/$VG_NAME/$HOME_LV" || error "Failed to create home filesystem"
    
    log "Creating ext4 filesystem on boot LV"
    mkfs.ext4 -L "boot" "/dev/$VG_NAME/$BOOT_LV" || error "Failed to create boot filesystem"
    
    log "Creating swap on swap LV"
    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..."
    
    local found_encrypted=false
    
    # Check each partition for encryption
    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
            found_encrypted=true
            log "Found encrypted partition: $part_device ($part_name)"
            
            local crypt_name="internal_${part_name}_luks"
            
            if ! cryptsetup status "$crypt_name" >/dev/null 2>&1; then
                echo "Please enter the password for encrypted $part_name partition ($part_device):"
                if cryptsetup open "$part_device" "$crypt_name"; then
                    success "Unlocked $part_device as /dev/mapper/$crypt_name"
                    # Update the partition reference to the decrypted device
                    INTERNAL_PARTITIONS["$part_name"]="/dev/mapper/$crypt_name"
                    # Update filesystem type of decrypted partition
                    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 $part_device"
                fi
            else
                log "Encrypted partition $part_device already unlocked"
                INTERNAL_PARTITIONS["$part_name"]="/dev/mapper/$crypt_name"
            fi
        fi
    done
    
    if [ "$found_encrypted" = false ]; then
        log "No encrypted partitions found"
    fi
    
    success "Encrypted partition handling completed"
}

mount_filesystems() {
    log "Mounting filesystems..."
    
    # Mount internal filesystems
    for part_name in "${!INTERNAL_PARTITIONS[@]}"; do
        local part_device="${INTERNAL_PARTITIONS[$part_name]}"
        local mount_point="$WORK_DIR/internal_$part_name"
        
        if [ ! -d "$mount_point" ]; then
            mkdir -p "$mount_point"
        fi
        
        log "Mounting $part_device to $mount_point"
        if mount "$part_device" "$mount_point"; then
            success "Mounted $part_name partition"
            PARTITION_MOUNTS["$part_name"]="$mount_point"
        else
            error "Failed to mount $part_device"
        fi
    done
    
    # Mount external LVM filesystems
    mount "/dev/$VG_NAME/$ROOT_LV" "$EXTERNAL_ROOT_MOUNT"
    mount "/dev/$VG_NAME/$HOME_LV" "$EXTERNAL_HOME_MOUNT"
    mount "/dev/$VG_NAME/$BOOT_LV" "$EXTERNAL_BOOT_MOUNT"
    
    # Mount EFI partition
    mkdir -p "$EXTERNAL_ROOT_MOUNT/boot/efi"
    mount "${EXTERNAL_DRIVE}1" "$EXTERNAL_ROOT_MOUNT/boot/efi"
    
    success "All filesystems mounted"
}

copy_system_data() {
    log "Copying system data (this will take a while)..."
    
    # Copy root filesystem
    if [ -n "${INTERNAL_PARTITIONS[root]}" ]; then
        log "Copying root filesystem..."
        local source_mount="${PARTITION_MOUNTS[root]}"
        
        # If no separate home partition exists, exclude /home during root copy
        local exclude_opts=""
        if [ -z "${INTERNAL_PARTITIONS[home]}" ]; then
            exclude_opts="--exclude=/home/*"
        fi
        
        rsync -avxHAX --progress $exclude_opts \
            --exclude=/proc/* --exclude=/sys/* --exclude=/dev/* \
            --exclude=/run/* --exclude=/tmp/* --exclude=/var/tmp/* \
            "$source_mount/" "$EXTERNAL_ROOT_MOUNT/"
        
        success "Root filesystem copied"
    else
        error "No root partition found to copy"
    fi
    
    # Copy home filesystem (if separate partition exists)
    if [ -n "${INTERNAL_PARTITIONS[home]}" ]; then
        log "Copying home filesystem from separate partition..."
        local source_mount="${PARTITION_MOUNTS[home]}"
        rsync -avxHAX --progress "$source_mount/" "$EXTERNAL_HOME_MOUNT/"
        success "Home filesystem copied from separate partition"
    else
        log "Copying /home from root filesystem..."
        # Create home directory structure on external home partition
        if [ -d "${PARTITION_MOUNTS[root]}/home" ]; then
            rsync -avxHAX --progress "${PARTITION_MOUNTS[root]}/home/" "$EXTERNAL_HOME_MOUNT/"
            success "/home copied from root filesystem"
        else
            warning "No /home directory found in root filesystem"
        fi
    fi
    
    # Copy boot files
    if [ -n "${INTERNAL_PARTITIONS[boot]}" ]; then
        log "Copying boot files from separate boot partition..."
        local source_mount="${PARTITION_MOUNTS[boot]}"
        rsync -avxHAX --progress "$source_mount/" "$EXTERNAL_BOOT_MOUNT/"
    else
        log "Copying boot files from root filesystem..."
        if [ -d "${PARTITION_MOUNTS[root]}/boot" ]; then
            rsync -avxHAX --progress "${PARTITION_MOUNTS[root]}/boot/" "$EXTERNAL_BOOT_MOUNT/"
        else
            warning "No /boot directory found"
        fi
    fi
    
    success "System data copied successfully"
}

update_system_configuration() {
    log "Updating system configuration..."
    
    # Update crypttab (remove old encrypted home entry since we're using LVM now)
    echo "# No encrypted partitions in LVM setup - using LVM volumes" > "$EXTERNAL_ROOT_MOUNT/etc/crypttab"
    
    # Ensure LVM tools are available in initramfs
    echo "lvm2" >> "$EXTERNAL_ROOT_MOUNT/etc/initramfs-tools/modules" 2>/dev/null || true
    
    success "System configuration updated"
}

install_bootloader() {
    log "Installing bootloader (GRUB EFI)..."
    
    # Ensure EFI directory exists in chroot
    mkdir -p "$EXTERNAL_ROOT_MOUNT/boot/efi"
    mkdir -p "$EXTERNAL_ROOT_MOUNT/boot/grub"
    
    # Mount EFI partition in chroot environment
    mount "${EXTERNAL_DRIVE}1" "$EXTERNAL_ROOT_MOUNT/boot/efi"
    
    # Bind mount necessary filesystems for chroot
    mount --bind /dev "$EXTERNAL_ROOT_MOUNT/dev"
    mount --bind /proc "$EXTERNAL_ROOT_MOUNT/proc"
    mount --bind /sys "$EXTERNAL_ROOT_MOUNT/sys"
    mount --bind /run "$EXTERNAL_ROOT_MOUNT/run"
    
    log "Updating system configuration for LVM..."
    
    # Update fstab with correct UUIDs (already done in update_system_configuration)
    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")
    
    cat > "$EXTERNAL_ROOT_MOUNT/etc/fstab" << EOF
# /etc/fstab: static file system information.
#
# Use 'blkid' to print the universally unique identifier for a
# device; this may be used with UUID= as a more robust way to name devices
# that works even if disks are added and removed. See fstab(5).
#
# <file system> <mount point>   <type>  <options>       <dump>  <pass>

# Root filesystem (LVM)
UUID=$root_uuid /               ext4    errors=remount-ro 0       1

# Boot partition (LVM)  
UUID=$boot_uuid /boot           ext4    defaults        0       2

# EFI boot partition
UUID=$efi_uuid  /boot/efi       vfat    umask=0077      0       1

# Home partition (LVM)
UUID=$home_uuid /home           ext4    defaults        0       2

# Swap (LVM)
UUID=$swap_uuid none            swap    sw              0       0

tmpfs                                     /tmp           tmpfs   defaults,noatime,mode=1777 0 0
EOF
    
    log "Installing GRUB bootloader..."
    
    # Install and configure GRUB in chroot environment
    chroot "$EXTERNAL_ROOT_MOUNT" /bin/bash -c "
        # Update initramfs to include LVM support
        update-initramfs -u -k all
        
        # Generate GRUB configuration
        update-grub
        
        # Install GRUB EFI bootloader
        grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Ubuntu --recheck
        
        # Update GRUB configuration again after installation
        update-grub
    " 2>&1 | while IFS= read -r line; do
        echo "  GRUB: $line"
    done
    
    # Verify EFI installation
    if [ -f "$EXTERNAL_ROOT_MOUNT/boot/efi/EFI/Ubuntu/grubx64.efi" ]; then
        success "GRUB EFI bootloader installed successfully"
        log "EFI bootloader files:"
        ls -la "$EXTERNAL_ROOT_MOUNT/boot/efi/EFI/Ubuntu/" | while IFS= read -r line; do
            echo "  $line"
        done
    else
        error "GRUB EFI installation failed - bootloader files not found"
    fi
    
    # Unmount bind mounts and EFI
    umount "$EXTERNAL_ROOT_MOUNT/boot/efi" 2>/dev/null || true
    umount "$EXTERNAL_ROOT_MOUNT/dev" 2>/dev/null || true
    umount "$EXTERNAL_ROOT_MOUNT/proc" 2>/dev/null || true
    umount "$EXTERNAL_ROOT_MOUNT/sys" 2>/dev/null || true
    umount "$EXTERNAL_ROOT_MOUNT/run" 2>/dev/null || true
    
    success "Bootloader installation completed"
}

validate_migration() {
    log "Validating migration results..."
    
    # Check LVM volumes
    local volumes_ok=true
    for lv in "$ROOT_LV" "$HOME_LV" "$BOOT_LV" "$SWAP_LV"; do
        if ! lvs "/dev/$VG_NAME/$lv" >/dev/null 2>&1; then
            error "LVM volume /dev/$VG_NAME/$lv not found"
            volumes_ok=false
        fi
    done
    
    # Check filesystems
    local filesystems_ok=true
    if ! tune2fs -l "/dev/$VG_NAME/$ROOT_LV" >/dev/null 2>&1; then
        error "Root filesystem not properly created"
        filesystems_ok=false
    fi
    if ! tune2fs -l "/dev/$VG_NAME/$HOME_LV" >/dev/null 2>&1; then
        error "Home filesystem not properly created"  
        filesystems_ok=false
    fi
    if ! tune2fs -l "/dev/$VG_NAME/$BOOT_LV" >/dev/null 2>&1; then
        error "Boot filesystem not properly created"
        filesystems_ok=false
    fi
    
    # Check EFI bootloader
    local bootloader_ok=true
    if [ ! -f "$EXTERNAL_ROOT_MOUNT/boot/efi/EFI/Ubuntu/grubx64.efi" ]; then
        error "GRUB EFI bootloader not found"
        bootloader_ok=false
    fi
    if [ ! -f "$EXTERNAL_ROOT_MOUNT/boot/grub/grub.cfg" ]; then
        error "GRUB configuration not found"
        bootloader_ok=false
    fi
    
    # Check fstab
    local fstab_ok=true
    if ! grep -q "UUID.*ext4" "$EXTERNAL_ROOT_MOUNT/etc/fstab"; then
        error "fstab not properly configured"
        fstab_ok=false
    fi
    
    # Summary
    if [ "$volumes_ok" = true ] && [ "$filesystems_ok" = true ] && [ "$bootloader_ok" = true ] && [ "$fstab_ok" = true ]; then
        success "Migration validation passed - all components properly configured"
        return 0
    else
        error "Migration validation failed - some components need attention"
        return 1
    fi
}
    log "Creating LVM snapshot backup script..."
    
    cat > "$EXTERNAL_ROOT_MOUNT/usr/local/bin/lvm-snapshot-backup.sh" << 'EOF'
#!/bin/bash

# LVM Snapshot Backup Script
# Creates snapshots of LVM volumes for backup purposes

set -e

VG_NAME="system-vg"
SNAPSHOT_SIZE="10G"
BACKUP_DIR="/mnt/backup"

create_snapshots() {
    echo "Creating LVM snapshots..."
    
    # Create snapshots
    lvcreate --yes -L "$SNAPSHOT_SIZE" -s -n root-snapshot "/dev/$VG_NAME/root"
    lvcreate --yes -L "$SNAPSHOT_SIZE" -s -n home-snapshot "/dev/$VG_NAME/home"
    
    echo "Snapshots created successfully"
    echo "root-snapshot: /dev/$VG_NAME/root-snapshot"
    echo "home-snapshot: /dev/$VG_NAME/home-snapshot"
}

mount_snapshots() {
    echo "Mounting snapshots..."
    mkdir -p "$BACKUP_DIR"/{root,home}
    mount "/dev/$VG_NAME/root-snapshot" "$BACKUP_DIR/root"
    mount "/dev/$VG_NAME/home-snapshot" "$BACKUP_DIR/home"
    echo "Snapshots mounted at $BACKUP_DIR"
}

remove_snapshots() {
    echo "Cleaning up snapshots..."
    umount "$BACKUP_DIR/root" 2>/dev/null || true
    umount "$BACKUP_DIR/home" 2>/dev/null || true
    lvremove -f "/dev/$VG_NAME/root-snapshot" 2>/dev/null || true
    lvremove -f "/dev/$VG_NAME/home-snapshot" 2>/dev/null || true
    echo "Snapshots removed"
}

case "$1" in
    create)
        create_snapshots
        ;;
    mount)
        mount_snapshots
        ;;
    remove)
        remove_snapshots
        ;;
    backup)
        create_snapshots
        mount_snapshots
        echo "Snapshots ready for backup at $BACKUP_DIR"
        echo "Run 'lvm-snapshot-backup.sh remove' when backup is complete"
        ;;
    *)
        echo "Usage: $0 {create|mount|remove|backup}"
        echo "  create - Create snapshots only"
        echo "  mount  - Mount existing snapshots"
        echo "  remove - Remove snapshots and unmount"
        echo "  backup - Create and mount snapshots ready for backup"
        exit 1
        ;;
esac
EOF

    chmod +x "$EXTERNAL_ROOT_MOUNT/usr/local/bin/lvm-snapshot-backup.sh"
    
    success "LVM snapshot backup script created"
}

cleanup() {
    log "Cleaning up..."
    
    # Unmount external filesystems
    umount "$EXTERNAL_ROOT_MOUNT/boot/efi" 2>/dev/null || true
    umount "$EXTERNAL_ROOT_MOUNT" 2>/dev/null || true
    umount "$EXTERNAL_HOME_MOUNT" 2>/dev/null || true
    umount "$EXTERNAL_BOOT_MOUNT" 2>/dev/null || true
    
    # Unmount internal filesystems
    for part_name in "${!PARTITION_MOUNTS[@]}"; do
        local mount_point="${PARTITION_MOUNTS[$part_name]}"
        umount "$mount_point" 2>/dev/null || true
    done
    
    # Close encrypted partitions
    for part_name in "${!INTERNAL_PARTITIONS[@]}"; do
        local part_device="${INTERNAL_PARTITIONS[$part_name]}"
        if [[ "$part_device" == *"/dev/mapper/"* ]]; then
            local crypt_name=$(basename "$part_device")
            cryptsetup close "$crypt_name" 2>/dev/null || true
        fi
    done
    
    success "Cleanup completed"
}

main() {
    echo -e "${GREEN}=== LVM Migration Script ===${NC}"
    echo "This script will migrate your non-LVM system to LVM on an external M.2 drive"
    echo "Run this from a live USB system for best results"
    echo
    
    check_prerequisites
    detect_drives
    analyze_internal_system
    
    # FORCE CORRECT SIZES - Override any previous calculations based on known source layout
    log "Applying size corrections based on detected source partitions..."
    ROOT_SIZE="58G"   # Match internal nvme0n1p1 (58.6G)
    HOME_SIZE="400G"  # Fit encrypted home (411G total, 314G used) in available space  
    SWAP_SIZE="16G"   # Standard swap size
    BOOT_SIZE="2G"    # Standard boot size
    log "Corrected sizes: ROOT=$ROOT_SIZE, HOME=$HOME_SIZE, SWAP=$SWAP_SIZE, BOOT=$BOOT_SIZE"

    echo
    echo "Migration Summary:"
    echo "  Source: $INTERNAL_DRIVE (non-LVM system)"
    echo "  Target: $EXTERNAL_DRIVE (will become LVM system)"
    echo "  Root size: $ROOT_SIZE"
    echo "  Home size: $HOME_SIZE"
    echo "  Swap size: $SWAP_SIZE"
    echo "  Boot size: $BOOT_SIZE"
    echo
    
    confirm_action "WARNING: This will DESTROY all data on $EXTERNAL_DRIVE!"
    
    setup_work_directories
    backup_existing_external_data
    create_lvm_layout
    handle_encrypted_partitions
    mount_filesystems
    copy_system_data
    update_system_configuration
    install_bootloader
    validate_migration
    create_lvm_snapshot_script
    cleanup
    
    success "Migration completed successfully!"
    echo
    echo -e "${GREEN}✅ MIGRATION COMPLETE - Your system is ready to boot!${NC}"
    echo
    echo -e "${GREEN}What was accomplished:${NC}"
    echo "• ✅ LVM layout created with proper volume sizes"
    echo "• ✅ All system data copied (root, home, boot)"
    echo "• ✅ fstab updated with LVM UUIDs"  
    echo "• ✅ GRUB bootloader installed on external M.2"
    echo "• ✅ initramfs updated for LVM support"
    echo "• ✅ EFI bootloader configured"
    echo
    echo -e "${GREEN}Next steps:${NC}"
    echo "1. 🔌 Keep the external M.2 drive connected"
    echo "2. 🔄 Reboot your system"
    echo "3. ⚡ Enter BIOS/UEFI boot menu (F12/F8/ESC during startup)"
    echo "4. 🎯 Select the external USB/M.2 drive to boot"
    echo "5. 🚀 Your system should boot normally with LVM!"
    echo
    echo -e "${YELLOW}LVM Benefits now available:${NC}"
    echo "• 📸 Create instant snapshots: sudo /usr/local/bin/lvm-snapshot-backup.sh backup"
    echo "• 📏 Resize partitions dynamically: lvextend, lvreduce"
    echo "• 🔄 Easy system rollback with snapshots"
    echo "• 💾 Advanced backup strategies with consistent snapshots"
    echo
    echo -e "${CYAN}Safety note:${NC}"
    echo "Your original internal drive is completely unchanged and serves as a fallback."
}

# Trap to ensure cleanup on exit
trap cleanup EXIT

main "$@"