#!/usr/bin/env python3
"""Run a command inside a cgroup-v2 MemoryMax scope and report peak memory."""

import argparse
import json
import os
import subprocess
import sys
import time


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit-gib", type=int, required=True)
    parser.add_argument("--timeout", type=float, default=7200)
    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")
    if args.limit_gib not in (2, 4, 8, 16, 32):
        parser.error("limit must be 2, 4, 8, 16, or 32 GiB")

    unit = f"flint-epq-bench-{os.getpid()}"
    runner = [
        "systemd-run", "--quiet", "--wait", "--collect", "--pipe",
        "--unit", unit, "--property", f"MemoryMax={args.limit_gib}G",
        "--property", "MemorySwapMax=0", "--", *command,
    ]
    started = time.monotonic()
    process = subprocess.Popen(runner, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    time.sleep(1)
    if process.poll() is not None:
        output, _ = process.communicate()
        if output:
            sys.stdout.write(output)
        print("UNTESTED: systemd/cgroup MemoryMax unavailable; RLIMIT_DATA is not a substitute", file=sys.stderr)
        return 77

    peak = swap_peak = 0
    while process.poll() is None:
        try:
            cgroup = subprocess.check_output(
                ["systemctl", "show", f"{unit}.scope", "-p", "ControlGroup", "--value"],
                text=True, stderr=subprocess.DEVNULL,
            ).strip()
            if cgroup:
                path = os.path.join("/sys/fs/cgroup", cgroup.lstrip("/"))
                for name, target in (("memory.peak", "peak"), ("memory.swap.peak", "swap")):
                    try:
                        with open(os.path.join(path, name), encoding="ascii") as stream:
                            value = int(stream.read().strip())
                        if target == "peak":
                            peak = max(peak, value)
                        else:
                            swap_peak = max(swap_peak, value)
                    except OSError:
                        pass
        except (OSError, subprocess.CalledProcessError):
            pass
        if time.monotonic() - started > args.timeout:
            process.kill()
            break
        time.sleep(0.25)

    output, _ = process.communicate()
    if output:
        sys.stdout.write(output)
    print("RUNMEM " + json.dumps({
        "limit_gib": args.limit_gib,
        "peak_bytes": peak,
        "peak_gib": round(peak / (1024 ** 3), 3),
        "swap_peak_bytes": swap_peak,
        "wall_s": round(time.monotonic() - started, 2),
        "exit": process.returncode,
        "cmd": command,
        "method": "cgroup-v2-MemoryMax",
    }))
    return process.returncode


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