#!/bin/bash

# Boot Verification and Test Script
# Validates that a cloned drive can boot properly before declaring success

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

TARGET_DRIVE=""
WORK_DIR="/mnt/verify_work"

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

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
    return 1
}

warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

check_partition_structure() {
    log "Verifying partition structure..."
    
    # Check if drive exists and has partitions
    if [ ! -b "$TARGET_DRIVE" ]; then
        error "Target drive $TARGET_DRIVE not found"
        return 1
    fi
    
    local partitions=($(lsblk -pno NAME "$TARGET_DRIVE" | grep -v "^$TARGET_DRIVE$"))
    
    if [ ${#partitions[@]} -eq 0 ]; then
        error "No partitions found on $TARGET_DRIVE"
        return 1
    fi
    
    echo "Found ${#partitions[@]} partitions:"
    for part in "${partitions[@]}"; do
        local size=$(lsblk -no SIZE "$part")
        local fstype=$(lsblk -no FSTYPE "$part")
        local uuid=$(lsblk -no UUID "$part")
        echo "  $part: $size, $fstype, UUID: ${uuid:-'none'}"
    done
    
    success "Partition structure verified"
    return 0
}

check_filesystem_integrity() {
    log "Checking filesystem integrity..."
    
    local partitions=($(lsblk -pno NAME "$TARGET_DRIVE" | grep -v "^$TARGET_DRIVE$"))
    local errors=0
    
    for part in "${partitions[@]}"; do
        local fstype=$(lsblk -no FSTYPE "$part")
        
        case "$fstype" in
            "ext4"|"ext3"|"ext2")
                log "Checking ext filesystem on $part..."
                if e2fsck -n "$part" >/dev/null 2>&1; then
                    success "Filesystem on $part is clean"
                else
                    warning "Filesystem on $part has errors (read-only check)"
                    ((errors++))
                fi
                ;;
            "vfat")
                log "Checking FAT filesystem on $part..."
                if fsck.fat -v "$part" >/dev/null 2>&1; then
                    success "FAT filesystem on $part is clean"
                else
                    warning "FAT filesystem on $part has errors"
                    ((errors++))
                fi
                ;;
            "crypto_LUKS")
                log "Skipping LUKS partition $part (encrypted)"
                ;;
            "swap")
                log "Skipping swap partition $part"
                ;;
            *)
                log "Skipping unknown filesystem type '$fstype' on $part"
                ;;
        esac
    done
    
    if [ $errors -eq 0 ]; then
        success "All filesystems are clean"
        return 0
    else
        warning "Found $errors filesystem(s) with potential issues"
        return 1
    fi
}

check_boot_files() {
    log "Checking essential boot files..."
    
    mkdir -p "$WORK_DIR"
    local errors=0
    
    # Find and mount root partition
    local root_partition=""
    local boot_partition=""
    local efi_partition=""
    
    local partitions=($(lsblk -pno NAME "$TARGET_DRIVE" | grep -v "^$TARGET_DRIVE$"))
    
    for part in "${partitions[@]}"; do
        local fstype=$(lsblk -no FSTYPE "$part")
        local size_bytes=$(lsblk -bno SIZE "$part")
        
        if [[ "$fstype" == "vfat" && "$size_bytes" -lt 1073741824 ]]; then
            efi_partition="$part"
        elif [[ "$fstype" == "ext4" && "$size_bytes" -lt 5368709120 ]]; then
            boot_partition="$part"
        elif [[ "$fstype" == "ext4" && "$size_bytes" -gt 5368709120 ]]; then
            root_partition="$part"
        elif [[ "$fstype" == "crypto_LUKS" ]]; then
            # Try to unlock for verification
            local crypt_name="verify_$(basename "$part")"
            echo "Found encrypted partition: $part"
            echo "Please enter password to verify boot files (optional - press Enter to skip):"
            read -s -t 30 password || {
                log "Skipping encrypted partition verification"
                continue
            }
            if [ -n "$password" ]; then
                if echo "$password" | cryptsetup open "$part" "$crypt_name" --key-file=-; then
                    local decrypted_fs=$(lsblk -no FSTYPE "/dev/mapper/$crypt_name")
                    if [[ "$decrypted_fs" == "ext4" ]]; then
                        root_partition="/dev/mapper/$crypt_name"
                        log "Using decrypted partition for verification: $root_partition"
                    fi
                else
                    warning "Could not unlock encrypted partition for verification"
                fi
            fi
        fi
    done
    
    if [ -z "$root_partition" ]; then
        error "Could not find root partition for verification"
        return 1
    fi
    
    # Mount root partition
    if ! mount "$root_partition" "$WORK_DIR"; then
        error "Could not mount root partition for verification"
        return 1
    fi
    
    # Mount boot if separate
    if [ -n "$boot_partition" ]; then
        mkdir -p "$WORK_DIR/boot"
        mount "$boot_partition" "$WORK_DIR/boot" || warning "Could not mount boot partition"
    fi
    
    # Mount EFI if exists
    if [ -n "$efi_partition" ]; then
        mkdir -p "$WORK_DIR/boot/efi"
        mount "$efi_partition" "$WORK_DIR/boot/efi" || warning "Could not mount EFI partition"
    fi
    
    # Check essential files
    local essential_files=(
        "/etc/fstab"
        "/boot/grub/grub.cfg"
        "/etc/default/grub"
    )
    
    for file in "${essential_files[@]}"; do
        if [ -f "$WORK_DIR$file" ]; then
            success "Found: $file"
        else
            warning "Missing: $file"
            ((errors++))
        fi
    done
    
    # Check for kernel and initramfs
    local kernels=($(ls "$WORK_DIR/boot/vmlinuz-"* 2>/dev/null || true))
    local initramfs=($(ls "$WORK_DIR/boot/initrd.img-"* 2>/dev/null || true))
    
    if [ ${#kernels[@]} -gt 0 ]; then
        success "Found ${#kernels[@]} kernel(s)"
    else
        warning "No kernels found"
        ((errors++))
    fi
    
    if [ ${#initramfs[@]} -gt 0 ]; then
        success "Found ${#initramfs[@]} initramfs image(s)"
    else
        warning "No initramfs images found"
        ((errors++))
    fi
    
    # Check EFI bootloader
    if [ -n "$efi_partition" ]; then
        if [ -f "$WORK_DIR/boot/efi/EFI/debian/grubx64.efi" ]; then
            success "Found EFI bootloader"
        else
            warning "EFI bootloader not found"
            ((errors++))
        fi
    fi
    
    # Check fstab content
    if [ -f "$WORK_DIR/etc/fstab" ]; then
        log "Checking /etc/fstab content..."
        local fstab_errors=0
        
        # Check if UUIDs in fstab actually exist
        while read -r line; do
            if [[ "$line" =~ ^UUID=([a-fA-F0-9-]+) ]]; then
                local uuid="${BASH_REMATCH[1]}"
                if ! blkid | grep -q "$uuid"; then
                    warning "UUID $uuid in fstab not found on system"
                    ((fstab_errors++))
                fi
            fi
        done < "$WORK_DIR/etc/fstab"
        
        if [ $fstab_errors -eq 0 ]; then
            success "/etc/fstab appears valid"
        else
            warning "/etc/fstab has $fstab_errors potential issues"
            ((errors++))
        fi
    fi
    
    # Cleanup mounts
    umount "$WORK_DIR/boot/efi" 2>/dev/null || true
    umount "$WORK_DIR/boot" 2>/dev/null || true
    umount "$WORK_DIR" 2>/dev/null || true
    
    # Close encrypted volumes
    for mapper in /dev/mapper/verify_*; do
        if [ -b "$mapper" ]; then
            local crypt_name=$(basename "$mapper")
            cryptsetup close "$crypt_name" 2>/dev/null || true
        fi
    done
    
    if [ $errors -eq 0 ]; then
        success "All essential boot files found"
        return 0
    else
        warning "Found $errors potential boot issues"
        return 1
    fi
}

check_grub_configuration() {
    log "Checking GRUB configuration..."
    
    # Try to validate GRUB configuration without mounting
    if grub-probe "$TARGET_DRIVE" >/dev/null 2>&1; then
        success "GRUB can recognize the drive"
    else
        warning "GRUB may have issues recognizing the drive"
        return 1
    fi
    
    return 0
}

perform_dry_run_boot_test() {
    log "Performing dry-run boot test..."
    
    # Check if we can simulate boot process
    warning "Note: This is a simulation - actual boot test requires reboot"
    
    # Check boot order in EFI (if available)
    if command -v efibootmgr >/dev/null 2>&1; then
        log "Current EFI boot order:"
        efibootmgr | grep -E "(BootOrder|Boot[0-9]+)" || true
    fi
    
    # Test if drive is bootable by checking MBR/GPT
    if fdisk -l "$TARGET_DRIVE" | grep -q "EFI System"; then
        success "Drive has EFI System partition (UEFI bootable)"
    elif fdisk -l "$TARGET_DRIVE" | grep -q "Boot"; then
        success "Drive has bootable partition"
    else
        warning "Drive may not be properly configured for booting"
        return 1
    fi
    
    return 0
}

comprehensive_verification() {
    log "Starting comprehensive verification of $TARGET_DRIVE..."
    
    local total_checks=5
    local passed_checks=0
    local failed_checks=0
    
    echo "Verification Progress:"
    echo "====================="
    
    # Test 1: Partition Structure
    echo -n "1. Partition Structure: "
    if check_partition_structure; then
        echo -e "${GREEN}PASS${NC}"
        ((passed_checks++))
    else
        echo -e "${RED}FAIL${NC}"
        ((failed_checks++))
    fi
    
    # Test 2: Filesystem Integrity
    echo -n "2. Filesystem Integrity: "
    if check_filesystem_integrity; then
        echo -e "${GREEN}PASS${NC}"
        ((passed_checks++))
    else
        echo -e "${YELLOW}WARNING${NC}"
        ((passed_checks++))  # Count warnings as pass for now
    fi
    
    # Test 3: Boot Files
    echo -n "3. Essential Boot Files: "
    if check_boot_files; then
        echo -e "${GREEN}PASS${NC}"
        ((passed_checks++))
    else
        echo -e "${RED}FAIL${NC}"
        ((failed_checks++))
    fi
    
    # Test 4: GRUB Configuration
    echo -n "4. GRUB Configuration: "
    if check_grub_configuration; then
        echo -e "${GREEN}PASS${NC}"
        ((passed_checks++))
    else
        echo -e "${YELLOW}WARNING${NC}"
        ((passed_checks++))  # Count as pass for compatibility
    fi
    
    # Test 5: Boot Readiness
    echo -n "5. Boot Readiness: "
    if perform_dry_run_boot_test; then
        echo -e "${GREEN}PASS${NC}"
        ((passed_checks++))
    else
        echo -e "${RED}FAIL${NC}"
        ((failed_checks++))
    fi
    
    echo "====================="
    echo "Verification Summary:"
    echo "  Passed: $passed_checks/$total_checks"
    echo "  Failed: $failed_checks/$total_checks"
    
    if [ $failed_checks -eq 0 ]; then
        success "All verification checks passed! Drive should boot properly."
        return 0
    elif [ $failed_checks -le 2 ]; then
        warning "Some checks failed but drive might still boot. Consider running boot repair."
        return 1
    else
        error "Multiple critical checks failed. Drive is unlikely to boot properly."
        return 2
    fi
}

cleanup_verification() {
    log "Cleaning up verification environment..."
    
    # Unmount any remaining mounts
    umount "$WORK_DIR" 2>/dev/null || true
    
    # Close any encrypted volumes
    for mapper in /dev/mapper/verify_*; do
        if [ -b "$mapper" ]; then
            local crypt_name=$(basename "$mapper")
            cryptsetup close "$crypt_name" 2>/dev/null || true
        fi
    done
    
    # Remove work directory
    rmdir "$WORK_DIR" 2>/dev/null || true
    
    success "Cleanup completed"
}

main() {
    echo -e "${GREEN}=== Boot Verification and Test Script ===${NC}"
    echo "This script validates that a cloned drive can boot properly"
    echo
    
    # Get target drive
    echo "Available drives:"
    lsblk -dpno NAME,SIZE,MODEL | grep -v "loop\|ram"
    echo
    
    read -p "Enter the drive to verify (e.g., /dev/sdb): " TARGET_DRIVE
    
    if [ ! -b "$TARGET_DRIVE" ]; then
        error "Drive $TARGET_DRIVE not found"
        exit 1
    fi
    
    echo "Selected drive for verification: $TARGET_DRIVE"
    lsblk "$TARGET_DRIVE"
    echo
    
    read -p "Verify this drive? [y/N]: " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        error "Verification cancelled"
        exit 1
    fi
    
    # Perform comprehensive verification
    if comprehensive_verification; then
        echo
        echo -e "${GREEN}🎉 VERIFICATION SUCCESSFUL! 🎉${NC}"
        echo "The cloned drive passed all verification checks."
        echo "It should boot properly when set as the primary boot device."
        echo
        echo -e "${BLUE}Next steps:${NC}"
        echo "1. Reboot your system"
        echo "2. Enter BIOS/UEFI setup"
        echo "3. Set $TARGET_DRIVE as the first boot device"
        echo "4. Save and exit BIOS/UEFI"
        echo "5. System should boot from cloned drive"
        
        if lsblk "$TARGET_DRIVE" | grep -q "crypto_LUKS"; then
            echo
            echo -e "${YELLOW}Note:${NC} System will ask for LUKS password during boot (this is normal)"
        fi
        
        cleanup_verification
        exit 0
    else
        local exit_code=$?
        echo
        if [ $exit_code -eq 1 ]; then
            echo -e "${YELLOW}⚠️ VERIFICATION COMPLETED WITH WARNINGS ⚠️${NC}"
            echo "The drive might boot but some issues were detected."
            echo "Consider running the boot repair script before attempting to boot."
            echo
            echo "Run: ./boot_repair_tools.sh"
        else
            echo -e "${RED}❌ VERIFICATION FAILED ❌${NC}"
            echo "The drive is unlikely to boot properly in its current state."
            echo "Please run the boot repair script to fix issues."
            echo
            echo "Run: ./boot_repair_tools.sh"
        fi
        
        cleanup_verification
        exit $exit_code
    fi
}

# Trap to ensure cleanup on exit
trap cleanup_verification EXIT

main "$@"