#!/bin/bash

# LVM Migration Script: External M.2 to Internal NVMe
# This script migrates the complete LVM structure from external M.2 to internal NVMe
# with snapshot capabilities for future backups

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"
VG_NAME="migration-vg"
NEW_VG_NAME="internal-vg"  # New VG name for internal drive
BACKUP_DIR="/tmp/lvm_migration_backup"

# Logging
LOG_FILE="/var/log/lvm_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"
        log_error "Please boot from external M.2 drive first"
        exit 1
    fi
    
    # Check drives exist
    if [[ ! -b "$EXTERNAL_DRIVE" ]]; then
        log_error "External drive $EXTERNAL_DRIVE not found"
        exit 1
    fi
    
    if [[ ! -b "$INTERNAL_DRIVE" ]]; then
        log_error "Internal drive $INTERNAL_DRIVE not found"
        exit 1
    fi
    
    # Check available space
    vg_size=$(vgs --noheadings --units g --nosuffix -o vg_size "$VG_NAME" | tr -d ' ')
    internal_size=$(lsblk -b -n -o SIZE "$INTERNAL_DRIVE" | awk '{printf "%.0f", $1/1024/1024/1024}')
    
    if (( $(echo "$vg_size > $internal_size" | bc -l) )); then
        log_error "Internal drive ($internal_size GB) is smaller than LVM structure ($vg_size GB)"
        exit 1
    fi
    
    log_info "Safety checks passed"
}

# Create backup directory
create_backup_dir() {
    log_step "Creating backup directory..."
    mkdir -p "$BACKUP_DIR"
    
    # Backup current LVM configuration
    vgcfgbackup -f "$BACKUP_DIR/vg_backup" "$VG_NAME"
    
    # Save partition tables
    sfdisk -d "$EXTERNAL_DRIVE" > "$BACKUP_DIR/external_partition_table.txt"
    sfdisk -d "$INTERNAL_DRIVE" > "$BACKUP_DIR/internal_partition_table.txt" 2>/dev/null || true
    
    log_info "Backup directory created at $BACKUP_DIR"
}

# Wipe internal drive and create new partition structure
prepare_internal_drive() {
    log_step "Preparing internal drive..."
    
    # Confirmation
    echo -e "${RED}WARNING: This will completely wipe $INTERNAL_DRIVE${NC}"
    echo "Current internal drive content will be lost!"
    read -p "Type 'YES' to continue: " confirm
    if [[ "$confirm" != "YES" ]]; then
        log_error "Migration cancelled by user"
        exit 1
    fi
    
    # Unmount any mounted partitions from internal drive
    for mount in $(findmnt -n -o TARGET -S "$INTERNAL_DRIVE"* 2>/dev/null || true); do
        log_info "Unmounting $mount"
        umount "$mount" || true
    done
    
    # Deactivate any LVM on internal drive
    for pv in $(pvs --noheadings -o pv_name | grep -E "^[[:space:]]*$INTERNAL_DRIVE" || true); do
        log_info "Deactivating PV $pv"
        vgchange -an $(pvs --noheadings -o vg_name "$pv") || true
    done
    
    # Wipe filesystem signatures and partition table
    wipefs -af "$INTERNAL_DRIVE"
    
    # Create new GPT partition table
    parted "$INTERNAL_DRIVE" --script mklabel gpt
    
    # Create EFI boot partition (512MB)
    parted "$INTERNAL_DRIVE" --script mkpart primary fat32 1MiB 513MiB
    parted "$INTERNAL_DRIVE" --script set 1 esp on
    
    # Create LVM partition (rest of the drive)
    parted "$INTERNAL_DRIVE" --script mkpart primary 513MiB 100%
    parted "$INTERNAL_DRIVE" --script set 2 lvm on
    
    # Wait for partitions to be recognized
    sleep 2
    partprobe "$INTERNAL_DRIVE"
    sleep 2
    
    # Format EFI partition
    mkfs.fat -F32 "${INTERNAL_DRIVE}p1"
    
    log_info "Internal drive prepared with new partition structure"
}

# Create LVM structure on internal drive
create_lvm_structure() {
    log_step "Creating LVM structure on internal drive..."
    
    # Create physical volume
    pvcreate "${INTERNAL_DRIVE}p2"
    
    # Create volume group
    vgcreate "$NEW_VG_NAME" "${INTERNAL_DRIVE}p2"
    
    # Get sizes from external LVM
    root_size=$(lvs --noheadings --units g --nosuffix -o lv_size "$VG_NAME/root" | tr -d ' ')
    home_size=$(lvs --noheadings --units g --nosuffix -o lv_size "$VG_NAME/home" | tr -d ' ')
    boot_size=$(lvs --noheadings --units g --nosuffix -o lv_size "$VG_NAME/boot" | tr -d ' ')
    swap_size=$(lvs --noheadings --units g --nosuffix -o lv_size "$VG_NAME/swap" | tr -d ' ')
    
    # Reserve 20% free space for snapshots
    vg_free=$(vgs --noheadings --units g --nosuffix -o vg_free "$NEW_VG_NAME" | tr -d ' ')
    snapshot_reserve=$(echo "$vg_free * 0.2" | bc -l | cut -d. -f1)
    
    log_info "Reserving ${snapshot_reserve}GB for LVM snapshots"
    
    # Create logical volumes with same sizes as external
    lvcreate -L "${root_size}G" -n root "$NEW_VG_NAME"
    lvcreate -L "${boot_size}G" -n boot "$NEW_VG_NAME"
    lvcreate -L "${swap_size}G" -n swap "$NEW_VG_NAME"
    
    # Create home with remaining space minus snapshot reserve
    remaining_space=$(echo "$vg_free - $snapshot_reserve" | bc -l | cut -d. -f1)
    adjusted_home_size=$(echo "$remaining_space - $root_size - $boot_size - $swap_size" | bc -l | cut -d. -f1)
    
    if (( $(echo "$adjusted_home_size > 0" | bc -l) )); then
        lvcreate -L "${adjusted_home_size}G" -n home "$NEW_VG_NAME"
    else
        log_warn "Using original home size, may reduce snapshot space"
        lvcreate -L "${home_size}G" -n home "$NEW_VG_NAME"
    fi
    
    # Create filesystems
    mkfs.ext4 -L root "/dev/$NEW_VG_NAME/root"
    mkfs.ext4 -L boot "/dev/$NEW_VG_NAME/boot"
    mkfs.ext4 -L home "/dev/$NEW_VG_NAME/home"
    mkswap -L swap "/dev/$NEW_VG_NAME/swap"
    
    log_info "LVM structure created on internal drive"
}

# Mount internal drive and copy data
copy_data() {
    log_step "Copying data from external to internal drive..."
    
    # Create mount points
    mkdir -p /mnt/internal_root
    mkdir -p /mnt/internal_boot
    mkdir -p /mnt/internal_home
    
    # Mount internal LVM
    mount "/dev/$NEW_VG_NAME/root" /mnt/internal_root
    mount "/dev/$NEW_VG_NAME/boot" /mnt/internal_boot
    mount "/dev/$NEW_VG_NAME/home" /mnt/internal_home
    
    # Copy data with progress
    log_info "Copying root filesystem..."
    rsync -avHAXS --progress / /mnt/internal_root/ \
        --exclude=/dev/* \
        --exclude=/proc/* \
        --exclude=/sys/* \
        --exclude=/tmp/* \
        --exclude=/run/* \
        --exclude=/mnt/* \
        --exclude=/media/* \
        --exclude=/lost+found \
        --exclude=/boot/*
    
    log_info "Copying boot filesystem..."
    rsync -avHAXS --progress /boot/ /mnt/internal_boot/
    
    log_info "Copying home filesystem..."
    rsync -avHAXS --progress /home/ /mnt/internal_home/
    
    log_info "Data copy completed"
}

# Configure system for internal LVM boot
configure_boot() {
    log_step "Configuring boot for internal LVM..."
    
    # Mount EFI partition
    mkdir -p /mnt/internal_root/boot/efi
    mount "${INTERNAL_DRIVE}p1" /mnt/internal_root/boot/efi
    
    # Chroot preparations
    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 fstab
    cat > /mnt/internal_root/etc/fstab << EOF
# Internal LVM Configuration
/dev/$NEW_VG_NAME/root     /           ext4    defaults                    0 1
/dev/$NEW_VG_NAME/boot     /boot       ext4    defaults                    0 2
/dev/$NEW_VG_NAME/home     /home       ext4    defaults                    0 2
/dev/$NEW_VG_NAME/swap     none        swap    sw                          0 0
${INTERNAL_DRIVE}p1        /boot/efi   vfat    umask=0077                  0 1
EOF
    
    # Update initramfs to include LVM
    chroot /mnt/internal_root /bin/bash -c "update-initramfs -u -k all"
    
    # Install and configure GRUB
    chroot /mnt/internal_root /bin/bash -c "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian"
    chroot /mnt/internal_root /bin/bash -c "update-grub"
    
    # Cleanup mounts
    umount /mnt/internal_root/dev
    umount /mnt/internal_root/proc
    umount /mnt/internal_root/sys
    umount /mnt/internal_root/run
    umount /mnt/internal_root/boot/efi
    
    log_info "Boot configuration completed"
}

# Setup LVM snapshots
setup_snapshots() {
    log_step "Setting up LVM snapshot capability..."
    
    # Create snapshot management script
    cat > /mnt/internal_root/usr/local/bin/lvm-snapshot-manager << 'EOF'
#!/bin/bash
# LVM Snapshot Manager

VG_NAME="internal-vg"
SNAPSHOT_SIZE="10G"  # Adjust based on available space

case "$1" in
    create)
        echo "Creating LVM snapshots..."
        lvcreate -L $SNAPSHOT_SIZE -s -n root_snapshot $VG_NAME/root
        lvcreate -L $SNAPSHOT_SIZE -s -n home_snapshot $VG_NAME/home
        echo "Snapshots created successfully"
        ;;
    remove)
        echo "Removing LVM snapshots..."
        lvremove -f $VG_NAME/root_snapshot 2>/dev/null || true
        lvremove -f $VG_NAME/home_snapshot 2>/dev/null || true
        echo "Snapshots removed"
        ;;
    list)
        echo "Current snapshots:"
        lvs $VG_NAME | grep snapshot || 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 (requires reboot to complete)"
        ;;
    *)
        echo "Usage: $0 {create|remove|list|merge <snapshot_name>}"
        echo "  create - Create snapshots of root and home"
        echo "  remove - Remove all snapshots"
        echo "  list   - List existing snapshots"
        echo "  merge  - Merge a snapshot back to origin"
        exit 1
        ;;
esac
EOF
    
    chmod +x /mnt/internal_root/usr/local/bin/lvm-snapshot-manager
    
    # Create backup script using snapshots
    cat > /mnt/internal_root/usr/local/bin/snapshot-backup << 'EOF'
#!/bin/bash
# Snapshot-based backup script

BACKUP_DIR="/backup"
VG_NAME="internal-vg"
DATE=$(date +%Y%m%d_%H%M%S)

# Create snapshots
echo "Creating snapshots..."
/usr/local/bin/lvm-snapshot-manager create

# Mount snapshots and backup
mkdir -p $BACKUP_DIR/root_$DATE
mkdir -p $BACKUP_DIR/home_$DATE

mount /dev/$VG_NAME/root_snapshot $BACKUP_DIR/root_$DATE
mount /dev/$VG_NAME/home_snapshot $BACKUP_DIR/home_$DATE

echo "Snapshots mounted. You can now backup from:"
echo "  Root: $BACKUP_DIR/root_$DATE"
echo "  Home: $BACKUP_DIR/home_$DATE"
echo ""
echo "When done, run: umount $BACKUP_DIR/root_$DATE $BACKUP_DIR/home_$DATE"
echo "Then run: /usr/local/bin/lvm-snapshot-manager remove"
EOF
    
    chmod +x /mnt/internal_root/usr/local/bin/snapshot-backup
    
    log_info "LVM snapshot tools installed"
}

# Cleanup and unmount
cleanup() {
    log_step "Cleaning up..."
    
    # Unmount all internal mounts
    umount /mnt/internal_home || true
    umount /mnt/internal_boot || true
    umount /mnt/internal_root || true
    
    # Remove mount points
    rmdir /mnt/internal_home /mnt/internal_boot /mnt/internal_root 2>/dev/null || true
    
    log_info "Cleanup completed"
}

# Main migration function
main() {
    log_info "Starting LVM migration from external M.2 to internal NVMe"
    log_info "External: $EXTERNAL_DRIVE -> Internal: $INTERNAL_DRIVE"
    
    safety_checks
    create_backup_dir
    prepare_internal_drive
    create_lvm_structure
    copy_data
    configure_boot
    setup_snapshots
    cleanup
    
    log_info "Migration completed successfully!"
    echo
    echo -e "${GREEN}SUCCESS!${NC} LVM migration completed"
    echo "Next steps:"
    echo "1. Reboot and select internal drive in BIOS/UEFI"
    echo "2. Verify system boots correctly from internal LVM"
    echo "3. Test snapshot functionality with: sudo lvm-snapshot-manager create"
    echo "4. Once confirmed working, you can repurpose the external M.2"
    echo
    echo "Snapshot management commands:"
    echo "  sudo lvm-snapshot-manager create    # Create snapshots"
    echo "  sudo lvm-snapshot-manager list      # List snapshots"
    echo "  sudo lvm-snapshot-manager remove    # Remove snapshots"
    echo "  sudo snapshot-backup                # Backup using snapshots"
}

# Handle interruption
trap cleanup EXIT

# Run main function
main "$@"