#!/bin/bash

# Configuration variables
RAID_DEVICE="/dev/md0"
EFI_PARTITIONS=(
    "/dev/sdd2"
    "/dev/sdb2"
    "/dev/nvme2n1p2"
    "/dev/sdc2"
    "/dev/nvme1n1p2"
    "/dev/sda2"
)
MOUNT_POINT="/mnt/efi"
BACKUP_DIR="/tmp/efi_backup"
LOG_FILE="/var/log/raid_bootloader_install.log"

# Logging function with timestamp
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') : $1" | tee -a "$LOG_FILE"
}

# Error handling function
handle_error() {
    log "ERROR: $1"
    exit 1
}

# Check if running with root privileges
check_root() {
    if [ "$(id -u)" != "0" ]; then
        handle_error "This script must be run as root"
    fi
}

# Ensure mount point is unmounted
ensure_unmounted() {
    if mountpoint -q "$MOUNT_POINT"; then
        umount "$MOUNT_POINT" || handle_error "Failed to unmount $MOUNT_POINT"
    fi
}

# Mount partition safely
safe_mount() {
    local partition=$1
    local mount_point=$2
    mount "$partition" "$mount_point" || handle_error "Failed to mount $partition on $mount_point"
}

# Unmount partition safely
safe_unmount() {
    local mount_point=$1
    umount "$mount_point" || handle_error "Failed to unmount $mount_point"
}

# Generic bootloader installation function
install_bootloader() {
    local partition=$1
    local bootloader=$2
    local install_cmd=$3
    
    log "Installing $bootloader on $partition"
    safe_mount "$partition" "$MOUNT_POINT"
    eval "$install_cmd" || handle_error "Failed to install $bootloader on $partition"
    safe_unmount "$MOUNT_POINT"
    log "$bootloader successfully installed on $partition"
}

# Install GRUB
install_grub() {
    install_bootloader "$1" "GRUB" "grub-install --target=x86_64-efi --efi-directory=$MOUNT_POINT --bootloader-id=GRUB --recheck"
}

# Install systemd-boot
install_systemd_boot() {
    install_bootloader "$1" "systemd-boot" "bootctl --path=$MOUNT_POINT install"
}

# Install rEFInd
install_refind() {
    install_bootloader "$1" "rEFInd" "refind-install --root $MOUNT_POINT"
}

# Install Limine
install_limine() {
    install_bootloader "$1" "Limine" "limine-install $MOUNT_POINT"
}

# Install LILO
install_lilo() {
    install_bootloader "$1" "LILO" "lilo -M $1"
}

# Copy EFI contents
copy_efi_contents() {
    local partition=$1
    log "Copying EFI contents to $partition"
    safe_mount "$partition" "$MOUNT_POINT"
    cp -r /boot/efi/* "$MOUNT_POINT/" || handle_error "Failed to copy EFI contents to $partition"
    safe_unmount "$MOUNT_POINT"
    log "EFI contents successfully copied to $partition"
}

# Backup EFI partition
backup_efi() {
    local partition=$1
    mkdir -p "$BACKUP_DIR"
    log "Backing up EFI partition from $partition"
    safe_mount "$partition" "$MOUNT_POINT"
    cp -r "$MOUNT_POINT/"* "$BACKUP_DIR/" || handle_error "Failed to backup EFI contents from $partition"
    safe_unmount "$MOUNT_POINT"
    log "EFI partition backup completed"
}

# Restore EFI contents to RAID
restore_efi() {
    log "Restoring EFI contents to RAID array"
    safe_mount "$RAID_DEVICE" "$MOUNT_POINT"
    cp -r "$BACKUP_DIR/"* "$MOUNT_POINT/" || handle_error "Failed to restore EFI contents to RAID"
    safe_unmount "$MOUNT_POINT"
    log "EFI contents restored to RAID array"
}

main() {
    # Check root privileges
    check_root

    # Create mount point if it doesn't exist
    mkdir -p "$MOUNT_POINT"

    # Stop existing RAID array
    log "Stopping RAID array $RAID_DEVICE"
    mdadm --stop "$RAID_DEVICE" || handle_error "Failed to stop RAID array"

    # Ensure mount point is unmounted
    ensure_unmounted

    # Install bootloaders on each partition
    for partition in "${EFI_PARTITIONS[@]}"; do
        log "Processing partition: $partition"
        install_grub "$partition"
        install_systemd_boot "$partition"
        install_refind "$partition"
        install_limine "$partition"
        install_lilo "$partition"
        copy_efi_contents "$partition"
    done

    # Backup EFI contents from first partition
    backup_efi "${EFI_PARTITIONS[0]}"

    # Create new RAID array with metadata 1.0
    log "Creating new RAID1 array with metadata 1.0"
    mdadm --create --verbose "$RAID_DEVICE" --level=1 --raid-devices=${#EFI_PARTITIONS[@]} \
        --metadata=1.0 "${EFI_PARTITIONS[@]}" || handle_error "Failed to create RAID array"

    # Restore EFI contents to RAID
    restore_efi

    # Mount RAID array to /boot/efi
    log "Mounting RAID array to /boot/efi"
    mount "$RAID_DEVICE" /boot/efi || handle_error "Failed to mount RAID array to /boot/efi"

    log "Bootloader installation and RAID array rebuild completed successfully"
}

# Execute main function
main