"""Watch forked ZeroGPU workers from the parent process (Linux /proc only).

A worker that stalls before its first output leaves nothing in the logs, and a
frozen process cannot report on itself. While a GPU call is pending, the parent
logs each python child's state, per-thread kernel wait channels, CPU time and
I/O. When a worker makes no progress between two looks, the parent sends it
SIGUSR1 once; faulthandler, registered before the fork, then dumps every
Python thread's stack from inside the worker.
"""
from __future__ import annotations

from collections import Counter
import os
from pathlib import Path
import signal

PROC = Path("/proc")
# No progress: under 5% of one CPU and under 1 MiB of reads/writes per interval.
STALL_CPU_FRACTION = 0.05
STALL_IO_BYTES = 1 << 20


def _stat(text):
    """(comm, fields after comm); comm may itself contain spaces or parentheses."""
    end = text.rindex(")")
    return text[text.index("(") + 1:end], text[end + 2:].split()


def child_pids(parent=None, proc=PROC):
    parent = os.getpid() if parent is None else parent
    pids = []
    for entry in proc.iterdir():
        if not entry.name.isdigit():
            continue
        try:
            comm, fields = _stat((entry / "stat").read_text())
        except (OSError, ValueError):
            continue
        if int(fields[1]) == parent and "python" in comm:
            pids.append(int(entry.name))
    return sorted(pids)


def snapshot(pid, proc=PROC):
    base = proc / str(pid)
    _, fields = _stat((base / "stat").read_text())
    try:
        io = dict(line.split(": ", 1) for line in (base / "io").read_text().splitlines())
        io_bytes = int(io["rchar"]) + int(io["wchar"])
    except (OSError, KeyError, ValueError):
        io_bytes = None
    threads = Counter()
    for task in (base / "task").iterdir():
        try:
            state = _stat((task / "stat").read_text())[1][0]
            wchan = (task / "wchan").read_text().strip() or "-"
        except (OSError, ValueError, IndexError):
            continue
        threads[f"{state}:{wchan}"] += 1
    return {"state": fields[0], "cpu_ticks": int(fields[11]) + int(fields[12]),
            "io_bytes": io_bytes, "threads": dict(threads)}


def stalled(before, now, interval, ticks_per_s):
    if before is None:
        return False
    cpu_s = (now["cpu_ticks"] - before["cpu_ticks"]) / ticks_per_s
    io = None if None in (now["io_bytes"], before["io_bytes"]) else now["io_bytes"] - before["io_bytes"]
    return cpu_s < STALL_CPU_FRACTION * interval and (io is None or io < STALL_IO_BYTES)


def watch(stop, interval, proc=PROC, kill=os.kill):
    """Until stop is set, log every python child each interval; dump a stalled one once."""
    ticks_per_s = os.sysconf("SC_CLK_TCK")
    last, dumped = {}, set()
    while not stop.wait(interval):
        for pid in child_pids(proc=proc):
            try:
                now = snapshot(pid, proc)
            except (OSError, ValueError, IndexError):
                continue
            stuck = stalled(last.get(pid), now, interval, ticks_per_s)
            last[pid] = now
            threads = ",".join(f"{key}x{count}" for key, count in sorted(now["threads"].items()))
            io_mib = "?" if now["io_bytes"] is None else now["io_bytes"] >> 20
            print(f"worker_probe pid={pid} state={now['state']} cpu_s={now['cpu_ticks'] / ticks_per_s:.0f} "
                  f"io_mib={io_mib} stalled={stuck} threads={threads}", flush=True)
            if stuck and pid not in dumped:
                dumped.add(pid)
                try:
                    kill(pid, signal.SIGUSR1)
                except OSError:
                    pass