#!/usr/bin/env python3
"""Exhaustion-trajectory memory watcher.

Alerts on projected time-to-exhaustion (runway falling toward a floor), not on
process size or raw growth. Names the fastest-growing group as the culprit.
Plateauing work never fires; a sustained runaway fires with lead time.
"""
import os, glob, time, subprocess
from collections import deque, defaultdict

def env_int(k, d):
    try: return int(os.environ.get(k, d))
    except ValueError: return d

POLL          = env_int("POLL", 5)             # seconds between scans
RUNWAY_WINDOW = env_int("RUNWAY_WINDOW", 180)  # runway history kept (s)
RATE_WINDOW   = env_int("RATE_WINDOW", 90)     # slope measured over recent (s)
CRIT_ETA_S    = env_int("CRIT_ETA_S", 90)      # critical if ETA < this (90 s)
FLOOR_PCT     = env_int("FLOOR_PCT", 8)        # "exhausted" = runway below this % of RAM+swap
HARD_AVAIL    = env_int("HARD_AVAIL_PCT", 8)   # critical if free RAM below this % regardless
CRIT_PSI      = env_int("CRIT_PSI", 20)        # critical if PSI full avg10 >= this %
GROW_MIN_KB   = env_int("GROW_MIN_KB", 262144) # a group must have grown >=256M to be "the grower"
MIN_GROUP_KB  = env_int("MIN_GROUP_KB", 524288)# ignore groups under 512M when attributing
GROW_WINDOW   = env_int("GROW_WINDOW", 120)    # per-group growth measured over (s)

UID_ME = os.getuid()
STATE  = os.path.join(os.environ.get("XDG_STATE_HOME", os.path.expanduser("~/.local/state")), "mem-guard")
os.makedirs(STATE, exist_ok=True)
LOG = os.path.join(STATE, "events.log")

PROTECT = {
    "systemd","systemd-oomd","pipewire","pipewire-pulse","wireplumber",
    "dbus-daemon","dbus-broker","gnome-shell","plasmashell","kwin_x11",
    "kwin_wayland","mutter","Xorg","Xwayland","gnome-session","xfwm4",
    "xfce4-session","xfsettingsd","lightdm","sddm","gdm",
    "gnome-keyring-daemon","polkitd","bash","notify-send","python3",
}

# Known-heavy processes the user runs on purpose (VMs, builders, agents, runtimes).
# Excluded from culprit attribution and the kill action, so their expected memory
# use is never alerted on. A real OOM caused only by these stays silent (systemd-oomd
# is the safety net); an UNEXPECTED process driving exhaustion still alerts.
IGNORE = {
    "qemu-system-x86_64","qemu-system-aarch64","qemu-system-i386","qemu",
    "node","nodejs","deno","bun","electron",
    "esbuild","rollup","webpack","tsc","tsserver","vite","turbo","next",
    "cargo","rustc","go","gopls","cc","clang","gcc","g++","ld","ld.lld",
    "make","ninja","cmake","java","dotnet","chromium","chrome",
    "chrome-headless-shell","docker","dockerd","containerd","podman",
}
IGNORE |= {s.strip() for s in os.environ.get("IGNORE_EXTRA", "").split(",") if s.strip()}

_SKIP = PROTECT | IGNORE

def is_skipped(key):
    # /proc/PID/status Name is capped at 15 chars; when readlink(exe) fails the
    # fallback key is truncated, so match a 15-char key as a prefix of any full name.
    if key in _SKIP:
        return True
    if len(key) == 15:
        return any(n.startswith(key) for n in _SKIP)
    return False

HAVE_NOTIFY = subprocess.run(["sh","-c","command -v notify-send"],
                             capture_output=True).returncode == 0
HAVE_WAIT = False
if HAVE_NOTIFY:
    h = subprocess.run(["notify-send","--help"], capture_output=True, text=True)
    HAVE_WAIT = "--wait" in (h.stdout + h.stderr)

def logline(s):
    with open(LOG, "a") as f:
        f.write(time.strftime("%F %T") + "  " + s + "\n")

def read_meminfo():
    mt=ma=st=sf=0
    with open("/proc/meminfo") as f:
        for ln in f:
            k,_,v = ln.partition(":")
            v = int(v.strip().split()[0])
            if   k=="MemTotal": mt=v
            elif k=="MemAvailable": ma=v
            elif k=="SwapTotal": st=v
            elif k=="SwapFree": sf=v
    avail_pct = 100*ma//mt if mt else 100
    swap_pct  = 100*(st-sf)//st if st else 0
    return avail_pct, swap_pct, ma, sf, mt, st

def psi_full_avg10():
    try:
        with open("/proc/pressure/memory") as f:
            for ln in f:
                if ln.startswith("full"):
                    for tok in ln.split():
                        if tok.startswith("avg10="):
                            return float(tok[6:])
    except OSError: pass
    return 0.0

def scan_groups():
    groups = defaultdict(lambda: {"kb":0, "pids":[]})
    for d in glob.glob("/proc/[0-9]*"):
        pid = d[6:]
        try:
            name=""; uid=-1; rss=0; swap=0
            with open(d+"/status") as f:
                for ln in f:
                    if   ln.startswith("Name:"):  name = ln[5:].strip()
                    elif ln.startswith("Uid:"):   uid  = int(ln.split()[1])
                    elif ln.startswith("VmRSS:"): rss  = int(ln.split()[1])
                    elif ln.startswith("VmSwap:"):swap = int(ln.split()[1])
            if uid != UID_ME:
                continue
            try:
                key = os.path.basename(os.readlink(d+"/exe"))
                if key.endswith(" (deleted)"): key = key[:-10]
            except OSError:
                key = name
            if not key or is_skipped(key):
                continue
            g = groups[key]
            g["kb"] += rss + swap
            g["pids"].append(pid)
        except (OSError, ValueError):
            continue
    return groups

hist   = defaultdict(deque)   # key -> deque[(ts, kb)]  per-group, for attribution
runway = deque()              # (ts, eff_kb) system runway, for the trigger
alerted = set()   # groups already alerted on, until pressure clears

def group_growth(key, now):
    dq = hist[key]
    while dq and now - dq[0][0] > GROW_WINDOW:
        dq.popleft()
    if len(dq) < 2:
        return 0
    return dq[-1][1] - min(k for _,k in dq)

def expire_history(now):
    """Drop samples past the growth window and forget groups left with none.
    rank_growers reaches group_growth only for groups in the current scan that
    clear MIN_GROUP_KB, so every other key would retain its samples forever."""
    for key in list(hist):
        dq = hist[key]
        while dq and now - dq[0][0] > GROW_WINDOW:
            dq.popleft()
        if not dq:
            del hist[key]
            alerted.discard(key)

def rank_growers(now, groups):
    out=[]
    for k, gr in groups.items():
        if gr["kb"] < MIN_GROUP_KB:
            continue
        d = group_growth(k, now)
        if d >= GROW_MIN_KB:
            out.append((k, d, gr["kb"]))
    out.sort(key=lambda x: -x[1])
    return out

def biggest(groups):
    return [(k, 0, v["kb"]) for k,v in
            sorted(groups.items(), key=lambda kv: -kv[1]["kb"])]

def runway_eta(now, eff_kb, floor_kb):
    """Projected seconds to hit floor at the recent burn rate. None if not falling."""
    while runway and now - runway[0][0] > RUNWAY_WINDOW:
        runway.popleft()
    # pick the sample ~RATE_WINDOW ago (recent slope, so a plateau reads ~flat)
    ref = None
    for ts, kb in runway:
        if now - ts <= RATE_WINDOW:
            ref = (ts, kb); break
    if ref is None and runway:
        ref = runway[0]
    if ref is None or now - ref[0] < 20:
        return None
    burn = (ref[1] - eff_kb) / (now - ref[0])   # kB/s, positive = shrinking
    if burn <= 0:
        return None
    return max(0.0, (eff_kb - floor_kb)) / burn

def g(kb): return f"{kb/1048576:.1f}"

def notify(tier, urg, reason, ranked, top_key, top_pids):
    body = [reason]
    if ranked:
        body.append("Culprit (fastest-growing):")
        for k, d, cur in ranked[:3]:
            tag = f"+{g(d)}G, " if d else ""
            body.append(f"{k}  ({tag}now {g(cur)}G)")
    logline(f"{tier}  {reason}  top={top_key}")
    if not HAVE_NOTIFY:
        return
    summary = f"Memory {tier} — {top_key}"
    if HAVE_WAIT and top_pids:
        r = subprocess.run(
            ["notify-send","--wait","-u",urg,"-t","20000",
             "-A", f"kill=Kill {top_key} ({len(top_pids)} procs)",
             "-A", "dismiss=Dismiss",
             summary, "\n".join(body)],
            capture_output=True, text=True)
        if r.stdout.strip() == "kill":
            kill_pids(top_key, top_pids)
    else:
        subprocess.run(["notify-send","-u",urg,summary,"\n".join(body)])

def kill_pids(label, pids):
    for p in pids:
        try: os.kill(int(p), 15)
        except (OSError, ValueError): pass
    alive = pids
    for _ in range(3):
        time.sleep(1)
        alive = [p for p in pids if os.path.isdir(f"/proc/{p}")]
        if not alive: break
    for p in alive:
        try: os.kill(int(p), 9)
        except (OSError, ValueError): pass
    logline(f"KILLED {label} pids={' '.join(pids)}")
    if HAVE_NOTIFY:
        subprocess.run(["notify-send","-u","normal",f"Killed {label}",
                        f"Freed memory from {label}."])

def main():
    while True:
        now = time.time()
        groups = scan_groups()
        for k, gr in groups.items():
            hist[k].append((now, gr["kb"]))
        expire_history(now)
        avail, swap, ma, sf, mt, st = read_meminfo()
        eff  = ma + sf
        floor = (mt + st) * FLOOR_PCT // 100
        runway.append((now, eff))
        psi = psi_full_avg10()
        eta = runway_eta(now, eff, floor)

        growers = rank_growers(now, groups)

        tier = urg = reason = None
        if avail < HARD_AVAIL or (eta is not None and eta < CRIT_ETA_S):
            tier, urg = "CRITICAL", "critical"
            if eta is not None and eta < CRIT_ETA_S:
                reason = f"~{int(eta)}s to exhaustion at current rate — {avail}% RAM free (PSI {psi:.0f}%)."
            else:
                reason = f"Only {avail}% RAM free (PSI {psi:.0f}%)."

        # Genuine runaway only: real exhaustion AND a non-ignored process actively
        # growing toward it. Abundant RAM with high PSI (normal swap paging under
        # VM/agent load) is not an emergency and stays silent.
        if tier and growers:
            top_key  = growers[0][0]
            top_pids = groups.get(top_key, {}).get("pids", [])
            if top_key not in alerted:
                alerted.add(top_key)
                notify(tier, urg, reason, growers, top_key, top_pids)
        elif not tier:
            alerted.clear()

        time.sleep(POLL)

if __name__ == "__main__":
    try: main()
    except KeyboardInterrupt: pass