#!/usr/bin/env python3
"""Turn silent multi-hour stalls into immediate loud failures.

Liveness is proven by PROGRESS, never by process existence: a crash-looping or
epoll-blocked child holds a live PID for its whole timeout while doing nothing.
Four orthogonal signals are sampled across the whole process tree. A run is
stuck when EVERY one of them is flat for the full idle window, or when output
alone is flat for the (much longer) output ceiling — the ceiling exists because
ANDing the signals lets one spinning signal veto detection forever.

  output   bytes written to the captured log
  cpu      utime+stime jiffies, recursive
  io       rchar+wchar, recursive (counts socket traffic; read_bytes does not)
  beat     mtime of a heartbeat file a legitimate waiter touches

Subcommands:
  run     supervise a command
  replay  re-run the detector over recorded traces (false-positive verification)
"""

import argparse
import json
import os
import signal
import subprocess
import sys
import threading
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "lib"))
from stall_detect import accumulate

STALL_EXIT = 91
CLOCK_TICKS = os.sysconf("SC_CLK_TCK")


def read_proc_tree(root_pid):
    """Every live descendant of root_pid, inclusive."""
    children = {}
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        try:
            with open(f"/proc/{entry}/stat", "rb") as handle:
                fields = handle.read().rsplit(b")", 1)[1].split()
            children.setdefault(int(fields[1]), []).append(int(entry))
        except (OSError, IndexError, ValueError):
            continue
    seen, stack = [], [root_pid]
    while stack:
        pid = stack.pop()
        seen.append(pid)
        stack.extend(children.get(pid, []))
    return seen


def sample_tree(root_pid):
    """Summed cpu jiffies and syscall bytes over the tree."""
    cpu = io_bytes = 0
    for pid in read_proc_tree(root_pid):
        try:
            with open(f"/proc/{pid}/stat", "rb") as handle:
                fields = handle.read().rsplit(b")", 1)[1].split()
            cpu += int(fields[11]) + int(fields[12])
        except (OSError, IndexError, ValueError):
            pass
        try:
            with open(f"/proc/{pid}/io", "rb") as handle:
                for line in handle:
                    name, _, value = line.partition(b":")
                    if name in (b"rchar", b"wchar"):
                        io_bytes += int(value)
        except (OSError, ValueError):
            pass
    return cpu, io_bytes


def size_of(path):
    try:
        return os.path.getsize(path)
    except OSError:
        return 0


def mtime_of(path):
    try:
        return os.path.getmtime(path)
    except OSError:
        return 0.0


def classify(samples, idle_window_s, output_ceiling_s=0.0):
    """Pure detector. Returns a stall reason, or None.

    Two independent rules:
      1. every tracked signal flat across the trailing idle_window
      2. `out` alone flat across output_ceiling_s, whatever cpu/io do

    Rule 2 exists because rule 1 ANDs the signals, so a spin loop, retry loop
    or keepalive that moves ONE signal forever would veto detection for good.
    An unreadable counter is not treated as flat-and-therefore-stalled: `cpu`
    or `io` all-zero for the whole run disables that signal rather than firing.
    """
    if len(samples) < 2:
        return None
    last = samples[-1]
    cumulative = accumulate(samples, ("out", "cpu", "io", "beat"))

    def anchor_index(window_s):
        """Newest sample at or before the window start, or None if too young."""
        window_start = last["t"] - window_s
        found = None
        for i, sample in enumerate(samples):
            if sample["t"] > window_start:
                break
            found = i
        return found

    if output_ceiling_s:
        index = anchor_index(output_ceiling_s)
        if index is not None and cumulative["out"][index] == cumulative["out"][-1]:
            return f"no output for {last['t'] - samples[index]['t']:.0f}s (ceiling)"

    index = anchor_index(idle_window_s)
    if index is None:
        return None

    cpu_ever = any(s["cpu"] for s in samples)
    io_ever = any(s["io"] for s in samples)

    tracked, flat = [], []
    for key, enabled in (("out", True), ("cpu", cpu_ever), ("io", io_ever), ("beat", True)):
        if not enabled:
            continue
        tracked.append(key)
        if cumulative[key][index] == cumulative[key][-1]:
            flat.append(key)
    if not tracked or len(flat) != len(tracked):
        return None
    idle_for = last["t"] - samples[index]["t"]
    return f"no progress on {'+'.join(tracked)} for {idle_for:.0f}s"


def kill_tree(root_pid, grace_s):
    """SIGTERM the whole tree, wait for it to actually die, then SIGKILL.

    killpg alone misses grandchildren that called setsid, and an unconditional
    sleep before SIGKILL lets a recycled PGID absorb the second signal.
    """
    def signal_all(sig):
        pids = read_proc_tree(root_pid)
        try:
            os.killpg(os.getpgid(root_pid), sig)
        except OSError:
            pass
        for pid in pids:
            try:
                os.kill(pid, sig)
            except OSError:
                pass

    signal_all(signal.SIGTERM)
    deadline = time.monotonic() + grace_s
    while time.monotonic() < deadline:
        if not [p for p in read_proc_tree(root_pid) if os.path.exists(f"/proc/{p}")]:
            return
        time.sleep(0.5)
    signal_all(signal.SIGKILL)


def dump_diagnostics(stream, samples, log_path, reason, root_pid):
    stream.write(f"\n[stall-guard] STALL DETECTED: {reason}\n")
    stream.write(f"[stall-guard] log: {log_path}\n")
    stream.write("[stall-guard] trailing samples (t, out, cpu, io, beat):\n")
    for s in samples[-8:]:
        stream.write(f"  {s['t']:8.0f} {s['out']:12d} {s['cpu']:10d} {s['io']:14d} {s['beat']:.0f}\n")
    try:
        tree = read_proc_tree(root_pid)
        stream.write(f"[stall-guard] live pids: {tree}\n")
        for pid in tree[:12]:
            try:
                cmd = open(f"/proc/{pid}/cmdline", "rb").read().replace(b"\0", b" ").decode(errors="replace")
                state = open(f"/proc/{pid}/stat", "rb").read().rsplit(b")", 1)[1].split()[0].decode()
                wchan = open(f"/proc/{pid}/wchan").read().strip() or "-"
                stream.write(f"  pid {pid} state={state} wchan={wchan} :: {cmd[:120]}\n")
            except OSError:
                pass
    except OSError:
        pass
    try:
        with open(log_path, "rb") as handle:
            handle.seek(max(0, size_of(log_path) - 4000))
            tail = handle.read().decode(errors="replace")
        stream.write("[stall-guard] last output:\n" + tail + "\n")
    except OSError:
        pass
    stream.flush()


def notify(title, body):
    if os.environ.get("STALL_GUARD_NOTIFY") != "1":
        sys.stderr.write(f"[stall-guard] desktop notification suppressed: {title} — {body}\n")
        return
    try:
        subprocess.run(["notify-send", "-u", "critical", title, body], timeout=10)
    except (OSError, subprocess.SubprocessError):
        pass


def cmd_run(args):
    log_path = args.log or os.path.join(
        os.environ.get("XDG_RUNTIME_DIR") or "/tmp", f"stall-guard-{args.key}-{os.getpid()}.log"
    )
    os.makedirs(os.path.dirname(log_path), exist_ok=True)
    trace_path = args.trace or os.path.join(
        os.path.expanduser("~/.cache/stall-guard"), f"{args.key}-{int(time.time())}.jsonl"
    )
    os.makedirs(os.path.dirname(trace_path), exist_ok=True)

    log = open(log_path, "wb")
    tee = subprocess.Popen(["tee", "-a", log_path], stdin=subprocess.PIPE)
    child = subprocess.Popen(args.command, stdout=tee.stdin, stderr=subprocess.STDOUT, start_new_session=True)
    log.close()

    samples = []
    verdict = {"stalled": False, "reason": None}
    noticed = []
    started = time.monotonic()
    stop = threading.Event()

    def watch():
        while not stop.wait(args.sample_interval):
            now = time.monotonic() - started
            cpu, io_bytes = sample_tree(child.pid)
            sample = {
                "t": now,
                "out": size_of(log_path),
                "cpu": cpu,
                "io": io_bytes,
                "beat": mtime_of(args.heartbeat) if args.heartbeat else 0.0,
            }
            samples.append(sample)
            with open(trace_path, "a") as handle:
                handle.write(json.dumps(sample) + "\n")

            if args.wall_notice and now > args.wall_notice and not noticed:
                noticed.append(now)
                notice = (
                    f"still running after {now / 60:.0f}m — slower than 99.9% of past runs. "
                    f"Progress signals are moving, so this is NOT a stall; check it yourself."
                )
                sys.stderr.write(f"[stall-guard] SLOW: {notice}\n")
                sys.stderr.flush()
                notify(f"stall-guard: {args.key} unusually slow", notice)

            reason = classify(samples, args.idle_window, args.output_ceiling)
            if reason is None and args.wall_deadline and now > args.wall_deadline:
                reason = f"wall-clock deadline exceeded ({now:.0f}s > {args.wall_deadline}s)"
            if reason and not verdict["stalled"]:
                verdict["stalled"] = True
                verdict["reason"] = reason
                dump_diagnostics(sys.stderr, samples, log_path, reason, child.pid)
                notify(f"stall-guard: {args.key} stuck", reason)
                if args.enforce:
                    kill_tree(child.pid, args.kill_grace)
                    return
                sys.stderr.write("[stall-guard] observe mode: NOT killing; recording only\n")

    watcher = threading.Thread(target=watch, daemon=True)
    watcher.start()
    status = child.wait()
    stop.set()
    watcher.join(timeout=args.sample_interval + 5)
    try:
        tee.stdin.close()
        tee.wait(timeout=10)
    except (OSError, subprocess.SubprocessError):
        pass

    sys.stderr.write(f"[stall-guard] trace {trace_path} ({len(samples)} samples)\n")
    if verdict["stalled"] and args.enforce:
        sys.stderr.write(f"[stall-guard] FAILED LOUDLY: {verdict['reason']}\n")
        return STALL_EXIT
    if verdict["stalled"]:
        sys.stderr.write(f"[stall-guard] would have failed: {verdict['reason']} (observe mode)\n")
    return status


def cmd_replay(args):
    """Offline detector verification over recorded traces.

    Both arms are required before arming enforcement: zero firings across traces
    of runs that completed successfully, and confirmed firing on known-stuck ones.
    """
    fired = []
    for path in args.traces:
        samples = []
        with open(path) as handle:
            for line in handle:
                line = line.strip()
                if line:
                    samples.append(json.loads(line))
        reason = None
        for i in range(2, len(samples) + 1):
            reason = classify(samples[:i], args.idle_window, args.output_ceiling)
            if reason:
                break
        status = "FIRED" if reason else "quiet"
        if reason:
            fired.append(path)
        print(f"{status:6} {os.path.basename(path)} ({len(samples)} samples) {reason or ''}")
    print(f"\n{len(fired)} of {len(args.traces)} traces fired at idle_window={args.idle_window}s")
    if args.expect == "quiet" and fired:
        print("FALSE POSITIVES — do not arm enforcement at this window")
        return 1
    if args.expect == "fired" and len(fired) != len(args.traces):
        print("MISSED DETECTIONS — detector does not catch known stalls")
        return 1
    return 0


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="subcommand", required=True)

    run = sub.add_parser("run", help="supervise a command")
    run.add_argument("--key", required=True, help="short label for traces and notifications")
    run.add_argument("--idle-window", type=float, default=900.0, help="seconds of total flatness before stall (default 900)")
    run.add_argument("--output-ceiling", type=float, default=0.0, help="seconds of zero output before stall, whatever cpu/io do (0 = none)")
    run.add_argument("--wall-notice", type=float, default=0.0, help="seconds after which to warn LOUDLY without killing (0 = none)")
    run.add_argument("--wall-deadline", type=float, default=0.0, help="hard wall-clock cap in seconds (0 = none)")
    run.add_argument("--sample-interval", type=float, default=15.0)
    run.add_argument("--heartbeat", help="file a legitimate long wait touches to declare itself alive")
    run.add_argument("--log", help="captured output path")
    run.add_argument("--trace", help="sample trace path (JSONL)")
    run.add_argument("--kill-grace", type=float, default=10.0)
    run.add_argument("--enforce", action="store_true", help="kill on stall; default is observe-only")
    run.add_argument("command", nargs=argparse.REMAINDER)
    run.set_defaults(func=cmd_run)

    replay = sub.add_parser("replay", help="re-run the detector over recorded traces")
    replay.add_argument("traces", nargs="+")
    replay.add_argument("--idle-window", type=float, default=900.0)
    replay.add_argument("--output-ceiling", type=float, default=0.0)
    replay.add_argument("--expect", choices=["quiet", "fired", "any"], default="any")
    replay.set_defaults(func=cmd_replay)

    args = parser.parse_args()
    if args.subcommand == "run":
        if args.command and args.command[0] == "--":
            args.command = args.command[1:]
        if not args.command:
            parser.error("run requires a command after --")
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())