#!/bin/sh
# oomd-notify: pop a desktop notification whenever oomd kills processes.
# Runs as a SYSTEM service (root); tails oomd's system journal via
# journalctl --follow (event-driven, no polling, near-zero idle cost).
#
# oomd is the Meta/facebook build (/usr/bin/oomd), NOT systemd-oomd, so there
# is no org.freedesktop.oom1 Killed() D-Bus signal. Kill events are detected
# from oomd's OLOG lines in the journal. oomd reports each kill event as:
#   [BaseKillPlugin.cpp:554] Trying to kill /sys/fs/cgroup/<...>/<scope>
#   [BaseKillPlugin.cpp:668] Killed 20: 1937(ai.opencode.des) 2061(...) ...
#   [Ruleset.cpp:302] Action=kill_by_swap_usage returned STOP. ...
# so one notification is produced per kill event, titled with the app that
# owns the scope, and listing every process oomd killed inside it.

# --- options (not used by the systemd unit) ---
#   --stdin    read oomd log lines from stdin (for testing)
#   --dry-run  print what would be notified instead of notifying
STREAM="journalctl --follow --output=cat --unit=oomd --since now"
DRY=0
for arg in "$@"; do
    case "$arg" in
        --stdin) STREAM="cat" ;;
        --dry-run) DRY=1 ;;
    esac
done

# Print "user uid" of the first graphical (seated) user session, else "".
pick_user_session() {
    loginctl list-sessions --no-legend --no-pager 2>/dev/null |
        awk '$6 == "user" && $4 != "-" { print $3, $2; exit }'
}

notify_user() {
    # $1 = urgency, $2 = summary, $3 = body
    urgency="$1"
    summary="$2"
    body="$3"
    if [ "$DRY" = 1 ]; then
        printf '%s\n%s\n' "$summary" "$body"
        return 0
    fi
    session=$(pick_user_session)
    user=${session%% *}
    uid=${session##* }
    [ -n "$user" ] || return 1
    [ -S "/run/user/$uid/bus" ] || return 1
    runuser -u "$user" -- env \
        DBUS_SESSION_BUS_ADDRESS="unix:path=/run/user/$uid/bus" \
        notify-send -u "$urgency" -i dialog-warning "$summary" "$body"
}

# Derive the app name that owns a scope from its cgroup path, e.g.
#   .../app.slice/app-ai.opencode.desktop-1937.scope  -> ai.opencode.desktop
#   .../app.slice/run-p9611-i9084.scope               -> run-p9611-i9084
scope_app() {
    base=$(basename "$1")
    base=${base%.scope}
    case "$base" in
        *vte-spawn-* | run-* | session-*)
            # terminal-wrapper/transient/session scope; use process comm instead
            base=""
            ;;
        app-*)
            base=${base#app-}
            base=${base%-[0-9]*}
            ;;
    esac
    echo "$base"
}

# --- per-event state ---
pending=$(mktemp /tmp/oomd-notify.XXXXXX)   # one "pid comm" per line, for current event
cur_scope=""                                 # scope path of current event ("" = none)

flush_event() {
    [ -s "$pending" ] || return 0
    total=$(sort -u "$pending" | wc -l)
    # collapse to "count comm" lines, dedup pids
    body=$(sort -u "$pending" | awk '{ print $2 }' | sort | uniq -c |
        sort -rn | awk '{ printf "%s x %s\n", $1, $2 }' | head -12)
    app=$(scope_app "$cur_scope")
    if [ -z "$app" ]; then
        case "$cur_scope" in
            *vte-spawn-*)
                sorted=$(sort -u "$pending" | awk '{print $2}' | sort | uniq -c | sort -rn)
                main=$(echo "$sorted" | head -1 | awk '{print $2}')
                term=$(echo "$sorted" | awk '{print $2}' |
                    grep -E '^xfce4-terminal$|^gnome-terminal$|^konsole$|^alacritty$|^kitty$|^wezterm$|^terminator$|^lxterminal$|^mate-terminal$|^foot$|^st$|^ptyxis$' |
                    head -1)
                if [ -n "$term" ] && [ "$term" != "$main" ]; then
                    app="$main and $term"
                else
                    app="$main"
                fi
                ;;
            *)
                app=$(sort -u "$pending" | awk '{print $2}' | sort | uniq -c |
                    sort -rn | head -1 | awk '{print $2}')
                ;;
        esac
    fi
    [ -n "$app" ] || app="unknown"
    summary="oomd stopped $app"
    body="$app was using too much memory and was killed. Your other apps are fine."
    notify_user critical "$summary" "$body"
    : > "$pending"
    cur_scope=""
}

trap 'flush_event; rm -f "$pending"' EXIT INT TERM

eval "$STREAM" | while IFS= read -r line; do
    case "$line" in
        *"Trying to kill "*)     # new kill event begins
            flush_event
            cur_scope=$(echo "$line" | sed 's/.*Trying to kill //')
            : > "$pending"
            ;;
        *"Killed "*)             # oomd reports killed pids: "Killed 20: pid(comm) ..."
            rest=$(echo "$line" | sed 's/.*Killed [0-9]*: //')
            echo "$rest" | grep -oE '[0-9]+\([^)]*\)' |
                sed -E 's/([0-9]+)\(([^)]*)\)/\1 \2/' >> "$pending"
            [ -n "$cur_scope" ] || cur_scope="unknown"
            ;;
        *"returned STOP"* | *"Terminating action chain"*)  # kill event ends
            flush_event
            ;;
    esac
done
