#!/bin/bash

# Snapshot-Based LVM Migration Script: External M.2 to Internal NVMe
# Creates snapshots of ALL partitions and transfers them to internal drive

set -euo pipefail

# 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
EXTERNAL_DRIVE="/dev/sda"
INTERNAL_DRIVE="/dev/nvme0n1"
SOURCE_VG="migration-vg"
TARGET_VG="internal-vg"
BACKUP_DIR="/tmp/lvm_snapshot_migration"

# Logging
LOG_FILE="/var/log/lvm_snapshot_migration.log"
exec 1> >(tee -a "$LOG_FILE")
exec 2> >(tee -a "$LOG_FILE" >&2)

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

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

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

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

# Safety checks
safety_checks() {
    log_step "Performing safety checks..."
    
    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi
    
    # Verify we're booted from external drive
    root_device=$(findmnt -n -o SOURCE /)
    if [[ "$root_device" != "/dev/mapper/migration--vg-root" ]]; then
        log_error "Not booted from external LVM! Current root: $root_device"
        exit 1
    fi
    
    # Check drives exist
    if [[ ! -b "$EXTERNAL_DRIVE" ]] || [[ ! -b "$INTERNAL_DRIVE" ]]; then
        log_error "Required drives not found"
        exit 1
    fi
    
    # Check LVM tools
    if ! command -v lvcreate &> /dev/null; then
        log_error "LVM tools not found"
        exit 1
    fi
    
    log_info "Safety checks passed"
}

# Create backup directory and metadata
create_backup_metadata() {
    log_step "Creating backup metadata..."
    mkdir -p "$BACKUP_DIR"
    
    # Backup LVM metadata
    vgcfgbackup -f "$BACKUP_DIR/vg_backup" "$SOURCE_VG"
    
    # Save current LV information
    lvs "$SOURCE_VG" > "$BACKUP_DIR/lv_info.txt"
    vgs "$SOURCE_VG" > "$BACKUP_DIR/vg_info.txt"
    pvs > "$BACKUP_DIR/pv_info.txt"
    
    # Save partition information
    lsblk > "$BACKUP_DIR/lsblk_before.txt"
    
    log_info "Backup metadata created"
}

# Prepare internal drive with identical structure
prepare_internal_drive() {
    log_step "Preparing internal drive..."
    
    echo -e "${RED}WARNING: This will completely wipe $INTERNAL_DRIVE${NC}"
    read -p "Type 'YES' to continue: " confirm
    if [[ "$confirm" != "YES" ]]; then
        log_error "Migration cancelled"
        exit 1
    fi
    
    # Clean up any existing mounts/LVM on internal drive
    for mount in $(findmnt -n -o TARGET -S "$INTERNAL_DRIVE"* 2>/dev/null || true); do
        umount "$mount" || true
    done
    
    # Deactivate any LUKS devices on internal drive
    for luks_dev in $(lsblk -o NAME,TYPE | grep crypt | awk '{print $1}' | grep -v "migration" || true); do
        log_info "Closing LUKS device $luks_dev"
        cryptsetup close "$luks_dev" 2>/dev/null || true
    done
    
    # Deactivate any existing VGs on internal drive
    for vg in $(vgs --noheadings -o vg_name 2>/dev/null | grep -v "$SOURCE_VG" || true); do
        vgchange -an "$vg" 2>/dev/null || true
    done
    
    # Remove any PVs on internal drive
    for pv in $(pvs --noheadings -o pv_name 2>/dev/null | grep "$INTERNAL_DRIVE" || true); do
        log_info "Removing PV $pv"
        pvremove -ff "$pv" 2>/dev/null || true
    done
    
    # Wipe the drive completely
    wipefs -af "$INTERNAL_DRIVE"
    dd if=/dev/zero of="$INTERNAL_DRIVE" bs=1M count=100 2>/dev/null || true
    
    # Create identical partition structure to external drive
    log_info "Creating partition table..."
    sfdisk -d "$EXTERNAL_DRIVE" | sfdisk --force "$INTERNAL_DRIVE"
    
    # Wait for partitions to appear
    sleep 3
    partprobe "$INTERNAL_DRIVE"
    sleep 3
    
    log_info "Internal drive partitioned"
}

# Create LVM structure on internal drive
setup_internal_lvm() {
    log_step "Setting up LVM on internal drive..."
    
    # Create physical volume on the LVM partition
    pvcreate "${INTERNAL_DRIVE}p2" -ff
    
    # Create volume group with extra space for snapshots
    vgcreate "$TARGET_VG" "${INTERNAL_DRIVE}p2"
    
    # Get exact sizes from source LVs
    declare -A lv_sizes
    while IFS= read -r line; do
        lv_name=$(echo "$line" | awk '{print $1}')
        lv_size=$(echo "$line" | awk '{print $2}')
        lv_sizes["$lv_name"]="$lv_size"
    done < <(lvs --noheadings --units b --nosuffix -o lv_name,lv_size "$SOURCE_VG")
    
    # Create LVs with exact same sizes
    for lv in root home boot swap; do
        if [[ -n "${lv_sizes[$lv]:-}" ]]; then
            size_bytes="${lv_sizes[$lv]}"
            log_info "Creating LV $lv with size $size_bytes bytes"
            lvcreate -L "${size_bytes}b" -n "$lv" "$TARGET_VG"
        fi
    done
    
    log_info "Internal LVM structure created"
}

# Create temporary snapshots and transfer via dd
transfer_with_snapshots() {
    log_step "Transferring data using snapshots..."
    
    # Calculate optimal snapshot size (use available free space)
    free_space=$(vgs --noheadings --units m --nosuffix -o vg_free "$SOURCE_VG" | tr -d ' ' | tr ',' '.')
    # Use 80% of free space for snapshots, divided by number of LVs
    snapshot_size=$(echo "$free_space * 0.8 / 4" | bc | cut -d. -f1)
    
    if (( snapshot_size < 100 )); then
        log_error "Not enough free space for snapshots. Need at least 400MB free."
        exit 1
    fi
    
    log_info "Using ${snapshot_size}MB for each snapshot"
    
    # Process each logical volume
    for lv in root home boot swap; do
        log_step "Processing $lv..."
        
        # Create snapshot
        log_info "Creating snapshot of $lv..."
        if ! lvcreate -L "${snapshot_size}M" -s -n "${lv}_snapshot" "$SOURCE_VG/$lv"; then
            log_error "Failed to create snapshot for $lv"
            continue
        fi
        
        # Copy data using dd with progress
        log_info "Copying $lv data to internal drive..."
        if [[ "$lv" == "swap" ]]; then
            # For swap, just copy the LV structure
            dd if="/dev/$SOURCE_VG/${lv}_snapshot" of="/dev/$TARGET_VG/$lv" bs=1M status=progress
        else
            # For filesystems, use dd for exact copy
            dd if="/dev/$SOURCE_VG/${lv}_snapshot" of="/dev/$TARGET_VG/$lv" bs=1M status=progress
        fi
        
        # Remove snapshot to free space for next one
        log_info "Removing snapshot of $lv..."
        lvremove -f "$SOURCE_VG/${lv}_snapshot"
        
        log_info "Completed transfer of $lv"
    done
    
    log_info "All data transferred successfully"
}

# Setup boot for internal drive
configure_internal_boot() {
    log_step "Configuring boot for internal drive..."
    
    # Format and mount EFI partition
    mkfs.fat -F32 "${INTERNAL_DRIVE}p1" || true
    mkdir -p /mnt/internal_efi
    mount "${INTERNAL_DRIVE}p1" /mnt/internal_efi
    
    # Mount internal filesystems
    mkdir -p /mnt/internal_root /mnt/internal_boot /mnt/internal_home
    mount "/dev/$TARGET_VG/root" /mnt/internal_root
    mount "/dev/$TARGET_VG/boot" /mnt/internal_boot
    mount "/dev/$TARGET_VG/home" /mnt/internal_home
    
    # Mount EFI in the target system
    mkdir -p /mnt/internal_root/boot/efi
    mount "${INTERNAL_DRIVE}p1" /mnt/internal_root/boot/efi
    
    # Update fstab for new VG
    log_info "Updating fstab..."
    cat > /mnt/internal_root/etc/fstab << EOF
# Internal LVM Configuration
/dev/$TARGET_VG/root     /           ext4    defaults                    0 1
/dev/$TARGET_VG/boot     /boot       ext4    defaults                    0 2
/dev/$TARGET_VG/home     /home       ext4    defaults                    0 2
/dev/$TARGET_VG/swap     none        swap    sw                          0 0
${INTERNAL_DRIVE}p1      /boot/efi   vfat    umask=0077                  0 1
EOF
    
    # Prepare chroot environment
    mount --bind /dev /mnt/internal_root/dev
    mount --bind /proc /mnt/internal_root/proc
    mount --bind /sys /mnt/internal_root/sys
    mount --bind /run /mnt/internal_root/run
    
    # Update initramfs and install GRUB
    log_info "Updating initramfs and GRUB..."
    chroot /mnt/internal_root /bin/bash -c "update-initramfs -u -k all"
    chroot /mnt/internal_root /bin/bash -c "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Internal-LVM"
    chroot /mnt/internal_root /bin/bash -c "update-grub"
    
    # Clean up mounts
    umount /mnt/internal_root/dev /mnt/internal_root/proc /mnt/internal_root/sys /mnt/internal_root/run
    umount /mnt/internal_root/boot/efi
    umount /mnt/internal_efi /mnt/internal_root /mnt/internal_boot /mnt/internal_home
    
    log_info "Boot configuration completed"
}

# Setup snapshot capabilities on internal drive
setup_snapshot_tools() {
    log_step "Setting up snapshot tools on internal drive..."
    
    # Remount for tool installation
    mount "/dev/$TARGET_VG/root" /mnt/internal_root
    
    # Create snapshot management script
    cat > /mnt/internal_root/usr/local/bin/lvm-snapshot-manager << 'EOFSCRIPT'
#!/bin/bash
# LVM Snapshot Manager for Internal Drive

VG_NAME="internal-vg"

case "$1" in
    create)
        echo "Creating LVM snapshots..."
        # Calculate available space for snapshots
        free_space=$(vgs --noheadings --units g --nosuffix -o vg_free "$VG_NAME" | tr -d ' ')
        snapshot_size=$(echo "$free_space / 4" | bc)G
        
        lvcreate -L "$snapshot_size" -s -n root_backup "$VG_NAME/root"
        lvcreate -L "$snapshot_size" -s -n home_backup "$VG_NAME/home"
        lvcreate -L "$snapshot_size" -s -n boot_backup "$VG_NAME/boot"
        echo "Snapshots created with size: $snapshot_size each"
        ;;
    remove)
        echo "Removing LVM snapshots..."
        lvremove -f "$VG_NAME/root_backup" 2>/dev/null || true
        lvremove -f "$VG_NAME/home_backup" 2>/dev/null || true  
        lvremove -f "$VG_NAME/boot_backup" 2>/dev/null || true
        echo "Snapshots removed"
        ;;
    list)
        echo "Current snapshots:"
        lvs "$VG_NAME" | grep backup || echo "No snapshots found"
        ;;
    merge)
        if [[ -z "$2" ]]; then
            echo "Usage: $0 merge <snapshot_name>"
            exit 1
        fi
        echo "Merging snapshot $2..."
        lvconvert --merge "$VG_NAME/$2"
        echo "Snapshot merge initiated (reboot required to complete)"
        ;;
    *)
        echo "Usage: $0 {create|remove|list|merge <snapshot_name>}"
        exit 1
        ;;
esac
EOFSCRIPT
    
    chmod +x /mnt/internal_root/usr/local/bin/lvm-snapshot-manager
    
    umount /mnt/internal_root
    
    log_info "Snapshot tools installed"
}

# Cleanup function
cleanup() {
    log_info "Cleaning up..."
    
    # Remove any remaining snapshots
    for snap in $(lvs --noheadings -o lv_name "$SOURCE_VG" 2>/dev/null | grep snapshot || true); do
        lvremove -f "$SOURCE_VG/$snap" 2>/dev/null || true
    done
    
    # Unmount any remaining mounts
    for mount in /mnt/internal_*; do
        umount "$mount" 2>/dev/null || true
        rmdir "$mount" 2>/dev/null || true
    done
}

# Main migration function
main() {
    log_info "Starting snapshot-based LVM migration"
    log_info "Source: $EXTERNAL_DRIVE ($SOURCE_VG) -> Target: $INTERNAL_DRIVE ($TARGET_VG)"
    
    safety_checks
    create_backup_metadata
    prepare_internal_drive
    setup_internal_lvm
    transfer_with_snapshots
    configure_internal_boot
    setup_snapshot_tools
    cleanup
    
    log_info "Migration completed successfully!"
    echo
    echo -e "${GREEN}SUCCESS!${NC} Snapshot-based LVM migration completed"
    echo
    echo "Next steps:"
    echo "1. Reboot and select internal drive in BIOS/UEFI"
    echo "2. Verify all systems working from internal LVM"
    echo "3. Test snapshot functionality: sudo lvm-snapshot-manager create"
    echo "4. External M.2 can now be used as backup drive"
    echo
    echo "The internal drive now has:"
    echo "- Complete copy of your current system"
    echo "- LVM with snapshot capabilities"
    echo "- Reserved space for future snapshots"
}

# Handle interruption
trap cleanup EXIT

# Run migration
main "$@"