#!/bin/bash

# Boot Repair and Emergency Recovery Tools
# Provides tools to fix boot issues after cloning

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/repair_work"

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

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

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

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

detect_target_drive() {
    log "Detecting target drive to repair..."
    
    # List available drives
    echo "Available drives:"
    lsblk -dpno NAME,SIZE,MODEL | grep -v "loop\|ram"
    echo
    
    read -p "Enter the drive to repair (e.g., /dev/sdb): " TARGET_DRIVE
    
    if [ ! -b "$TARGET_DRIVE" ]; then
        error "Drive $TARGET_DRIVE not found"
    fi
    
    echo "Selected drive for repair: $TARGET_DRIVE"
    lsblk "$TARGET_DRIVE"
    echo
    
    read -p "Is this correct? [y/N]: " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        error "Operation cancelled"
    fi
    
    success "Target drive selected: $TARGET_DRIVE"
}

mount_target_system() {
    log "Mounting target system for repair..."
    
    mkdir -p "$WORK_DIR"
    
    # Find partitions
    local partitions=($(lsblk -pno NAME "$TARGET_DRIVE" | grep -v "^$TARGET_DRIVE$"))
    local root_partition=""
    local boot_partition=""
    local efi_partition=""
    
    for part in "${partitions[@]}"; do
        local fstype=$(lsblk -no FSTYPE "$part")
        local size_bytes=$(lsblk -bno SIZE "$part")
        
        # Detect partitions by filesystem and size
        if [[ "$fstype" == "vfat" && "$size_bytes" -lt 1073741824 ]]; then  # Less than 1GB
            efi_partition="$part"
        elif [[ "$fstype" == "ext4" && "$size_bytes" -lt 5368709120 ]]; then  # Between 100MB and 5GB
            boot_partition="$part"
        elif [[ "$fstype" == "ext4" && "$size_bytes" -gt 5368709120 ]]; then  # Larger than 5GB
            root_partition="$part"
        elif [[ "$fstype" == "crypto_LUKS" ]]; then
            # Try to unlock encrypted partition
            local crypt_name="repair_$(basename "$part")"
            echo "Found encrypted partition: $part"
            echo "Please enter the password to unlock for repair:"
            if cryptsetup open "$part" "$crypt_name"; 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 as root: $root_partition"
                fi
            else
                warning "Could not unlock encrypted partition"
            fi
        fi
    done
    
    if [ -z "$root_partition" ]; then
        error "Could not find root partition to repair"
    fi
    
    # Mount filesystems
    log "Mounting root partition: $root_partition"
    mount "$root_partition" "$WORK_DIR" || error "Failed to mount root partition"
    
    if [ -n "$boot_partition" ]; then
        log "Mounting boot partition: $boot_partition"
        mkdir -p "$WORK_DIR/boot"
        mount "$boot_partition" "$WORK_DIR/boot" || warning "Failed to mount boot partition"
    fi
    
    if [ -n "$efi_partition" ]; then
        log "Mounting EFI partition: $efi_partition"
        mkdir -p "$WORK_DIR/boot/efi"
        mount "$efi_partition" "$WORK_DIR/boot/efi" || warning "Failed to mount EFI partition"
    fi
    
    success "Target system mounted at $WORK_DIR"
}

repair_grub() {
    log "Repairing GRUB bootloader..."
    
    # Bind mount necessary filesystems
    mount --bind /dev "$WORK_DIR/dev"
    mount --bind /proc "$WORK_DIR/proc"
    mount --bind /sys "$WORK_DIR/sys"
    mount --bind /run "$WORK_DIR/run"
    
    # Repair GRUB
    chroot "$WORK_DIR" /bin/bash -c "
        echo 'Updating initramfs...'
        update-initramfs -u -k all
        
        echo 'Reinstalling GRUB...'
        grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian --recheck $TARGET_DRIVE
        
        echo 'Updating GRUB configuration...'
        update-grub
        
        echo 'GRUB repair completed'
    "
    
    # Unmount bind mounts
    umount "$WORK_DIR/dev" 2>/dev/null || true
    umount "$WORK_DIR/proc" 2>/dev/null || true
    umount "$WORK_DIR/sys" 2>/dev/null || true
    umount "$WORK_DIR/run" 2>/dev/null || true
    
    success "GRUB repair completed"
}

fix_fstab() {
    log "Checking and fixing /etc/fstab..."
    
    if [ ! -f "$WORK_DIR/etc/fstab" ]; then
        warning "/etc/fstab not found"
        return 0
    fi
    
    # Backup current fstab
    cp "$WORK_DIR/etc/fstab" "$WORK_DIR/etc/fstab.backup.$(date +%Y%m%d_%H%M%S)"
    
    # Get current UUIDs of all partitions
    local partitions=($(lsblk -pno NAME "$TARGET_DRIVE" | grep -v "^$TARGET_DRIVE$"))
    
    echo "Current partition UUIDs:"
    for part in "${partitions[@]}"; do
        local uuid=$(lsblk -no UUID "$part")
        local fstype=$(lsblk -no FSTYPE "$part")
        if [ -n "$uuid" ]; then
            echo "  $part ($fstype): $uuid"
        fi
    done
    
    echo
    echo "Current /etc/fstab content:"
    cat "$WORK_DIR/etc/fstab"
    echo
    
    read -p "Do you want to regenerate /etc/fstab with current UUIDs? [y/N]: " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        # Generate new fstab
        echo "# /etc/fstab: static file system information." > "$WORK_DIR/etc/fstab.new"
        echo "# Regenerated by repair script on $(date)" >> "$WORK_DIR/etc/fstab.new"
        echo "#" >> "$WORK_DIR/etc/fstab.new"
        echo "# <file system> <mount point> <type> <options> <dump> <pass>" >> "$WORK_DIR/etc/fstab.new"
        
        for part in "${partitions[@]}"; do
            local uuid=$(lsblk -no UUID "$part")
            local fstype=$(lsblk -no FSTYPE "$part")
            
            if [ -n "$uuid" ]; then
                case "$fstype" in
                    "vfat")
                        echo "UUID=$uuid /boot/efi vfat defaults 0 2" >> "$WORK_DIR/etc/fstab.new"
                        ;;
                    "ext4")
                        # Determine if it's boot or root by size
                        local size_bytes=$(lsblk -bno SIZE "$part")
                        if [ "$size_bytes" -lt 5368709120 ]; then  # Less than 5GB = boot
                            echo "UUID=$uuid /boot ext4 defaults 0 2" >> "$WORK_DIR/etc/fstab.new"
                        else  # Root partition
                            echo "UUID=$uuid / ext4 defaults 0 1" >> "$WORK_DIR/etc/fstab.new"
                        fi
                        ;;
                    "swap")
                        echo "UUID=$uuid none swap sw 0 0" >> "$WORK_DIR/etc/fstab.new"
                        ;;
                esac
            fi
        done
        
        # Add tmpfs
        echo "tmpfs /tmp tmpfs defaults,noatime,mode=1777 0 0" >> "$WORK_DIR/etc/fstab.new"
        
        # Show new fstab
        echo
        echo "New /etc/fstab content:"
        cat "$WORK_DIR/etc/fstab.new"
        echo
        
        read -p "Apply this new /etc/fstab? [y/N]: " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            mv "$WORK_DIR/etc/fstab.new" "$WORK_DIR/etc/fstab"
            success "/etc/fstab updated"
        else
            rm "$WORK_DIR/etc/fstab.new"
            log "New fstab not applied"
        fi
    fi
    
    success "fstab check completed"
}

check_boot_configuration() {
    log "Checking boot configuration..."
    
    # Check if EFI boot entry exists
    if command -v efibootmgr >/dev/null 2>&1; then
        echo "Current EFI boot entries:"
        efibootmgr
        echo
    fi
    
    # Check GRUB configuration
    if [ -f "$WORK_DIR/etc/default/grub" ]; then
        echo "GRUB configuration (/etc/default/grub):"
        cat "$WORK_DIR/etc/default/grub"
        echo
    fi
    
    # Check if initramfs exists
    echo "Available kernels and initramfs:"
    ls -la "$WORK_DIR/boot/vmlinuz-"* 2>/dev/null || echo "No kernels found"
    ls -la "$WORK_DIR/boot/initrd.img-"* 2>/dev/null || echo "No initramfs found"
    echo
    
    success "Boot configuration check completed"
}

regenerate_initramfs() {
    log "Regenerating initramfs..."
    
    # Bind mount necessary filesystems
    mount --bind /dev "$WORK_DIR/dev"
    mount --bind /proc "$WORK_DIR/proc"
    mount --bind /sys "$WORK_DIR/sys"
    mount --bind /run "$WORK_DIR/run"
    
    chroot "$WORK_DIR" /bin/bash -c "
        echo 'Regenerating initramfs for all kernels...'
        update-initramfs -u -k all
        echo 'Initramfs regeneration completed'
    "
    
    # Unmount bind mounts
    umount "$WORK_DIR/dev" 2>/dev/null || true
    umount "$WORK_DIR/proc" 2>/dev/null || true
    umount "$WORK_DIR/sys" 2>/dev/null || true
    umount "$WORK_DIR/run" 2>/dev/null || true
    
    success "Initramfs regenerated"
}

cleanup_repair() {
    log "Cleaning up repair environment..."
    
    # Unmount all filesystems
    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/repair_*; 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 Repair and Emergency Recovery Tools ===${NC}"
    echo "This script helps fix boot issues after cloning"
    echo
    
    echo "Available repair options:"
    echo "1. Full repair (mount system, fix fstab, repair GRUB, regenerate initramfs)"
    echo "2. GRUB repair only"
    echo "3. fstab fix only"
    echo "4. Initramfs regeneration only"
    echo "5. Check boot configuration"
    echo
    
    read -p "Select repair option [1-5]: " -n 1 -r
    echo
    
    case $REPLY in
        1)
            log "Performing full repair..."
            detect_target_drive
            mount_target_system
            fix_fstab
            regenerate_initramfs
            repair_grub
            check_boot_configuration
            cleanup_repair
            success "Full repair completed!"
            ;;
        2)
            log "Performing GRUB repair only..."
            detect_target_drive
            mount_target_system
            repair_grub
            cleanup_repair
            success "GRUB repair completed!"
            ;;
        3)
            log "Fixing fstab only..."
            detect_target_drive
            mount_target_system
            fix_fstab
            cleanup_repair
            success "fstab fix completed!"
            ;;
        4)
            log "Regenerating initramfs only..."
            detect_target_drive
            mount_target_system
            regenerate_initramfs
            cleanup_repair
            success "Initramfs regeneration completed!"
            ;;
        5)
            log "Checking boot configuration..."
            detect_target_drive
            mount_target_system
            check_boot_configuration
            cleanup_repair
            success "Boot configuration check completed!"
            ;;
        *)
            error "Invalid option selected"
            ;;
    esac
    
    echo
    echo -e "${BLUE}Repair completed. Next steps:${NC}"
    echo "1. Reboot your system"
    echo "2. Set the repaired drive as first boot device in BIOS/UEFI"
    echo "3. Try booting from the external drive"
    echo "4. If issues persist, run this script again or check system logs"
}

# Trap to ensure cleanup on exit
trap cleanup_repair EXIT

main "$@"