#!/bin/bash
set -u

# ---------------------------------------------------------------------------
# Config: defaults, overridden by the config file, then by the argument
# ---------------------------------------------------------------------------
THRESHOLD=30
AUR_REPO="aur-local"
# stock kernel package -> your custom build of it; none by default.
# Declared here so a [stock]=custom list in the config stays associative.
declare -A KERNELS=()
PACNEW_CHECK=1  # remind about .pacnew/.pacsave files (0 to disable)
REBOOT_CHECK=1  # remind to reboot after a kernel upgrade (0 to disable)
config_file="${XDG_CONFIG_HOME:-$HOME/.config}/update-notify/config"
state_dir="${XDG_STATE_HOME:-$HOME/.local/state}/update-notify"

if [[ -r $config_file ]]; then
    # shellcheck source=/dev/null
    source "$config_file"
fi
quiet=0  # refresh the pending file without notifying (path unit)
if [[ ${1:-} == -q || ${1:-} == --quiet ]]; then
    quiet=1
    shift
fi
config_threshold=$THRESHOLD  # written to the pending file, so a one-off argument doesn't change status bars
THRESHOLD=${1:-$THRESHOLD}

for t in "$config_threshold" "$THRESHOLD"; do
    [[ $t =~ ^[0-9]+$ ]] || { echo "Threshold must be a number" >&2; exit 1; }
done
# force base 10, so e.g. "08" isn't read as octal
config_threshold=$((10#$config_threshold))
THRESHOLD=$((10#$THRESHOLD))
mkdir -p "$state_dir"

# notify-send from cron/systemd needs the session bus address
export DBUS_SESSION_BUS_ADDRESS="${DBUS_SESSION_BUS_ADDRESS:-unix:path=/run/user/$(id -u)/bus}"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# If $1 has a pending update, print its new version and return 0.
# checkupdates lines look like: "pkgname oldver -> newver"
pending_version() {
    local pkg=$1 item new
    for item in "${arch_list[@]}"; do
        if [[ $item == "$pkg "* ]]; then
            read -r _ _ _ new <<< "$item"
            printf '%s\n' "$new"
            return 0
        fi
    done
    return 1
}

# Newest known version of $1: the pending version if an update is queued,
# otherwise the installed version. Returns 1 if the package isn't installed.
latest_version() {
    local pkg=$1 ver
    if ver=$(pending_version "$pkg"); then
        printf '%s\n' "$ver"
    elif ver=$(pacman -Q "$pkg" 2>/dev/null); then
        printf '%s\n' "${ver#* }"
    else
        return 1
    fi
}

# ---------------------------------------------------------------------------
# Gather updates
# ---------------------------------------------------------------------------

# Fetch checkupdates once. Exit codes: 0 = updates, 2 = none, 1 = error.
checkupdates_out=$(checkupdates 2>/dev/null)
checkupdates_rc=$?
if (( checkupdates_rc == 1 )); then
    # Bail out without touching the state, so the next good run doesn't re-notify
    echo "checkupdates failed (network or sync problem?)" >&2
    exit 1
fi

arch_list=()
if [[ -n $checkupdates_out ]]; then
    mapfile -t arch_list <<< "$checkupdates_out"
fi
arch_updates=${#arch_list[@]}

aur_list=()
if command -v aur >/dev/null 2>&1; then
    if aur_out=$(aur repo -u -d "$AUR_REPO"); then
        [[ -n $aur_out ]] && mapfile -t aur_list <<< "$aur_out"
    else
        echo "aur repo -u failed, AUR updates not counted" >&2
    fi
fi
aur_updates=${#aur_list[@]}
updates=$((arch_updates + aur_updates))

# ---------------------------------------------------------------------------
# Kernel check: the stock kernel updated, but your custom build didn't
# ---------------------------------------------------------------------------
kernel_lines=()
kernel_pending=()
for stock_pkg in "${!KERNELS[@]}"; do
    custom_pkg=${KERNELS[$stock_pkg]}

    # Skip kernels where either package isn't installed
    if ! stock_ver=$(latest_version "$stock_pkg") || ! custom_ver=$(latest_version "$custom_pkg"); then
        continue
    fi

    # Flag the kernel while the stock one is newer than the newest custom build
    # (installed, or already built into your repo but not yet installed).
    # pkgrel is ignored: a stock packaging-only bump doesn't need a rebuild.
    if (( $(vercmp "${stock_ver%-*}" "${custom_ver%-*}") > 0 )); then
        kernel_lines+=("Rebuild $custom_pkg ($stock_pkg $stock_ver)")
        kernel_pending+=("$custom_pkg $stock_ver")
    fi
done

# ---------------------------------------------------------------------------
# Maintenance: config files to merge, reboot after a kernel upgrade.
# maint_items are tracked in the state and the pending file, maint_lines are shown.
# ---------------------------------------------------------------------------
maint_items=()
maint_lines=()
if [[ $PACNEW_CHECK != 0 ]]; then
    mapfile -t pacnew_files < <(pacdiff -o 2>/dev/null)
    if (( ${#pacnew_files[@]} )); then
        maint_items+=("${pacnew_files[@]/#/pacnew }")
        maint_lines+=("Config files to merge with pacdiff: ${#pacnew_files[@]}")
    fi
fi
running_kernel=$(uname -r)
modules_dir=${UPDATE_NOTIFY_MODULES_DIR:-/usr/lib/modules}  # override for tests (build chroots have no kernel)
# vmlinuz belongs to the kernel package, while DKMS files can keep the folder around
if [[ $REBOOT_CHECK != 0 && ! -e $modules_dir/$running_kernel/vmlinuz ]]; then
    maint_items+=("reboot $running_kernel")
    maint_lines+=("Reboot to load the new kernel (running $running_kernel)")
fi

# ---------------------------------------------------------------------------
# Pending list for status bars, one "<source> <pkg> <newver>" per line
# ---------------------------------------------------------------------------
{
    printf 'threshold %s\n' "$config_threshold"
    for item in "${arch_list[@]}"; do
        read -r pkg _ _ new <<< "$item"
        printf 'repo %s %s\n' "$pkg" "$new"
    done
    for item in "${aur_list[@]}"; do
        read -r pkg _ _ new <<< "$item"
        printf 'aur %s %s\n' "$pkg" "$new"
    done
    for item in "${kernel_pending[@]}"; do
        printf 'kernel %s\n' "$item"
    done
    for item in "${maint_items[@]}"; do
        printf '%s\n' "$item"
    done
} > "$state_dir/pending.tmp" && mv "$state_dir/pending.tmp" "$state_dir/pending"  # readers never see a partial file

# ---------------------------------------------------------------------------
# Build a single notification
#   - kernel rebuild lines are always included
#   - update counts and reminders only when the threshold is reached
# ---------------------------------------------------------------------------
over_threshold=0
(( updates >= THRESHOLD )) && over_threshold=1

parts=()
if (( over_threshold )); then
    parts+=("Arch: $arch_updates, AUR: $aur_updates")
fi
parts+=("${kernel_lines[@]}")

if (( over_threshold )); then
    reminders=()

    # Always remind about git packages if git-monitor is installed
    if command -v git-monitor >/dev/null 2>&1; then
        reminders+=("git-monitor")
    fi

    # Firmware: only mention it when updates actually exist.
    # fwupdmgr get-updates exits 0 when updates are available, 2 when there are none.
    # --no-metadata-check avoids prompting/refreshing; the fwupd-refresh timer keeps metadata fresh.
    if command -v fwupdmgr >/dev/null 2>&1; then
        if timeout 30 fwupdmgr get-updates --no-metadata-check </dev/null >/dev/null 2>&1; then
            reminders+=("fwupdmgr (firmware updates available)")
        fi
    fi

    if (( ${#reminders[@]} )); then
        printf -v list '%s, ' "${reminders[@]}"
        parts+=("Also check: ${list%, }")
    fi
fi

# ---------------------------------------------------------------------------
# Notify when the update part changed since the last notification (maintenance
# lines included), or for maintenance alone when it has items not shown before
# ---------------------------------------------------------------------------
main_file="$state_dir/updates"
main_body=""
main_changed=0
if (( ${#parts[@]} )); then
    printf -v main_body '%s\n' "${parts[@]}"
    main_body=${main_body%$'\n'}
    [[ -f $main_file && $(<"$main_file") == "$main_body" ]] || main_changed=1
else
    rm -f "$main_file"  # so the next occurrence notifies again
fi

maint_file="$state_dir/maintenance"
declare -A maint_seen=()
if [[ -f $maint_file ]]; then
    while IFS= read -r item; do
        maint_seen["$item"]=1
    done < "$maint_file"
fi
maint_new=0
maint_kept=()  # current items already shown
for item in "${maint_items[@]}"; do
    if [[ -n ${maint_seen["$item"]:-} ]]; then
        maint_kept+=("$item")
    else
        maint_new=1
    fi
done

lines=()
if (( main_changed )); then
    lines=("${parts[@]}" "${maint_lines[@]}")
    title="Updates available"
    (( over_threshold )) || title="Kernel update available"
elif (( maint_new )); then
    lines=("${maint_lines[@]}")
    title="System maintenance"
fi

shown=0
if (( ${#lines[@]} )); then
    printf -v body '%s\n' "${lines[@]}"
    if (( quiet )); then
        # Track changes to a shown notification, so the timer doesn't re-notify them;
        # leave a new one for the timer to show
        (( main_changed )) && [[ -f $main_file ]] && printf '%s' "$main_body" > "$main_file"
    elif notify-send "$title" "${body%$'\n'}"; then
        shown=1
        [[ -n $main_body ]] && printf '%s' "$main_body" > "$main_file"
    fi
fi

# Remember the maintenance items shown: drop resolved ones, keep new ones for the next notification
maint_save=("${maint_kept[@]}")
(( shown )) && maint_save=("${maint_items[@]}")
if (( ${#maint_save[@]} )); then
    printf '%s\n' "${maint_save[@]}" > "$maint_file"
else
    rm -f "$maint_file"
fi
