#!/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="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")
    local internal_size_bytes=$(lsblk -bno SIZE "$INTERNAL_DRIVE")
    
    if [ "$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..."
    
    local partitions=($(lsblk -pno NAME "$INTERNAL_DRIVE" | grep -v "^$INTERNAL_DRIVE$"))
    
    echo "Source drive partitions:"
    for part in "${partitions[@]}"; do
        local size=$(lsblk -no SIZE "$part")
        local fstype=$(lsblk -no FSTYPE "$part")
        local label=$(lsblk -no LABEL "$part")
        local mountpoint=$(lsblk -no MOUNTPOINT "$part")
        
        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"
}

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..."
    
    # Calculate sizes based on source
    local root_size="70G"
    local home_size="100G"  
    local swap_size="8G"
    local boot_size="2G"
    
    # Adjust based on available space
    local total_space_gb=$(lsblk -bno SIZE "$EXTERNAL_DRIVE" | awk '{print int($1/1024/1024/1024)}')
    if [ "$total_space_gb" -gt 500 ]; then
        home_size="200G"
        root_size="100G"
    fi
    
    log "LVM sizes: Root=$root_size, Home=$home_size, Swap=$swap_size, Boot=$boot_size"
    
    # Wipe and create partition table
    wipefs -a "$EXTERNAL_DRIVE"
    parted -s "$EXTERNAL_DRIVE" mklabel gpt
    
    # Create EFI partition (512MB)
    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)
    parted -s "$EXTERNAL_DRIVE" mkpart primary 513MiB 100%
    parted -s "$EXTERNAL_DRIVE" set 2 lvm on
    
    # Wait for partitions
    sleep 3
    partprobe "$EXTERNAL_DRIVE"
    sleep 2
    
    # Verify partitions exist
    if [ ! -b "${EXTERNAL_DRIVE}1" ] || [ ! -b "${EXTERNAL_DRIVE}2" ]; then
        error "Partitions not created properly"
    fi
    
    # Create filesystems
    mkfs.fat -F32 "${EXTERNAL_DRIVE}1" || error "Failed to create EFI filesystem"
    
    # Setup LVM
    pvcreate "${EXTERNAL_DRIVE}2" || error "Failed to create physical volume"
    vgcreate "$VG_NAME" "${EXTERNAL_DRIVE}2" || error "Failed to create volume group"
    
    # 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
    umount "$WORK_DIR/external_root/boot/efi" 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
    
    # Close encrypted partitions
    for mapper in /dev/mapper/migration_*; do
        if [ -b "$mapper" ]; then
            cryptsetup close "$(basename "$mapper")" 2>/dev/null || true
        fi
    done
    
    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
    
    echo
    echo "Migration Summary:"
    echo "  Source: $INTERNAL_DRIVE (current system)"
    echo "  Target: $EXTERNAL_DRIVE (new LVM system)"
    echo "  VG Name: $VG_NAME"
    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 "$@"