#!/bin/bash

# Simplified LUKS Setup Script
# Wipes internal home, creates LUKS encryption, and restores from external 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'

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"
}

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

echo -e "${BLUE}=== Simple LUKS Encryption Setup ===${NC}"
echo
echo "This will:"
echo "1. Remove the current home LV on internal drive"
echo "2. Create a new LUKS-encrypted home LV"
echo "3. Copy your home data directly from external M.2"
echo "4. Update system configuration"
echo
echo -e "${YELLOW}Source:${NC} External M.2 (/dev/migration-vg/home)"
echo -e "${YELLOW}Target:${NC} Internal NVMe (/dev/internal-vg/home) - WILL BE WIPED"
echo

read -p "Continue with LUKS encryption setup? (yes/no): " confirm
if [[ "$confirm" != "yes" ]]; then
    log_info "Operation cancelled"
    exit 0
fi

MOUNT_POINT="/mnt/luks_setup"
EXTERNAL_HOME="/dev/migration-vg/home"
INTERNAL_VG="internal-vg"

log_step "Removing current internal home LV..."

# Remove the current home LV
lvremove -f "$INTERNAL_VG/home"

log_step "Creating new home LV..."

# Get the original home size from external drive
home_size=$(lvs --noheadings --units g --nosuffix -o lv_size migration-vg/home | tr -d ' ' | tr ',' '.')

# Create new home LV
lvcreate -L "${home_size}G" -n home "$INTERNAL_VG"

log_step "Setting up LUKS encryption..."

# Setup LUKS on the new LV
echo "Please enter your desired LUKS passphrase for home encryption:"
cryptsetup luksFormat "/dev/$INTERNAL_VG/home"

echo "Please enter your LUKS passphrase again to open the volume:"
cryptsetup open "/dev/$INTERNAL_VG/home" luks-home-internal

# Format the encrypted volume
mkfs.ext4 -L home /dev/mapper/luks-home-internal

log_step "Copying home data from external drive..."

# Mount source and target
mkdir -p "$MOUNT_POINT/external" "$MOUNT_POINT/encrypted"
mount "$EXTERNAL_HOME" "$MOUNT_POINT/external"
mount /dev/mapper/luks-home-internal "$MOUNT_POINT/encrypted"

# Copy data directly from external to encrypted volume
log_info "Copying ${home_size}GB of home data..."
rsync -avHAXS --progress "$MOUNT_POINT/external/" "$MOUNT_POINT/encrypted/"

# Clean up mounts
umount "$MOUNT_POINT/external" "$MOUNT_POINT/encrypted"
cryptsetup close luks-home-internal

log_step "Updating system configuration..."

# Get the UUID of the LUKS device
LUKS_UUID=$(cryptsetup luksUUID "/dev/$INTERNAL_VG/home")

# Mount the internal root to update configuration
mount "/dev/$INTERNAL_VG/root" "$MOUNT_POINT"

# Update /etc/crypttab
echo "luks-home-internal UUID=$LUKS_UUID none luks" >> "$MOUNT_POINT/etc/crypttab"

# Update /etc/fstab
cat > "$MOUNT_POINT/etc/fstab" << EOF
# Internal LVM Configuration with LUKS
/dev/$INTERNAL_VG/root        /           ext4    defaults                    0 1
/dev/$INTERNAL_VG/boot        /boot       ext4    defaults                    0 2
/dev/mapper/luks-home-internal /home      ext4    defaults                    0 2
/dev/$INTERNAL_VG/swap        none        swap    sw                          0 0
/dev/nvme0n1p1                /boot/efi   vfat    umask=0077                  0 1
EOF

# Update initramfs and GRUB to include LUKS support
mount --bind /dev "$MOUNT_POINT/dev"
mount --bind /proc "$MOUNT_POINT/proc"
mount --bind /sys "$MOUNT_POINT/sys"
mount --bind /run "$MOUNT_POINT/run"

log_info "Updating initramfs for LUKS support..."
chroot "$MOUNT_POINT" /bin/bash -c "update-initramfs -u -k all"

log_info "Updating GRUB configuration..."
chroot "$MOUNT_POINT" /bin/bash -c "update-grub"

# Clean up
umount "$MOUNT_POINT/dev" "$MOUNT_POINT/proc" "$MOUNT_POINT/sys" "$MOUNT_POINT/run"
umount "$MOUNT_POINT"
rmdir "$MOUNT_POINT/external" "$MOUNT_POINT/encrypted" "$MOUNT_POINT" 2>/dev/null || true

log_info "LUKS encryption setup completed successfully!"
echo
echo -e "${GREEN}SUCCESS!${NC} Home partition is now encrypted with LUKS"
echo
echo "Configuration summary:"
echo "• LUKS UUID: $LUKS_UUID"
echo "• Encrypted device: /dev/mapper/luks-home-internal"
echo "• Mount point: /home"
echo "• Data copied from external M.2"
echo
echo "Next steps:"
echo "1. Reboot and select internal NVMe drive in BIOS"
echo "2. You will be prompted for LUKS passphrase during boot"
echo "3. Verify that all your home data is accessible"
echo
echo -e "${YELLOW}Important:${NC} Remember your LUKS passphrase! Without it, your home data will be inaccessible."