#!/usr/bin/env python3
import os
import subprocess
import sys
import time
from datetime import datetime


def now_iso():
    return datetime.now().astimezone().isoformat(timespec="seconds")


def mem_available_mb():
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) // 1024
    return None


def proc_alive(name):
    return subprocess.run(["pgrep", "-x", name],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0


def main():
    threshold = int(sys.argv[1]) if len(sys.argv) > 1 else 1500
    proc = sys.argv[2] if len(sys.argv) > 2 else "deqp-vk"
    poll = int(sys.argv[3]) if len(sys.argv) > 3 else 3

    print(f"guarding: kill {proc} when MemAvailable < {threshold}M (poll {poll}s)")

    while True:
        avail = mem_available_mb()
        if avail is not None and avail < threshold:
            print(f"{now_iso()} avail={avail}M below {threshold}M — terminating {proc}")
            subprocess.run(["pkill", "-TERM", "-x", proc])
            time.sleep(10)
            if proc_alive(proc):
                print("still alive, SIGKILL")
                subprocess.run(["pkill", "-KILL", "-x", proc])

            outdir = os.environ.get("OUTDIR", ".")
            with open(os.path.join(outdir, "mem-guard.log"), "a") as f:
                f.write(f"{now_iso()} killed {proc} at avail={avail}M\n")
            subprocess.run(["sync"])
            sys.exit(1)

        time.sleep(poll)


if __name__ == "__main__":
    main()