#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-3.0-only
# eniatos-update - unattended system upgrade: pacman -Syu, restart of services running
# replaced code, boot-safety gate, reboot when required, post-boot health check
set -uo pipefail

PROG=eniatos-update
CONF=/etc/eniatos/update.conf
STATE=/var/lib/eniatos-update
MOTD=/etc/motd.d/50-eniatos-update
NEWS_URL=https://archlinux.org/feeds/news/

AUTO_REBOOT=yes
CRITICAL_UNITS="sshd.service"
NEWS_HOLD_DAYS=3
MAX_POOL_CAP=90
NO_RESTART_UNITS=""
# shellcheck source=/dev/null
[[ -r $CONF ]] && . "$CONF"

# units never restarted live - affected => reboot
NO_RESTART_BUILTIN="dbus.service dbus-broker.service systemd-logind.service user@*.service
getty@*.service serial-getty@*.service eniatos-update.service eniatos-update-health.service"

usage() {
    cat <<'EOF'
Usage: eniatos-update            run unattended upgrade (eniatos-update.timer)
       eniatos-update --check    show stale services, reboot need and boot gate; no changes
       eniatos-update --health   post-boot health check (eniatos-update-health.timer)
       eniatos-update --status   show last report
       eniatos-update --ack-news mark Arch News as read, release the update hold
Config: /etc/eniatos/update.conf
EOF
    exit "${1:-0}"
}

log() { echo "$*"; logger -t "$PROG" -p user.info -- "$*" 2>/dev/null || true; }
warn() { echo "WARNING: $*" >&2; logger -t "$PROG" -p user.warning -- "$*" 2>/dev/null || true; }
err() { echo "ERROR: $*" >&2; logger -t "$PROG" -p user.err -- "$*" 2>/dev/null || true; }

# --- report / motd ---------------------------------------------------------
# set_status STATUS MESSAGE...  - STATUS ok clears the login banner
set_status() {
    local status=$1; shift
    mkdir -p "$STATE"
    { echo "status=$status"; echo "time=$(date '+%Y-%m-%d %H:%M %Z')"; printf '%s\n' "$@"; } > "$STATE/status"
    local -a notes=()
    [[ $status != ok ]] && notes+=("eniatos-update [$status]: $*")
    local pacnew
    pacnew=$(find /etc -name '*.pacnew' 2>/dev/null | sort)
    [[ -n $pacnew ]] && notes+=("eniatos-update: .pacnew files to merge (pacdiff):" "$pacnew")
    if [[ ${#notes[@]} -gt 0 ]]; then
        mkdir -p "${MOTD%/*}"
        { echo; printf '%s\n' "${notes[@]}"; echo "Details: eniatos-update --status, journalctl -t $PROG"; echo; } > "$MOTD"
    else
        rm -f "$MOTD"
    fi
}

# --- Arch News hold --------------------------------------------------------
news_latest() {  # prints "EPOCH<TAB>TITLE" of the newest item
    local xml date title
    xml=$(curl -fsS --max-time 30 "$NEWS_URL") || return 1
    date=$(xmllint --xpath 'string(//item[1]/pubDate)' - <<< "$xml" 2>/dev/null) || return 1
    title=$(xmllint --xpath 'string(//item[1]/title)' - <<< "$xml" 2>/dev/null)
    [[ -n $date ]] || return 1
    printf '%s\t%s\n' "$(date -d "$date" +%s)" "$title"
}

# returns 0 when the update may proceed
news_gate() {
    local latest epoch title seen pending now
    if ! latest=$(news_latest); then
        warn "cannot fetch Arch News, proceeding"
        return 0
    fi
    IFS=$'\t' read -r epoch title <<< "$latest"
    mkdir -p "$STATE"
    if [[ ! -s $STATE/news-seen ]]; then
        echo "$epoch" > "$STATE/news-seen"   # first run: existing news counts as read
        return 0
    fi
    seen=$(< "$STATE/news-seen")
    if (( epoch <= seen )); then
        rm -f "$STATE/news-pending"
        return 0
    fi
    now=$(date +%s)
    [[ -s $STATE/news-pending ]] || echo "$now" > "$STATE/news-pending"
    pending=$(< "$STATE/news-pending")
    if (( now - pending >= NEWS_HOLD_DAYS * 86400 )); then
        warn "Arch News \"$title\" unread for $NEWS_HOLD_DAYS days, proceeding with the update"
        echo "$epoch" > "$STATE/news-seen"
        rm -f "$STATE/news-pending"
        return 0
    fi
    NEWS_TITLE=$title
    return 1
}

ack_news() {
    local latest epoch title
    latest=$(news_latest) || { err "cannot fetch Arch News"; exit 1; }
    IFS=$'\t' read -r epoch title <<< "$latest"
    mkdir -p "$STATE"
    echo "$epoch" > "$STATE/news-seen"
    rm -f "$STATE/news-pending"
    log "Arch News acknowledged up to: $title"
    [[ -s $STATE/status ]] && grep -q '^status=held' "$STATE/status" && set_status ok "news acknowledged"
}

# --- preconditions ---------------------------------------------------------
preflight() {
    local health cap
    health=$(zpool status -x 2>&1)
    [[ $health == "all pools are healthy" ]] || { echo "zpool not healthy: $health"; return 1; }
    while read -r pool cap; do
        cap=${cap%\%}
        (( cap < MAX_POOL_CAP )) || { echo "pool $pool is ${cap}% full"; return 1; }
    done < <(zpool list -H -o name,cap)
    [[ ! -e /var/lib/pacman/db.lck ]] || { echo "pacman database locked (/var/lib/pacman/db.lck)"; return 1; }
    return 0
}

# --- stale processes -------------------------------------------------------
# prints "PID<TAB>UNIT<TAB>COMM<TAB>LIBC" for processes mapping replaced files from /usr
stale_procs() {
    local p pid maps unit comm libc
    for p in /proc/[0-9]*; do
        pid=${p#/proc/}
        [[ $pid == "$$" ]] && continue
        maps=$(grep -E ' /usr/[^ ]* \(deleted\)$' "$p/maps" 2>/dev/null) || continue
        [[ -n $maps ]] || continue
        libc=0
        grep -qE '/libc\.so|/ld-linux' <<< "$maps" && libc=1
        comm=$(cat "$p/comm" 2>/dev/null) || comm="?"
        unit=$(awk -F: '$1 == 0 { print $3 }' "$p/cgroup" 2>/dev/null)
        if [[ $unit == /init.scope ]]; then
            unit=init.scope
        elif [[ $unit =~ ^/system\.slice/(.*/)?([^/]+\.service) ]]; then
            unit=${BASH_REMATCH[2]}
        else
            unit="-${unit}"   # not a system service (user session, scope)
        fi
        printf '%s\t%s\t%s\t%s\n' "$pid" "$unit" "$comm" "$libc"
    done
}

no_restart() {
    local u=$1 pat
    for pat in $NO_RESTART_BUILTIN $NO_RESTART_UNITS; do
        # shellcheck disable=SC2053
        [[ $u == $pat ]] && return 0
    done
    return 1
}

# fills REBOOT_REASONS, RESTART_UNITS, USER_PROCS
analyze() {
    REBOOT_REASONS=() RESTART_UNITS=() USER_PROCS=()
    local running
    running=$(uname -r)
    [[ -e /usr/lib/modules/$running/vmlinuz ]] || REBOOT_REASONS+=("kernel $running replaced")

    local stale pid unit comm libc
    stale=$(stale_procs)
    if grep -qP '\tinit\.scope\t' <<< "$stale"; then
        systemctl daemon-reexec && sleep 2
        stale=$(stale_procs)
    fi
    declare -A units=()
    while IFS=$'\t' read -r pid unit comm libc; do
        [[ -n $pid ]] || continue
        if [[ $libc == 1 ]]; then
            REBOOT_REASONS+=("glibc replaced (in use by $comm)")
        fi
        case $unit in
            init.scope) REBOOT_REASONS+=("systemd (PID 1) still runs replaced code") ;;
            -*) USER_PROCS+=("${comm}[${pid}] ${unit#-}") ;;
            *)
                if no_restart "$unit"; then
                    REBOOT_REASONS+=("$unit runs replaced code")
                else
                    units[$unit]=1
                fi ;;
        esac
    done <<< "$stale"
    RESTART_UNITS=("${!units[@]}")
    # dedupe reasons
    mapfile -t REBOOT_REASONS < <(printf '%s\n' "${REBOOT_REASONS[@]}" | awk 'NF && !s[$0]++')
}

# --- boot gate -------------------------------------------------------------
# returns 0 when every installed kernel looks bootable
boot_gate() {
    local ok=0 dir ver base vmlinuz initrd zv uv
    uv=$(pacman -Q zfs-utils 2>/dev/null | awk '{print $2}'); uv=${uv%%-*}
    for dir in /usr/lib/modules/*/; do
        dir=${dir%/}; ver=${dir##*/}
        [[ -r $dir/pkgbase && -e $dir/vmlinuz ]] || continue
        base=$(< "$dir/pkgbase")
        vmlinuz=/boot/vmlinuz-$base initrd=/boot/initramfs-$base.img
        if ! cmp -s "$dir/vmlinuz" "$vmlinuz"; then
            echo "GATE $base: $vmlinuz differs from $dir/vmlinuz"; ok=1; continue
        fi
        if [[ ! -s $initrd || $initrd -ot $vmlinuz ]]; then
            echo "GATE $base: $initrd missing or older than kernel"; ok=1; continue
        fi
        if ! lsinitrd "$initrd" 2>/dev/null | grep "modules/$ver/.*/zfs\.ko" >/dev/null; then
            echo "GATE $base: zfs.ko for $ver not in $initrd"; ok=1; continue
        fi
        zv=$(modinfo -k "$ver" -F version zfs 2>/dev/null); zv=${zv%%-*}
        if [[ -z $zv || $zv != "$uv" ]]; then
            echo "GATE $base: zfs module ${zv:-missing} != zfs-utils $uv"; ok=1; continue
        fi
        echo "gate $base ($ver): ok, zfs $zv"
    done
    return $ok
}

# --- main actions ----------------------------------------------------------
do_update() {
    exec 9> /run/eniatos-update.lock
    flock -n 9 || { err "already running"; exit 1; }
    log "=== eniatos-update $(date '+%F %T %Z') ==="

    local why
    if ! why=$(preflight); then
        err "preflight: $why"; set_status failed "preflight: $why"; exit 1
    fi
    NEWS_TITLE=""
    if ! news_gate; then
        warn "update held: unread Arch News \"$NEWS_TITLE\""
        set_status held "unread Arch News: \"$NEWS_TITLE\" - read https://archlinux.org/news/, then eniatos-update --ack-news (auto after $NEWS_HOLD_DAYS days)"
        exit 0
    fi

    log "--- keyring"
    pacman -Sy --needed --noconfirm archlinux-keyring cachyos-keyring 2>&1 \
        || { err "keyring update failed"; set_status failed "pacman -Sy keyring failed"; exit 1; }
    log "--- pacman -Su"
    local out rc
    out=$(pacman -Su --noconfirm 2>&1); rc=$?
    echo "$out"
    if [[ $rc -ne 0 ]]; then
        err "pacman -Su failed (rc=$rc)"
        set_status failed "pacman -Su failed (rc=$rc), system unchanged or partially upgraded - see journalctl -t $PROG"
        exit 1
    fi
    grep -q '^ there is nothing to do' <<< "$out" && log "nothing to update"

    analyze
    local u failed=() gate="" reboot=0
    if [[ ${#REBOOT_REASONS[@]} -gt 0 ]]; then
        log "reboot needed: ${REBOOT_REASONS[*]}"
        if [[ $AUTO_REBOOT == yes ]]; then
            if gate=$(boot_gate); then
                reboot=1
            else
                err "boot gate failed, not rebooting"
            fi
            echo "$gate"
        fi
    fi
    # restart services unless the machine reboots anyway
    if [[ $reboot -eq 0 ]]; then
        for u in "${RESTART_UNITS[@]}"; do
            if systemctl try-restart "$u"; then log "restarted $u"; else err "restart $u failed"; failed+=("$u"); fi
        done
    fi
    for u in "${USER_PROCS[@]}"; do warn "not a system service, still running replaced code: $u"; done

    if [[ $reboot -eq 1 ]]; then
        echo "reboot $(date +%s)" > "$STATE/reboot-pending"
        set_status rebooting "reboot after update: ${REBOOT_REASONS[*]}"
        log "rebooting"
        systemctl reboot --no-block
        exit 0
    fi
    if [[ ${#REBOOT_REASONS[@]} -gt 0 && $AUTO_REBOOT == yes ]]; then
        set_status failed "reboot needed (${REBOOT_REASONS[*]}) but boot gate failed:" "$gate"
        exit 1
    fi
    if [[ ${#REBOOT_REASONS[@]} -gt 0 ]]; then
        set_status reboot-needed "reboot needed: ${REBOOT_REASONS[*]}"
        exit 0
    fi

    if [[ ${#failed[@]} -gt 0 ]]; then
        set_status failed "service restart failed: ${failed[*]}"; exit 1
    fi
    set_status ok "update ok${RESTART_UNITS[*]:+, restarted: ${RESTART_UNITS[*]}}"
}

do_check() {
    analyze
    echo "reboot needed: ${REBOOT_REASONS[*]:-no}"
    echo "services to restart: ${RESTART_UNITS[*]:-none}"
    echo "non-service processes with replaced code: ${USER_PROCS[*]:-none}"
    boot_gate
}

do_health() {
    local problems=() u state after_update=""
    [[ -e $STATE/reboot-pending ]] && after_update=" (boot after update)"
    rm -f "$STATE/reboot-pending"
    state=$(systemctl is-system-running 2>/dev/null)
    [[ $state == running ]] || problems+=("system $state: $(systemctl --failed --no-legend --plain | awk '{print $1}' | xargs)")
    for u in $CRITICAL_UNITS; do
        systemctl is-active -q "$u" || problems+=("$u not active")
    done
    if [[ ${#problems[@]} -gt 0 ]]; then
        err "health check failed$after_update: ${problems[*]}"
        set_status unhealthy "health check failed$after_update: ${problems[*]}"
        exit 1
    fi
    log "health check ok$after_update, kernel $(uname -r)"
    # keep a held/failed status from the update run, only clear rebooting
    if [[ ! -s $STATE/status ]] || grep -qE '^status=(rebooting|unhealthy)' "$STATE/status"; then
        set_status ok "health check ok$after_update, kernel $(uname -r)"
    fi
}

[[ $(id -u) -eq 0 ]] || { echo "must be root" >&2; exit 1; }
case "${1:-}" in
    "") do_update ;;
    --check) do_check ;;
    --health) do_health ;;
    --status) cat "$STATE/status" 2>/dev/null || echo "no runs yet"; cat "$MOTD" 2>/dev/null ;;
    --ack-news) ack_news ;;
    -h|--help) usage 0 ;;
    *) usage 1 ;;
esac
