#!/bin/bash
#
# Gentoo Kernel Upgrade Script
# This script helps upgrade the gentoo-kernel-bin package safely.
#
# IMPORTANT: This script modifies boot files. A failed upgrade could
# prevent your system from booting. Ensure you have a backup plan.
#

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
NC='\033[0m' # No Color

EFI_DIR="/boot/efi"
BOOT_DIR="/boot"

# Print functions
info()    { echo -e "${BLUE}==>${NC} ${BOLD}$1${NC}"; }
warn()    { echo -e "${YELLOW}WARNING:${NC} $1"; }
error()   { echo -e "${RED}ERROR:${NC} $1" >&2; }
success() { echo -e "${GREEN}==>${NC} $1"; }

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root"
        exit 1
    fi
}

# Prompt user for confirmation
confirm() {
    local prompt="$1"
    local response
    echo -e "${YELLOW}${prompt}${NC}"
    read -p "Type 'yes' to confirm: " response
    [[ "$response" == "yes" ]]
}

# Get the currently running kernel version
get_current_version() {
    uname -r
}

# Get the latest installed kernel version from /usr/src
get_latest_installed_version() {
    local latest
    latest=$(ls -1d /usr/src/linux-*-gentoo-dist 2>/dev/null | sort -V | tail -1 | sed 's|.*/linux-||')
    echo "$latest"
}

# Check for available updates using emerge
check_for_updates() {
    info "Checking for gentoo-kernel-bin updates..."
    echo

    # Use emerge to check for updates (pretend mode)
    if emerge -pvu gentoo-kernel-bin 2>/dev/null | grep -q "gentoo-kernel-bin"; then
        local update_info
        update_info=$(emerge -pvu gentoo-kernel-bin 2>/dev/null | grep -E "gentoo-kernel-bin|ebuild")
        echo "$update_info"
        return 0
    else
        return 1
    fi
}

# Install the kernel update
install_kernel_update() {
    info "Installing gentoo-kernel-bin update..."
    echo

    echo -e "${BOLD}The following command will be executed:${NC}"
    echo -e "  ${BLUE}emerge gentoo-kernel-bin${NC}"
    echo

    if ! confirm "Proceed with kernel package installation?"; then
        warn "Installation cancelled by user"
        return 1
    fi

    echo
    emerge gentoo-kernel-bin

    success "Kernel package installed successfully"
}

# Select the new kernel with eselect
select_kernel() {
    local new_version="$1"

    info "Available kernel versions:"
    eselect kernel list
    echo

    # Find the number corresponding to the new version
    local kernel_num
    kernel_num=$(eselect kernel list | grep "$new_version" | sed 's/.*\[\([0-9]*\)\].*/\1/')

    if [[ -z "$kernel_num" ]]; then
        error "Could not find kernel version $new_version in eselect list"
        echo "Please select the kernel manually:"
        eselect kernel list
        read -p "Enter the number to select: " kernel_num
    fi

    echo -e "${BOLD}The following command will be executed:${NC}"
    echo -e "  ${BLUE}eselect kernel set $kernel_num${NC}"
    echo

    if ! confirm "Set kernel $new_version as active?"; then
        warn "Kernel selection cancelled by user"
        return 1
    fi

    eselect kernel set "$kernel_num"

    # Verify
    info "Current kernel symlink:"
    eselect kernel list | grep '\*'
    echo
}

# Backup current boot files
backup_boot_files() {
    info "Backing up current boot files..."

    echo -e "${BOLD}The following operations will be performed in ${EFI_DIR}:${NC}"
    echo -e "  ${BLUE}mv vmlinuz.efi -> vmlinuz.efi.bak${NC}"
    echo -e "  ${BLUE}mv initramfs.img -> initramfs.img.bak${NC}"
    echo

    if [[ -f "${EFI_DIR}/vmlinuz.efi.bak" ]] || [[ -f "${EFI_DIR}/initramfs.img.bak" ]]; then
        warn "Existing backup files will be overwritten!"
    fi

    if ! confirm "Backup current boot files?"; then
        warn "Backup cancelled by user"
        return 1
    fi

    cd "${EFI_DIR}"

    if [[ -f "vmlinuz.efi" ]]; then
        mv vmlinuz.efi vmlinuz.efi.bak
        success "Backed up vmlinuz.efi"
    fi

    if [[ -f "initramfs.img" ]]; then
        mv initramfs.img initramfs.img.bak
        success "Backed up initramfs.img"
    fi

    echo
}

# Generate new initramfs
generate_initramfs() {
    local version="$1"

    info "Generating new initramfs for kernel ${version}..."
    echo

    echo -e "${BOLD}The following command will be executed:${NC}"
    echo -e "  ${BLUE}${EFI_DIR}/generate_initramfs.sh ${version} ${EFI_DIR}/initramfs.img${NC}"
    echo

    if ! confirm "Generate initramfs?"; then
        warn "Initramfs generation cancelled by user"
        return 1
    fi

    bash "${EFI_DIR}/generate_initramfs.sh" "${version}" "${EFI_DIR}/initramfs.img"

    success "Initramfs generated successfully"
    echo
}

# Copy new kernel to EFI partition
copy_kernel() {
    local version="$1"
    local kernel_source="${BOOT_DIR}/vmlinuz-${version}"

    info "Copying new kernel to EFI partition..."
    echo

    if [[ ! -f "$kernel_source" ]]; then
        error "Kernel file not found: $kernel_source"
        return 1
    fi

    echo -e "${BOLD}The following command will be executed:${NC}"
    echo -e "  ${BLUE}cp ${kernel_source} ${EFI_DIR}/vmlinuz.efi${NC}"
    echo

    if ! confirm "Copy new kernel to EFI partition?"; then
        warn "Kernel copy cancelled by user"
        return 1
    fi

    cp "$kernel_source" "${EFI_DIR}/vmlinuz.efi"

    success "Kernel copied successfully"
    echo
}

# Cleanup old kernel versions from /boot
cleanup_old_versions() {
    info "Checking for old kernel versions to clean up..."
    echo

    # Find old initramfs and kernel files (excluding current version)
    local initramfs_to_delete=()
    local kernels_to_delete=()

    

    while IFS= read -r f; do
        initramfs_to_delete+=("$f")
    done < <(ls ${BOOT_DIR}/initramfs-*.img 2>/dev/null | grep -v "$current_version" || true)

    while IFS= read -r f; do
        kernels_to_delete+=("$f")
    done < <(ls ${BOOT_DIR}/vmlinuz-* 2>/dev/null | grep -v "$current_version" || true)

    # Check if there's anything to delete
    if [[ ${#initramfs_to_delete[@]} -eq 0 ]] && [[ ${#kernels_to_delete[@]} -eq 0 ]]; then
        success "No old kernel versions to clean up"
        echo
        return 0
    fi

    echo -e "${BOLD}The following files will be DELETED:${NC}"
    echo
    for f in "${initramfs_to_delete[@]}" "${kernels_to_delete[@]}"; do
        echo -e "  ${RED}$f${NC}"
    done
    echo

    warn "This action cannot be undone!"
    echo

    if ! confirm "Delete these old kernel files?"; then
        warn "Cleanup cancelled by user"
        return 1
    fi

    echo
    for f in "${initramfs_to_delete[@]}" "${kernels_to_delete[@]}"; do
        rm -v "$f"
    done

    echo
    success "Old kernel versions cleaned up"
    echo
}

# Print summary of what will be done
print_summary() {
    local current="$1"
    local new="$2"

    echo
    echo -e "${BOLD}╔════════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BOLD}║              GENTOO KERNEL UPGRADE SUMMARY                     ║${NC}"
    echo -e "${BOLD}╚════════════════════════════════════════════════════════════════╝${NC}"
    echo
    echo -e "  Current running kernel:  ${YELLOW}${current}${NC}"
    echo -e "  New kernel version:      ${GREEN}${new}${NC}"
    echo
    echo -e "${BOLD}The following steps will be performed:${NC}"
    echo "  1. Backup current vmlinuz.efi and initramfs.img"
    echo "  2. Generate new initramfs using dracut"
    echo "  3. Copy new kernel to EFI partition"
    echo
    echo -e "${RED}${BOLD}WARNING: This process modifies boot files!${NC}"
    echo -e "${RED}If something goes wrong, your system may not boot.${NC}"
    echo -e "${RED}Make sure you have a way to recover (live USB, backup, etc.)${NC}"
    echo
}

# Main function
main() {
    echo
    echo -e "${BOLD}╔════════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BOLD}║           GENTOO KERNEL UPGRADE HELPER                         ║${NC}"
    echo -e "${BOLD}╚════════════════════════════════════════════════════════════════╝${NC}"
    echo

    check_root

    local current_version
    current_version=$(get_current_version)
    info "Currently running kernel: ${current_version}"
    echo

    # Check for updates
    if ! check_for_updates; then
        success "No updates available for gentoo-kernel-bin"
        echo

        local latest_installed
        latest_installed=$(get_latest_installed_version)

        if [[ "$latest_installed" != "$current_version" ]] && [[ -n "$latest_installed" ]]; then
            # Installed kernel differs from running - offer to set up boot files
            warn "Installed kernel ($latest_installed) differs from running kernel ($current_version)"
            echo
            if confirm "Would you like to set up boot files for $latest_installed?"; then
                print_summary "$current_version" "$latest_installed"

                if ! confirm "Proceed with the upgrade process?"; then
                    echo "Upgrade cancelled."
                    exit 0
                fi

                select_kernel "$latest_installed"
                backup_boot_files
                generate_initramfs "$latest_installed"
                copy_kernel "$latest_installed"

                echo
                success "Kernel upgrade complete!"
                echo
                echo -e "${BOLD}Next steps:${NC}"
                echo "  1. Reboot your system"
                echo "  2. Verify the new kernel is running: uname -r"
                echo "  3. Run this script again to clean up old versions"
                echo
            fi
        else
            # Running kernel is the newest - offer cleanup
            cleanup_old_versions
        fi
        exit 0
    fi

    echo
    if ! confirm "Would you like to install this kernel update?"; then
        echo "Update cancelled."
        exit 0
    fi

    echo
    install_kernel_update

    # Get the new version after installation
    local new_version
    new_version=$(get_latest_installed_version)

    if [[ -z "$new_version" ]]; then
        error "Could not determine new kernel version"
        exit 1
    fi

    print_summary "$current_version" "$new_version"

    if ! confirm "Proceed with the upgrade process?"; then
        echo "Upgrade process cancelled. Kernel package is installed but boot files unchanged."
        exit 0
    fi

    echo
    select_kernel "$new_version"
    backup_boot_files
    generate_initramfs "$new_version"
    copy_kernel "$new_version"

    echo
    success "Kernel upgrade complete!"
    echo
    echo -e "${BOLD}╔════════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BOLD}║                         NEXT STEPS                             ║${NC}"
    echo -e "${BOLD}╚════════════════════════════════════════════════════════════════╝${NC}"
    echo
    echo "  1. ${YELLOW}Reboot your system${NC}"
    echo "  2. Verify the new kernel is running:"
    echo "       uname -r"
    echo "     Expected: ${GREEN}${new_version}${NC}"
    echo
    echo "  3. If the new kernel works correctly, run this script again"
    echo "     to clean up old versions"
    echo
    echo "  4. ${RED}If the system fails to boot:${NC}"
    echo "     - Boot from a live USB"
    echo "     - Mount your EFI partition"
    echo "     - Restore backups:"
    echo "         mv vmlinuz.efi.bak vmlinuz.efi"
    echo "         mv initramfs.img.bak initramfs.img"
    echo
}

main "$@"