#!/bin/bash
#
# dist-upgrade - update a Slackware -current system with slackpkg
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, see <https://www.gnu.org/licenses/>.

export LC_ALL=C

INITRD=${INITRD:-/boot/initrd-generic.img}
MKINITRD_GEN=/usr/share/mkinitrd/mkinitrd_command_generator.sh

INTERACTIVE=0
TMPDIR_SNAP=""
ORIG_ARGS=("$@")
RESTARTS=${DIST_UPGRADE_RESTARTS:-0}

die() { printf 'dist-upgrade: %s\n' "$*" >&2; exit 1; }
warn() { printf 'WARNING: %s\n' "$*" >&2; }
info() { printf '%s\n' "$*"; }

usage() {
  cat <<'EOF'
Usage: dist-upgrade [-i|--interactive] [-h|--help] [--self-test]

Update a Slackware -current system: slackpkg update, install-new,
upgrade-all, clean-system and new-config, then verify the initrd still
matches the newly installed kernel and still holds the modules needed
to boot.

  -i, --interactive   run slackpkg with its normal prompts
  -h, --help          show this help
      --self-test     run the internal logic checks and exit

Environment:
  INITRD              initrd to inspect (default /boot/initrd-generic.img)
EOF
}

norm() { tr '_' '-'; }

initrd_cat() {
  local f=$1 magic
  magic=$(head -c4 "$f" | od -An -tx1 | tr -d ' \n')
  case "$magic" in
    1f8b*)     gzip -dc "$f" ;;
    28b52ffd)  zstd -dc "$f" ;;
    fd377a58)  xz -dc "$f" ;;
    *)         cat "$f" ;;
  esac
}

# ponytail: only gzip/zstd/xz/uncompressed are detected, mkinitrd emits one of these
initrd_list() { initrd_cat "$1" | cpio -t 2>/dev/null; }

modules_from_listing() {
  grep -E '\.ko(\.(xz|zst|gz))?$' | sed -E 's#.*/##; s/\.ko.*$//'
}

version_from_listing() {
  grep -oE 'lib/modules/[^/]+' | head -n1 | sed 's#.*lib/modules/##'
}

explicit_modules() {
  awk '{
    for (i = 1; i <= NF; i++) {
      if ($i == "-f") { print $(i+1) }
      else if ($i == "-m") { n = split($(i+1), m, ":"); for (j = 1; j <= n; j++) print m[j] }
    }
  }' | tr -d '"' | sed '/^$/d'
}

initrd_modules() { initrd_list "$1" | modules_from_listing | sort -u; }
initrd_version() { initrd_list "$1" | version_from_listing; }

# ponytail: sort -V; rc-to-final ordering edge cases are not handled
newest_kernel() {
  local d
  for d in /lib/modules/*/; do
    [ -d "$d" ] || continue
    d=${d%/}; d=${d##*/}
    printf '%s\n' "$d"
  done | sort -V | tail -n1
}

needed_modules() { "$MKINITRD_GEN" -k "$1" 2>/dev/null | explicit_modules; }

builtin_modules() {
  grep -oE '[^/]+\.ko' "/lib/modules/$1/modules.builtin" 2>/dev/null | sed 's/\.ko$//' | norm | sort -u
}

run_slackpkg() {
  if [ "$INTERACTIVE" -eq 1 ]; then
    slackpkg "$@"
  else
    slackpkg -batch=on -default_answer=y "$@"
  fi
}

# slackpkg (via cleanup() in core-functions.sh) exits 20 when nothing matches,
# and 50 when it upgraded itself and wants the operation restarted.
handle_rc() {
  local cmd=$1 rc=$2
  case "$rc" in
    0)  return 0 ;;
    20) info "slackpkg $cmd: nothing to do." ;;
    50)
      if [ "$RESTARTS" -ge 1 ]; then
        warn "slackpkg upgraded itself again mid-run; re-run dist-upgrade to continue."
        exit 1
      fi
      info "slackpkg upgraded itself mid-run; restarting dist-upgrade to continue."
      RESTARTS=$((RESTARTS + 1))
      export DIST_UPGRADE_RESTARTS=$RESTARTS
      [ -n "$TMPDIR_SNAP" ] && rm -rf "$TMPDIR_SNAP"
      exec "$0" "${ORIG_ARGS[@]}"
      ;;
    *)  die "slackpkg $cmd failed (exit $rc)" ;;
  esac
}

run_step() {
  local cmd=$1
  shift
  run_slackpkg "$cmd" "$@"
  handle_rc "$cmd" $?
}

snapshot_initrd() {
  if [ ! -e "$INITRD" ] && [ ! -L "$INITRD" ]; then
    info "snapshot: no initrd at $INITRD, skipping initrd checks."
    return 0
  fi
  initrd_modules "$INITRD" > "$TMPDIR_SNAP/mods.before"
  initrd_version "$INITRD" > "$TMPDIR_SNAP/ver.before"
  info "snapshot: $(readlink -f "$INITRD") kernel $(cat "$TMPDIR_SNAP/ver.before"), $(wc -l < "$TMPDIR_SNAP/mods.before") modules"
}

check_initrd() {
  local target newver iver needed have missing dropped before gen_ok

  if [ ! -e "$INITRD" ] && [ ! -L "$INITRD" ]; then
    info "No initrd at $INITRD (huge-kernel setup?), skipping initrd checks."
    return 0
  fi

  target=$(readlink -f "$INITRD" 2>/dev/null || printf '%s' "$INITRD")
  newver=$(newest_kernel)
  iver=$(initrd_version "$INITRD")

  if [ -z "$iver" ]; then
    warn "could not determine the kernel version inside $target"
  elif [ -n "$newver" ] && [ "$iver" != "$newver" ]; then
    warn "initrd $target is built for kernel $iver but $newver is installed; regenerate it or the system may not boot"
  else
    info "initrd kernel version matches installed kernel: $iver"
  fi

  gen_ok=0
  if [ -x "$MKINITRD_GEN" ] && [ -n "$newver" ] && [ -d "/lib/modules/$newver" ]; then
    gen_ok=1
    needed=$(needed_modules "$newver" | norm | sort -u)
    if [ -f "/lib/modules/$newver/modules.builtin" ]; then
      needed=$(comm -23 <(printf '%s\n' "$needed") <(builtin_modules "$newver"))
    fi
  fi

  if [ "$gen_ok" -eq 0 ]; then
    info "note: mkinitrd_command_generator unavailable for $newver, skipping bootability check."
  elif [ -z "$needed" ]; then
    info "initrd needs no extra modules: everything the boot config requires is built into the kernel."
  else
    have=$(initrd_modules "$INITRD" | norm | sort -u)
    missing=$(comm -23 <(printf '%s\n' "$needed") <(printf '%s\n' "$have"))
    if [ -n "$missing" ]; then
      warn "initrd $target is missing modules the boot config needs: $(printf '%s' "$missing" | tr '\n' ' ')"
    else
      info "initrd contains every module mkinitrd_command_generator says are needed: $(printf '%s' "$needed" | tr '\n' ' ')"
    fi
  fi

  before="$TMPDIR_SNAP/mods.before"
  if [ -s "$before" ]; then
    dropped=$(comm -23 <(norm < "$before" | sort -u) <(initrd_modules "$INITRD" | norm | sort -u))
    if [ -n "$dropped" ]; then
      warn "initrd lost modules that were present before the upgrade: $(printf '%s' "$dropped" | tr '\n' ' ')"
    else
      info "initrd module set unchanged, nothing dropped."
    fi
  fi
}

self_test() {
  local sample expected got a b dropped d img

  sample='lib
lib/modules
lib/modules/6.18.50
lib/modules/6.18.50/kernel/fs/ext4/ext4.ko
lib/modules/6.18.50/kernel/drivers/usb/host/xhci-hcd.ko.xz
lib/modules/6.18.50/modules.dep'

  got=$(printf '%s\n' "$sample" | modules_from_listing | sort -u | tr '\n' ' ')
  expected="ext4 xhci-hcd "
  [ "$got" = "$expected" ] || die "self-test: module extraction failed: '$got'"

  got=$(printf '%s\n' "$sample" | version_from_listing)
  [ "$got" = "6.18.50" ] || die "self-test: version extraction failed: '$got'"

  a=$(printf 'ext4\nusbhid\nxhci-hcd\n' | sort)
  b=$(printf 'ext4\nxhci-hcd\n' | sort)
  dropped=$(comm -23 <(printf '%s\n' "$a") <(printf '%s\n' "$b"))
  [ "$dropped" = "usbhid" ] || die "self-test: module diff failed: '$dropped'"

  got=$(printf '# mkinitrd -c -k 6.18.50 -f ext4 -r /dev/sda2 -m ext4:usbhid \\\n#   -u -o /boot/initrd-6.18.50.img\n' | explicit_modules | sort -u | tr '\n' ' ')
  expected="ext4 usbhid "
  [ "$got" = "$expected" ] || die "self-test: mkinitrd parse failed: '$got'"

  got=$(printf 'mkinitrd -c -k 6.18.50 -f ext4 -r /dev/sda2 -u -o /boot/initrd-6.18.50.img\n' | explicit_modules | sort -u | tr '\n' ' ')
  expected="ext4 "
  [ "$got" = "$expected" ] || die "self-test: mkinitrd -f parse failed: '$got'"

  ( handle_rc test 0 ) >/dev/null 2>&1 || die "self-test: rc 0 should pass"
  ( handle_rc test 20 ) >/dev/null 2>&1 || die "self-test: rc 20 should pass"
  if ( RESTARTS=1; handle_rc test 50 ) >/dev/null 2>&1; then die "self-test: rc 50 should abort once restarted"; fi
  if ( handle_rc test 1 ) >/dev/null 2>&1; then die "self-test: rc 1 should abort"; fi

  d=$(mktemp -d) || die "self-test: mktemp failed"
  mkdir -p "$d/lib/modules/6.18.50/kernel/fs/ext4" "$d/lib/modules/6.18.50/kernel/drivers/usb/host"
  : > "$d/lib/modules/6.18.50/kernel/fs/ext4/ext4.ko"
  : > "$d/lib/modules/6.18.50/kernel/drivers/usb/host/xhci-hcd.ko.xz"
  img="$d/initrd.img"
  ( cd "$d" && find . -mindepth 1 | cpio -o -H newc 2>/dev/null | gzip -c ) > "$img"

  got=$(initrd_version "$img")
  [ "$got" = "6.18.50" ] || die "self-test: initrd version read failed: '$got'"

  got=$(initrd_modules "$img" | sort | tr '\n' ' ')
  expected="ext4 xhci-hcd "
  [ "$got" = "$expected" ] || die "self-test: initrd module read failed: '$got'"

  rm -rf "$d"

  info "self-test: all checks passed"
}

for arg in "$@"; do
  case "$arg" in
    -i|--interactive) INTERACTIVE=1 ;;
    -h|--help)        usage; exit 0 ;;
    --self-test)      self_test; exit 0 ;;
    *)                usage; die "unknown option: $arg" ;;
  esac
done

[ "$(id -u)" -eq 0 ] || die "must be run as root"

TMPDIR_SNAP=$(mktemp -d) || die "mktemp failed"
trap 'rm -rf "$TMPDIR_SNAP"' EXIT

snapshot_initrd

run_step update
run_step install-new
run_step upgrade-all
if [ "$INTERACTIVE" -eq 1 ]; then
  run_step clean-system
else
  info "clean-system: listing packages not in the official tree only, removing nothing (use -i to review)."
  slackpkg -batch=on -default_answer=n clean-system
  handle_rc clean-system $?
fi
run_step new-config

check_initrd

running=$(uname -r)
newver=$(newest_kernel)
if [ -n "$newver" ] && [ "$running" != "$newver" ]; then
  info "Reboot recommended: running kernel $running, newest installed $newver."
fi

info "Done."
