#!/bin/bash

# Resume snapshot migration - continue from data transfer phase
# Since internal LVM structure already exists

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
SOURCE_VG="migration-vg"
TARGET_VG="internal-vg"
INTERNAL_DRIVE="/dev/nvme0n1"

# Logging
LOG_FILE="/var/log/lvm_snapshot_migration_resume.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_step() {
    echo -e "${BLUE}[STEP]${NC} $(date '+%Y-%m-%d %H:%M:%S'): $1"
}

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

# Transfer data with snapshots (resume from where we left off)
transfer_with_snapshots() {
    log_step "Resuming data transfer 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 boot home 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 tools
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 function
main() {
    log_info "Resuming snapshot-based LVM migration"
    log_info "Internal LVM structure already exists, continuing with data transfer..."
    
    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"
}

# Handle interruption
trap cleanup EXIT

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root (use sudo)"
    exit 1
fi

# Run migration
main "$@"