#!/usr/bin/env python3
"""Kill only the processes opencode spawns, never opencode itself.

Logic:
  roots   = processes whose comm starts with "opencode"
  workers = every descendant of those roots (opencode itself excluded)
  total   = summed RSS of workers
  while total > LIMIT: kill the direct child of opencode holding the most RSS
                       (whole subtree, so the runaway program A dies with its kids)
"""
import json
import os
import re
import signal
import sys
import time

LIMIT = float(sys.argv[1] if len(sys.argv) > 1 else 80) * 1024 ** 3
INTERVAL = float(sys.argv[2] if len(sys.argv) > 2 else 2)
PER_PROC = float(sys.argv[3] if len(sys.argv) > 3 else 64) * 1024 ** 3
MARK = "/tmp/memkill.last.json"
HIST = "/tmp/memkill.jsonl"
PAGE = os.sysconf("SC_PAGE_SIZE")


def read_cmd(p):
    try:
        with open("/proc/%d/cmdline" % p, "rb") as f:
            return f.read().replace(b"\0", b" ").decode(errors="replace").strip()
    except Exception:
        return ""


def snapshot():
    ppid = {}
    rss = {}
    comm = {}
    for d in os.listdir("/proc"):
        if not d.isdigit():
            continue
        p = int(d)
        try:
            with open("/proc/%d/stat" % p) as f:
                s = f.read()
            ppid[p] = int(s[s.rindex(")") + 2:].split()[1])
            with open("/proc/%d/statm" % p) as f:
                rss[p] = int(f.read().split()[1]) * PAGE
            with open("/proc/%d/comm" % p) as f:
                comm[p] = f.read().strip()
        except Exception:
            continue
    return ppid, rss, comm


def protect():
    """never let the guard (or its own helpers) become a victim"""
    out = {os.getpid()}
    for d in os.listdir("/proc"):
        if d.isdigit():
            try:
                with open("/proc/%s/cmdline" % d, "rb") as f:
                    if b"memkill.py" in f.read():
                        out.add(int(d))
            except Exception:
                pass
    return out


def main():
    try:
        with open("/tmp/.memguard.pid", "w") as f:
            f.write(str(os.getpid()))
    except Exception:
        pass
    print("memkill: guard armed  total_limit=%.0fGiB  per_proc_limit=%.0fGiB"
          % (LIMIT / 1024 ** 3, PER_PROC / 1024 ** 3), flush=True)
    while True:
        ppid, rss, comm = snapshot()
        keep = protect()
        roots = {p for p, c in comm.items() if c.startswith("opencode")} - keep
        kids = {}
        for child, parent in ppid.items():
            kids.setdefault(parent, []).append(child)

        # subtree walk helpers
        def sub(p):
            out = {p}
            stack = [p]
            while stack:
                for k in kids.get(stack.pop(), ()):
                    if k not in out:
                        out.add(k)
                        stack.append(k)
            return out

        workers = set()
        top = []                      # direct children of opencode
        for r in roots:
            for k in kids.get(r, ()):
                if k not in workers:
                    workers |= sub(k)
                    top.append(k)
        workers -= roots | keep

        # threshold counts opencode itself too, so "opencode + everything it
        # started" stays under LIMIT -- but victims are only ever children
        total = sum(rss.get(p, 0) for p in workers | roots)
        print("%s opencode=%.2fGiB workers=%d total=%.2fGiB"
              % (time.strftime("%H:%M:%S"),
                 sum(rss.get(p, 0) for p in roots) / 1024 ** 3,
                 len(workers), total / 1024 ** 3), flush=True)

        victims = set()
        top_set = set(top)
        if total > LIMIT:
            while total > LIMIT and top:
                c = max(top, key=lambda p: sum(rss.get(x, 0) for x in sub(p)))
                if c in victims:
                    break
                victims |= sub(c)
                top.remove(c)
                total -= sum(rss.get(x, 0) for x in victims)
        if PER_PROC:
            for p in list(workers):
                if rss.get(p, 0) > PER_PROC:
                    victims |= sub(p)

        doomed = victims - roots - keep
        cmds = {v: read_cmd(v) for v in doomed}
        for v in sorted(doomed):
            try:
                os.kill(v, signal.SIGKILL)
                print("  KILL pid=%d rss=%.2fGiB (opencode untouched)  %s"
                      % (v, rss.get(v, 0) / 1024 ** 3, cmds[v][:90]), flush=True)
            except Exception as e:
                print("  skip pid=%d: %s" % (v, e), flush=True)

        if doomed:
            # the direct child of opencode == the command the assistant ran
            heads = [v for v in doomed if v in top_set] or sorted(doomed)
            rec = {
                "ts": round(time.time(), 3),
                "iso": time.strftime("%Y-%m-%d %H:%M:%S"),
                "head_pid": heads[0],
                "head_cmd": cmds.get(heads[0], ""),
                "cmds": [cmds[v] for v in sorted(doomed)][:8],
                "rss_gib": round(max(rss.get(v, 0) for v in doomed) / 1024 ** 3, 2),
                "total_gib": round(sum(rss.get(v, 0) for v in doomed) / 1024 ** 3, 2),
                "limit_gib": round(LIMIT / 1024 ** 3, 2),
                "per_proc_gib": round(PER_PROC / 1024 ** 3, 2),
                "reason": "over-limit",
            }
            try:
                with open(MARK + ".tmp", "w") as f:
                    json.dump(rec, f)
                os.replace(MARK + ".tmp", MARK)
                with open(HIST, "a") as f:
                    f.write(json.dumps(rec) + "\n")
            except Exception as e:
                print("  marker failed: %s" % e, flush=True)
            time.sleep(3)
        time.sleep(INTERVAL)


if __name__ == "__main__":
    main()