#!/usr/bin/env python3
"""Sample process-tree RSS/swap while running a command."""

import argparse
import json
import signal
import subprocess
import sys
import time


def process_tree(root):
    output = subprocess.check_output(["ps", "-e", "-o", "pid=,ppid="], text=True)
    children = {}
    for line in output.splitlines():
        pid, ppid = map(int, line.split())
        children.setdefault(ppid, []).append(pid)
    result, todo = [], [root]
    while todo:
        pid = todo.pop()
        result.append(pid)
        todo.extend(children.get(pid, ()))
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--interval", type=float, default=0.25)
    parser.add_argument("--timeout", type=float, default=7200)
    parser.add_argument("--rss-limit-gib", type=float)
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    command = args.command[1:] if args.command[:1] == ["--"] else args.command
    if not command:
        parser.error("missing command")

    started = time.monotonic()
    process = subprocess.Popen(command)
    peak_rss = peak_swap = 0
    exceeded = False
    while process.poll() is None:
        rss = swap = 0
        for pid in process_tree(process.pid):
            try:
                with open(f"/proc/{pid}/status", encoding="ascii") as status:
                    for line in status:
                        if line.startswith("VmRSS:"):
                            rss += int(line.split()[1])
                        elif line.startswith("VmSwap:"):
                            swap += int(line.split()[1])
            except (FileNotFoundError, ProcessLookupError):
                pass
        peak_rss = max(peak_rss, rss)
        peak_swap = max(peak_swap, swap)
        # VmRSS 以 kB 為單位，門檻須換算成 kB（1 GiB = 1024 * 1024 kB）
        if args.rss_limit_gib is not None and rss > args.rss_limit_gib * (1024 ** 2):
            exceeded = True
            process.send_signal(signal.SIGKILL)
            break
        if time.monotonic() - started > args.timeout:
            process.kill()
            break
        time.sleep(args.interval)
    result = process.wait()
    print("RUNMEM " + json.dumps({
        "peak_kb": peak_rss,
        "peak_gib": round(peak_rss / 1048576, 3),
        "limit_gib": args.rss_limit_gib,
        "limit_exceeded": exceeded,
        "swap_kb": peak_swap,
        "wall_s": round(time.monotonic() - started, 2),
        "exit": result,
        "cmd": command,
        "method": "procfs-rss-watchdog",
    }))
    return 0 if args.rss_limit_gib is None else (1 if exceeded else result)


if __name__ == "__main__":
    raise SystemExit(main())