#!/usr/bin/env python3
"""Evacuate processes from a pids-wedged cgroup, then kill it.

A cgroup at its pids.max denies fork() to everything inside it, so nothing in
it can repair itself. This runs from a HEALTHY cgroup (a terminal in
human.slice can still fork) and touches the wedged one through file I/O only:
write the pids to keep into a fresh leaf, read both cgroup.procs back to prove
they moved, then write cgroup.kill on the wedge.

Fail-closed: the target must be an agent-owned scope, the move must be
confirmed, and an unconfirmed move refuses the kill.

usage:
  pids-rescue <cgroup-path> [--keep PID[,PID...]] [--keep-comm NAME[,NAME...]]
              [--rescue-slice human.slice] [--dry-run]

<cgroup-path> is an absolute /sys/fs/cgroup path or a cgroup-root-relative one.
With no --keep/--keep-comm nothing is evacuated and the cgroup is killed whole.
"""
import argparse
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "lib"))
from pids_cgroup import is_killable  # noqa: E402

CG_ROOT = "/sys/fs/cgroup"
UID = os.getuid()


def normalize(path):
    """Absolute cgroupfs path and cgroup-root-relative path for one argument."""
    p = os.path.normpath(path)
    if p.startswith(CG_ROOT):
        return p, p[len(CG_ROOT):] or "/"
    if not p.startswith("/"):
        return None, None
    return CG_ROOT + p, p


def read_pids(path):
    try:
        with open(os.path.join(path, "cgroup.procs")) as f:
            return [int(line) for line in f if line.strip()]
    except (OSError, ValueError):
        return []


def comm(pid):
    try:
        with open(f"/proc/{pid}/comm") as f:
            return f.read().strip()
    except OSError:
        return ""


def select_keep(pids, keep_pids, keep_comms):
    """Pids to evacuate: an explicit pid, or a process whose comm was named."""
    return [p for p in pids if p in keep_pids or (keep_comms and comm(p) in keep_comms)]


def wedged(path):
    """True when this cgroup has already denied a fork (pids.events.local max)."""
    for name in ("pids.events.local", "pids.events"):
        try:
            with open(os.path.join(path, name)) as f:
                for line in f:
                    if line.startswith("max ") and int(line.split()[1]) > 0:
                        return True
        except (OSError, ValueError):
            continue
    return False


def move_pids(dest, pids):
    """Write pids one per line. Returns the pids the kernel accepted."""
    moved = []
    procs = os.path.join(dest, "cgroup.procs")
    for pid in pids:
        try:
            with open(procs, "w") as f:
                f.write(str(pid))
            moved.append(pid)
        except OSError:
            if not os.path.exists(f"/proc/{pid}"):
                continue  # exited on its own; nothing left to rescue
    return moved


def verify_move(wanted, source_pids, dest_pids):
    """Every wanted pid is in dest and none is left in source."""
    live = [p for p in wanted if os.path.exists(f"/proc/{p}")]
    return all(p in dest_pids for p in live) and not any(p in source_pids for p in live)


def main():
    ap = argparse.ArgumentParser(add_help=True)
    ap.add_argument("cgroup")
    ap.add_argument("--keep", default="")
    ap.add_argument("--keep-comm", default="")
    ap.add_argument("--rescue-slice", default="human.slice")
    ap.add_argument("--dry-run", action="store_true")
    ns = ap.parse_args()

    path, rel = normalize(ns.cgroup)
    if path is None:
        print(f"pids-rescue: not a cgroup path: {ns.cgroup}", file=sys.stderr)
        return 2
    if not is_killable(rel):
        print(
            f"pids-rescue: refusing {rel} — not a provably agent-owned scope.\n"
            "Only scopes under agent.slice/agent-seat.slice/build.slice named by an "
            "agent launcher may be killed. Human sessions are never a target.",
            file=sys.stderr,
        )
        return 3
    if not os.path.isdir(path):
        print(f"pids-rescue: no such cgroup: {path}", file=sys.stderr)
        return 2

    keep_pids = {int(p) for p in ns.keep.split(",") if p.strip()}
    keep_comms = {c for c in ns.keep_comm.split(",") if c.strip()}
    source = read_pids(path)
    wanted = select_keep(source, keep_pids, keep_comms)

    rescue = os.path.join(
        CG_ROOT, f"user.slice/user-{UID}.slice/user@{UID}.service",
        ns.rescue_slice, f"pids-rescue-{int(time.time())}",
    )

    print(f"cgroup   {rel}")
    print(f"wedged   {wedged(path)}  (pids.events max counter)")
    print(f"pids     {len(source)} present, {len(wanted)} to evacuate")
    if ns.dry_run:
        print(f"rescue   {rescue} (not created)")
        print("dry-run: nothing moved, nothing killed")
        return 0

    if wanted:
        try:
            os.makedirs(rescue, exist_ok=True)
        except OSError as err:
            print(f"pids-rescue: cannot create {rescue}: {err}", file=sys.stderr)
            return 4
        moved = move_pids(rescue, wanted)
        if not verify_move(moved, read_pids(path), read_pids(rescue)):
            print(
                "pids-rescue: move not confirmed by reading both cgroup.procs back; "
                "refusing to kill",
                file=sys.stderr,
            )
            return 5
        print(f"moved    {len(moved)} pids -> {rescue}")

    try:
        with open(os.path.join(path, "cgroup.kill"), "w") as f:
            f.write("1")
    except OSError as err:
        print(f"pids-rescue: cgroup.kill failed: {err}", file=sys.stderr)
        return 6
    print(f"killed   {rel}")
    return 0


if __name__ == "__main__":
    sys.exit(main())