#!/usr/bin/env bash
set -o errexit
set -o nounset
set -o pipefail

# APT autoremove respects package dependencies and kernel protection rules. That
# is safer than name-based purging on HWE hosts using NVIDIA, DKMS, or VFIO.

LOG_FILE="/var/log/ailab-kernel-cleanup.log"
execute=false
non_interactive=false

usage() {
  printf 'Usage: %s [--execute [--non-interactive]]\n' "$(basename "$0")"
}

kernel_packages() {
  dpkg-query -W -f='${db:Status-Abbrev} ${binary:Package}\n' \
    'linux-image*' 'linux-headers*' 'linux-modules*' 2>/dev/null \
    | awk '$1 ~ /^ii/ {print $2}' \
    | sort -u || true
}

versioned_kernel_images() {
  dpkg-query -W -f='${db:Status-Abbrev} ${binary:Package}\n' 'linux-image-[0-9]*' 2>/dev/null \
    | awk '$1 ~ /^ii/ {sub(/:.*/, "", $2); print $2}' \
    | sort -u || true
}

while (($# > 0)); do
  case "$1" in
    --execute) execute=true ;;
    --non-interactive) non_interactive=true ;;
    -h|--help) usage; exit 0 ;;
    *) printf 'CRITICAL: unknown argument: %s\n' "$1" >&2; usage >&2; exit 2 ;;
  esac
  shift
done

if [[ "$non_interactive" == true && "$execute" != true ]]; then
  printf 'CRITICAL: --non-interactive requires --execute\n' >&2
  exit 2
fi
if ((EUID != 0)); then
  printf 'CRITICAL: this script must run as root\n' >&2
  exit 2
fi
for command_name in apt dpkg-query uname; do
  if ! command -v "$command_name" >/dev/null 2>&1; then
    printf 'CRITICAL: required command is missing: %s\n' "$command_name" >&2
    exit 2
  fi
done

exec > >(tee -a "$LOG_FILE") 2>&1
printf '\n[%s] Kernel cleanup\n' "$(date --iso-8601=seconds)"
printf 'Running kernel: %s\n' "$(uname -r)"
printf '\nInstalled kernel-related packages before cleanup:\n'
kernel_packages

simulation="$(LC_ALL=C apt -s autoremove --purge)"
printf '\nAPT autoremove simulation:\n%s\n' "$simulation"

mapfile -t installed_images < <(versioned_kernel_images)
mapfile -t removed_images < <(
  awk '$1 == "Remv" && $2 ~ /^linux-image-[0-9]/ {sub(/:.*/, "", $2); print $2}' <<<"$simulation" | sort -u
)

remaining_images=0
for image in "${installed_images[@]}"; do
  remove_image=false
  for removed in "${removed_images[@]}"; do
    if [[ "$image" == "$removed" ]]; then
      remove_image=true
      break
    fi
  done
  if [[ "$remove_image" != true ]]; then
    remaining_images=$((remaining_images + 1))
  fi
done

printf 'Kernel image safety check: installed=%d simulated-removals=%d remaining=%d\n' \
  "${#installed_images[@]}" "${#removed_images[@]}" "$remaining_images"

if ((${#installed_images[@]} < 2 || remaining_images < 2)); then
  printf 'CRITICAL: cleanup would not leave at least two versioned kernel images; refusing execution\n'
  exit 1
fi

if [[ "$execute" != true ]]; then
  printf 'INFO: dry-run mode; no packages were removed\n'
  printf 'INFO: rerun with --execute and confirm to apply the simulated cleanup\n'
  exit 0
fi

if [[ "$non_interactive" != true ]]; then
  printf 'WARNING: APT will remove the packages shown in the simulation above.\n'
  printf 'Type EXECUTE to continue: '
  read -r confirmation
  if [[ "$confirmation" != "EXECUTE" ]]; then
    printf 'CRITICAL: confirmation failed; no changes made\n'
    exit 2
  fi
fi

apt autoremove --purge -y
apt autoclean -y
if command -v update-grub >/dev/null 2>&1; then
  update-grub || true
else
  printf 'WARNING: update-grub is not installed\n'
fi

printf '\nInstalled kernel-related packages after cleanup:\n'
kernel_packages
printf 'OK: kernel cleanup completed with APT-managed package selection\n'