#!/bin/bash

# LVM Migration Validation Script
# Comprehensive validation that the migration from non-LVM to LVM was successful

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
VG_NAME="system-vg"
EXTERNAL_DRIVE="/dev/sda"
EXPECTED_LVS=("root" "home" "boot" "swap")
VALIDATION_LOG="/var/log/lvm-migration-validation.log"

log() {
    local message="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "${BLUE}$message${NC}"
    echo "$message" >> "$VALIDATION_LOG" 2>/dev/null || true
}

error() {
    local message="[ERROR] $1"
    echo -e "${RED}$message${NC}" >&2
    echo "$message" >> "$VALIDATION_LOG" 2>/dev/null || true
}

warning() {
    local message="[WARNING] $1"
    echo -e "${YELLOW}$message${NC}"
    echo "$message" >> "$VALIDATION_LOG" 2>/dev/null || true
}

success() {
    local message="[SUCCESS] $1"
    echo -e "${GREEN}$message${NC}"
    echo "$message" >> "$VALIDATION_LOG" 2>/dev/null || true
}

check_lvm_volumes() {
    log "Checking LVM volumes..."
    
    if ! vgs "$VG_NAME" >/dev/null 2>&1; then
        error "Volume group $VG_NAME not found"
        return 1
    fi
    
    local expected_lvs=("root" "home" "swap" "boot")
    for lv in "${expected_lvs[@]}"; do
        if ! lvs "$VG_NAME/$lv" >/dev/null 2>&1; then
            error "Logical volume $VG_NAME/$lv not found"
            return 1
        else
            success "Found logical volume: $VG_NAME/$lv"
        fi
    done
}

check_filesystems() {
    log "Checking filesystems..."
    
    local volumes=("root" "home" "boot")
    for vol in "${volumes[@]}"; do
        local device="/dev/$VG_NAME/$vol"
        if fsck.ext4 -n "$device" >/dev/null 2>&1; then
            success "Filesystem on $device is clean"
        else
            error "Filesystem on $device has errors"
            return 1
        fi
    done
    
    # Check EFI partition
    if fsck.fat -v "${EXTERNAL_DRIVE}1" >/dev/null 2>&1; then
        success "EFI filesystem is clean"
    else
        warning "EFI filesystem check failed (this might be normal)"
    fi
}

check_boot_files() {
    log "Checking boot files..."
    
    local mount_point="/mnt/validation_check"
    mkdir -p "$mount_point"
    
    # Check root partition for essential directories
    mount "/dev/$VG_NAME/root" "$mount_point"
    
    local essential_dirs=("/bin" "/sbin" "/etc" "/usr" "/var")
    for dir in "${essential_dirs[@]}"; do
        if [ -d "$mount_point$dir" ]; then
            success "Found essential directory: $dir"
        else
            error "Missing essential directory: $dir"
            umount "$mount_point"
            return 1
        fi
    done
    
    # Check for key system files
    local key_files=("/etc/fstab" "/etc/passwd" "/etc/group")
    for file in "${key_files[@]}"; do
        if [ -f "$mount_point$file" ]; then
            success "Found key system file: $file"
        else
            error "Missing key system file: $file"
            umount "$mount_point"
            return 1
        fi
    done
    
    umount "$mount_point"
}

check_grub_installation() {
    log "Checking GRUB installation..."
    
    local efi_mount="/mnt/efi_check"
    mkdir -p "$efi_mount"
    mount "${EXTERNAL_DRIVE}1" "$efi_mount"
    
    if [ -d "$efi_mount/EFI/debian" ]; then
        success "GRUB EFI installation found"
    else
        error "GRUB EFI installation not found"
        umount "$efi_mount"
        return 1
    fi
    
    if [ -f "$efi_mount/EFI/debian/grubx64.efi" ]; then
        success "GRUB EFI binary found"
    else
        error "GRUB EFI binary not found"
        umount "$efi_mount"
        return 1
    fi
    
    umount "$efi_mount"
}

check_fstab() {
    log "Checking /etc/fstab configuration..."
    
    local mount_point="/mnt/fstab_check"
    mkdir -p "$mount_point"
    mount "/dev/$VG_NAME/root" "$mount_point"
    
    if [ -f "$mount_point/etc/fstab" ]; then
        success "Found /etc/fstab"
        
        # Check if fstab contains LVM UUIDs
        local root_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/root")
        local home_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/home")
        local boot_uuid=$(blkid -s UUID -o value "/dev/$VG_NAME/boot")
        
        if grep -q "$root_uuid" "$mount_point/etc/fstab"; then
            success "Root UUID found in fstab"
        else
            error "Root UUID not found in fstab"
        fi
        
        if grep -q "$home_uuid" "$mount_point/etc/fstab"; then
            success "Home UUID found in fstab"
        else
            error "Home UUID not found in fstab"
        fi
        
        if grep -q "$boot_uuid" "$mount_point/etc/fstab"; then
            success "Boot UUID found in fstab"
        else
            error "Boot UUID not found in fstab"
        fi
    else
        error "/etc/fstab not found"
    fi
    
    umount "$mount_point"
}

check_snapshot_capability() {
    log "Checking LVM snapshot capability..."
    
    # Check free space for snapshots
    local vg_free=$(vgs --noheadings -o vg_free --units g "$VG_NAME" | tr -d ' G')
    local vg_free_int=${vg_free%.*}
    
    if [ "$vg_free_int" -ge 20 ]; then
        success "Sufficient free space for snapshots: ${vg_free}G available"
    else
        warning "Limited free space for snapshots: ${vg_free}G available (recommend 20G+)"
    fi
    
    # Test snapshot creation and removal
    log "Testing snapshot creation..."
    if lvcreate -L 1G -s -n test-snapshot "/dev/$VG_NAME/root" >/dev/null 2>&1; then
        success "Snapshot creation test successful"
        if lvremove -f "/dev/$VG_NAME/test-snapshot" >/dev/null 2>&1; then
            success "Snapshot removal test successful"
        else
            error "Snapshot removal test failed"
        fi
    else
        error "Snapshot creation test failed"
        return 1
    fi
}

check_lvm_tools() {
    log "Checking for LVM snapshot script..."
    
    local mount_point="/mnt/script_check"
    mkdir -p "$mount_point"
    mount "/dev/$VG_NAME/root" "$mount_point"
    
    if [ -f "$mount_point/usr/local/bin/lvm-snapshot-backup.sh" ]; then
        success "LVM snapshot backup script found"
        if [ -x "$mount_point/usr/local/bin/lvm-snapshot-backup.sh" ]; then
            success "LVM snapshot backup script is executable"
        else
            error "LVM snapshot backup script is not executable"
        fi
    else
        error "LVM snapshot backup script not found"
    fi
    
    umount "$mount_point"
}

main() {
    echo -e "${GREEN}=== LVM Migration Validation ===${NC}"
    echo "Validating the migrated LVM system..."
    echo
    
    local failed_checks=0
    
    check_lvm_volumes || ((failed_checks++))
    check_filesystems || ((failed_checks++))
    check_boot_files || ((failed_checks++))
    check_grub_installation || ((failed_checks++))
    check_fstab || ((failed_checks++))
    check_snapshot_capability || ((failed_checks++))
    check_lvm_tools || ((failed_checks++))
    
    echo
    if [ $failed_checks -eq 0 ]; then
        success "All validation checks passed!"
        echo -e "${GREEN}The migration appears to be successful.${NC}"
        echo "You can now:"
        echo "1. Update BIOS/UEFI boot order to boot from external M.2"
        echo "2. Test booting from the external drive"
        echo "3. Use 'lvm-snapshot-backup.sh backup' for backups"
    else
        error "$failed_checks validation check(s) failed!"
        echo -e "${RED}The migration may have issues. Review the errors above.${NC}"
        return 1
    fi
}

main "$@"