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


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


def meminfo():
    d = {}
    with open("/proc/meminfo") as f:
        for line in f:
            k, v = line.split(":", 1)
            d[k.strip()] = int(v.split()[0])
    return d


def main():
    outdir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.cwd() / "hangs"
    poll = int(sys.argv[2]) if len(sys.argv) > 2 else 10

    err = ""
    for c in sorted(glob.glob("/sys/class/drm/card*/error")):
        err = c
        break

    if not err:
        print("no /sys/class/drm/card*/error — is i915 loaded?", file=sys.stderr)
    elif not os.access(err, os.R_OK):
        print(f"cannot read {err} — rerun with sudo for GPU hang capture", file=sys.stderr)
        err = ""

    outdir.mkdir(parents=True, exist_ok=True)
    hangs_txt = outdir / "hangs.txt"
    mem_log = outdir / "mem.log"

    print(f"watching {err or '(memory only)'} every {poll}s -> {outdir}")
    with open(hangs_txt, "a") as f:
        f.write(f"{now_iso()} watch started\n")

    n = 0
    while True:
        if err and os.path.getsize(err) > 0:
            n += 1
            stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            hang_f = outdir / f"hang-{stamp}.txt"
            data = open(err, "rb").read()
            hang_f.write_bytes(data)

            with open(hangs_txt, "a") as hf:
                for line in data.decode(errors="replace").splitlines():
                    if line.startswith(("GPU HANG", "Kernel", "Time")):
                        hf.write(line + "\n")
                        break
                msg = f"{now_iso()} hang #{n} -> {hang_f} ({hang_f.stat().st_size} B)"
                hf.write(msg + "\n")
            print(msg)

            dmesg = subprocess.run(["dmesg"], capture_output=True, text=True).stdout
            matched = [l for l in dmesg.splitlines()
                       if any(k.lower() in l.lower() for k in ("GPU HANG", "reset", "hangcheck"))]
            (outdir / f"dmesg-{stamp}.txt").write_text("\n".join(matched[-20:]) + ("\n" if matched else ""))

            with open(err, "w") as ef:
                ef.write("1")

        mi = meminfo()
        avail = mi.get("MemAvailable", 0)
        swtot = mi.get("SwapTotal", 0)
        swfree = mi.get("SwapFree", 0)
        sunrecl = mi.get("SUnreclaim", 0)
        ptbl = mi.get("PageTables", 0)

        ps = subprocess.run(["ps", "-eo", "rss,comm", "--sort=-rss", "--no-headers"],
                             capture_output=True, text=True).stdout
        top3 = []
        for line in ps.splitlines()[:3]:
            parts = line.split(None, 1)
            if len(parts) == 2:
                rss, comm = parts
                top3.append(f"{comm}:{int(rss) // 1024}M")
        top3_str = " ".join(top3) + (" " if top3 else "")

        line = (f"{now_iso()} avail={avail // 1024}M swapused={(swtot - swfree) // 1024}M "
                f"sunreclaim={sunrecl // 1024}M pagetables={ptbl // 1024}M {top3_str}\n")
        with open(mem_log, "a") as mf:
            mf.write(line)
        try:
            subprocess.run(["sync", "-d", str(mem_log)], check=True,
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except subprocess.CalledProcessError:
            subprocess.run(["sync"])

        time.sleep(poll)


if __name__ == "__main__":
    main()