#!/bin/bash
# =============================================================================
#  Proxmox VE Kernel Manager
#  Purpose : List all installed PVE kernels and selectively remove old ones.
#            Always protects: (1) running kernel  (2) previous kernel (backup)
# =============================================================================
clear
main() {

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m'

print_banner() {
    echo ""
    echo -e "${CYAN}${BOLD}╔══════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}${BOLD}║         Proxmox VE  ·  Kernel Manager                ║${NC}"
    echo -e "${CYAN}${BOLD}╚══════════════════════════════════════════════════════╝${NC}"
    echo ""
}

print_step() { echo -e "\n${CYAN}▶  ${BOLD}$1${NC}"; }
ok()         { echo -e "  ${GREEN}✔  $1${NC}"; }
warn()       { echo -e "  ${YELLOW}⚠  $1${NC}"; }
err()        { echo -e "  ${RED}✘  $1${NC}"; }
info()       { echo -e "  ${DIM}ℹ  $1${NC}"; }

# ── Root check ────────────────────────────────────────────────────────────────
if [[ "$EUID" -ne 0 ]]; then
    err "This script must be run as root."
    exit 1
fi

print_banner

# ── Step 1: Detect the running kernel ─────────────────────────────────────────
print_step "Detecting currently running kernel ..."
RUNNING_FULL=$(uname -r)
ok "Running kernel : ${BOLD}${RUNNING_FULL}${NC}"

# ── Step 2: Scan installed PVE kernel packages ────────────────────────────────
print_step "Scanning installed Proxmox kernel packages ..."

mapfile -t PKG_LIST < <(
    dpkg --list \
    | awk '/^ii/ && /proxmox-kernel-[0-9]/ && /-pve-signed/ {print $2}' \
    | sort --version-sort
)

# Fallback: unsigned packages (older PVE installs)
if [[ ${#PKG_LIST[@]} -eq 0 ]]; then
    mapfile -t PKG_LIST < <(
        dpkg --list \
        | awk '/^ii/ && /proxmox-kernel-[0-9]/ && /-pve/ && !/-signed/ {print $2}' \
        | sort --version-sort
    )
fi

if [[ ${#PKG_LIST[@]} -eq 0 ]]; then
    err "No PVE kernels found. Exiting."
    exit 1
fi

ok "Found ${#PKG_LIST[@]} installed kernel package(s)."

# ── Step 3: Sort newest-first; define the two protected slots ─────────────────
mapfile -t PKG_SORTED < <(printf '%s\n' "${PKG_LIST[@]}" | tac)

CURRENT_PKG="${PKG_SORTED[0]}"    # slot 1 - running kernel  -> ALWAYS PROTECTED
PREVIOUS_PKG="${PKG_SORTED[1]}"   # slot 2 - backup kernel   -> ALWAYS PROTECTED

# ── Step 4: Build menu ────────────────────────────────────────────────────────
print_step "Installed Kernels (newest first)"
echo ""
printf "  ${BOLD}%-6s  %-52s  %s${NC}\n" "Opt." "Package Name" "Status"
echo -e "  ${DIM}------  ----------------------------------------------------  ----------------------${NC}"

declare -A MENU_MAP
VALID_OPTS=()
OPT=1

for pkg in "${PKG_SORTED[@]}"; do

    # PROTECTED: running kernel (slot 1)
    if [[ "$pkg" == "$CURRENT_PKG" ]] || [[ "$pkg" == *"$RUNNING_FULL"* ]]; then
        STATUS="${GREEN}${BOLD}● In Use  (protected)${NC}"
        OPT_LABEL=" -- "
        printf "  ${BOLD}%-6s${NC}  %-52s  " "$OPT_LABEL" "$pkg"
        echo -e "$STATUS"
        OPT=$((OPT + 1))
        continue
    fi

    # PROTECTED: previous/backup kernel (slot 2) - always protected
    if [[ "$pkg" == "$PREVIOUS_PKG" ]] && [[ -n "$PREVIOUS_PKG" ]]; then
        STATUS="${YELLOW}◑ Previous (backup, protected)${NC}"
        OPT_LABEL=" -- "
        printf "  ${BOLD}%-6s${NC}  %-52s  " "$OPT_LABEL" "$pkg"
        echo -e "$STATUS"
        OPT=$((OPT + 1))
        continue
    fi

    # REMOVABLE: anything older than the top two
    STATUS="${RED}✖ Removable${NC}"
    OPT_LABEL="  ${OPT} "
    MENU_MAP[$OPT]="$pkg"
    VALID_OPTS+=("$OPT")
    printf "  ${BOLD}%-6s${NC}  %-52s  " "$OPT_LABEL" "$pkg"
    echo -e "$STATUS"
    OPT=$((OPT + 1))

done

echo ""

# ── Step 5: Nothing removable? ────────────────────────────────────────────────
if [[ ${#VALID_OPTS[@]} -eq 0 ]]; then
    ok "Nothing to remove - only the running kernel and its backup are installed."
    echo ""
    exit 0
fi

OPTS_STR=$(IFS=" or "; echo "${VALID_OPTS[*]}")

echo -e "  ${BOLD}Options:${NC}"
for n in "${VALID_OPTS[@]}"; do
    echo -e "    ${CYAN}${n}${NC}  ->  Remove: ${BOLD}${MENU_MAP[$n]}${NC}"
done
echo -e "    ${CYAN}A${NC}  ->  Remove ALL old kernels (keeps running + previous backup only)"
echo -e "    ${CYAN}N${NC}  ->  Exit without changes"
echo ""

# ── Step 6: Read user choice ──────────────────────────────────────────────────
# NOTE: echo -ne + read -r used intentionally (no nested quotes = no stream parse errors)
while true; do
    echo -ne "  ${BOLD}Choose option [${OPTS_STR} or A or N]: ${NC}"
    read -r CHOICE
    CHOICE="${CHOICE^^}"

    if [[ "$CHOICE" == "N" ]]; then
        info "No changes made. Exiting."
        echo ""
        exit 0
    fi

    if [[ "$CHOICE" == "A" ]]; then
        TO_REMOVE=()
        for n in "${VALID_OPTS[@]}"; do
            TO_REMOVE+=("${MENU_MAP[$n]}")
        done
        break
    fi

    if [[ "$CHOICE" =~ ^[0-9]+$ ]] && [[ -n "${MENU_MAP[$CHOICE]+_}" ]]; then
        TO_REMOVE=("${MENU_MAP[$CHOICE]}")
        break
    fi

    warn "Invalid option '${CHOICE}'. Please enter one of: ${OPTS_STR} or A or N"
done

# ── Step 7: Confirm ───────────────────────────────────────────────────────────
echo ""
echo -e "  ${YELLOW}${BOLD}The following kernel(s) will be permanently removed:${NC}"
for pkg in "${TO_REMOVE[@]}"; do
    echo -e "    ${RED}✖  ${pkg}${NC}"
done
echo ""
echo -ne "  ${BOLD}Are you sure? [y/N]: ${NC}"
read -r CONFIRM
CONFIRM="${CONFIRM,,}"

if [[ "$CONFIRM" != "y" && "$CONFIRM" != "yes" ]]; then
    info "Removal cancelled. No changes made."
    echo ""
    exit 0
fi

# ── Step 8: Remove ────────────────────────────────────────────────────────────
echo ""
print_step "Removing selected kernel(s) ..."

REMOVED=0
FAILED=0

for pkg in "${TO_REMOVE[@]}"; do

    # Hard safety net - never touch the two protected kernels under any circumstance
    if [[ "$pkg" == "$CURRENT_PKG" ]] || [[ "$pkg" == "$PREVIOUS_PKG" ]] || [[ "$pkg" == *"$RUNNING_FULL"* ]]; then
        warn "Skipping protected kernel: ${pkg}"
        continue
    fi

    BASE_PKG="${pkg%-signed}"
    PKGS_TO_PURGE=("$pkg")
    if dpkg -l "$BASE_PKG" 2>/dev/null | grep -q "^ii"; then
        PKGS_TO_PURGE+=("$BASE_PKG")
        info "Also removing base package: ${BASE_PKG}"
    fi

    echo -ne "  ${CYAN}Removing:${NC} ${BOLD}${pkg}${NC} ... "

    if apt-get purge -y "${PKGS_TO_PURGE[@]}" >> /tmp/pve-kernel-rm.log 2>&1; then
        echo -e "${GREEN}${BOLD}Done${NC}"
        ok "Successfully removed: ${pkg}"
        REMOVED=$((REMOVED + 1))
    else
        echo -e "${RED}${BOLD}Failed${NC}"
        err "Failed via apt. Trying dpkg fallback ..."
        if dpkg --purge "${PKGS_TO_PURGE[@]}" >> /tmp/pve-kernel-rm.log 2>&1; then
            ok "Successfully removed via dpkg: ${pkg}"
            REMOVED=$((REMOVED + 1))
        else
            err "dpkg fallback also failed: ${pkg}"
            info "See /tmp/pve-kernel-rm.log"
            FAILED=$((FAILED + 1))
        fi
    fi

done

# ── Step 9: Autoremove orphans ────────────────────────────────────────────────
if [[ $REMOVED -gt 0 ]]; then
    print_step "Cleaning up orphaned kernel dependencies ..."
    apt-get autoremove --purge -y >> /tmp/pve-kernel-rm.log 2>&1 \
        && ok "Orphaned packages cleaned up." \
        || warn "autoremove had warnings (see /tmp/pve-kernel-rm.log)."
fi

# ── Step 10: Refresh bootloader ───────────────────────────────────────────────
if [[ $REMOVED -gt 0 ]]; then
    print_step "Updating bootloader ..."
    if command -v proxmox-boot-tool &>/dev/null; then
        proxmox-boot-tool refresh >> /tmp/pve-kernel-rm.log 2>&1 \
            && ok "proxmox-boot-tool refresh completed." \
            || warn "proxmox-boot-tool refresh had warnings."
    fi
    if command -v update-grub &>/dev/null; then
        update-grub >> /tmp/pve-kernel-rm.log 2>&1 \
            && ok "GRUB updated successfully." \
            || warn "update-grub had warnings."
    fi
fi

# ── Step 11: Verify removal ───────────────────────────────────────────────────
if [[ $REMOVED -gt 0 ]]; then
    print_step "Verifying removal ..."
    for pkg in "${TO_REMOVE[@]}"; do
        if dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
            warn "Package still present: ${pkg} - manual intervention may be needed."
        else
            ok "Confirmed removed: ${pkg}"
        fi
    done
fi

# ── Step 12: Summary ──────────────────────────────────────────────────────────
echo ""
echo -e "${CYAN}${BOLD}╔══════════════════════════════════════════╗${NC}"
echo -e "${CYAN}${BOLD}║              Summary                     ║${NC}"
echo -e "${CYAN}${BOLD}╚══════════════════════════════════════════╝${NC}"
[[ $REMOVED -gt 0 ]] && ok "Kernels removed  : ${REMOVED}"
[[ $FAILED  -gt 0 ]] && err "Kernels failed   : ${FAILED}"
ok "Running kernel   : ${RUNNING_FULL}  (untouched)"
[[ -n "$PREVIOUS_PKG" ]] && ok "Backup kernel    : ${PREVIOUS_PKG}  (untouched)"
echo ""
info "Full log saved to /tmp/pve-kernel-rm.log"
echo ""

} # end of main()

main "$@"