"""Issue #30: a stall watcher that only notifies.

Runs next to the delegate wrapper's subscriber for a local herdr job and tells
people when the agent looks stuck. It never closes the job: its only output is
a `progress` event (through `_notify`, the single publisher) plus log lines,
and one short nudge prompt per job. Closing stays with the protocol events and
the two clocks of #18/#33.

  * blocked  -> `progress stalled:blocked <dialog line>` once per episode, no reply.
  * idle for MAM_STALL_IDLE_S (600) with no background children of the agent
    process and an empty input box -> `progress stalled:idle` and one nudge.
  * herdr failing (TRANSPORT_ERROR / AGENT_NOT_FOUND) is never read as "no
    event": warning + `progress stall_watch:herdr_unavailable:<code>`, backoff,
    and exit 2 after MAM_STALL_HERDR_FAIL_S (120) without recovery.

Signals are collected in one place (`HerdrSource.collect`) and judged by the
pure `judge`, so a future remote gateway can hand over the same `Signals`.

Exit codes: 0 the job reached a terminal status, 2 herdr stayed unavailable.

It lives at the top of lib_py, not in `delegation` or `loop`: it needs the
agent adapters, which neither package may import (their sibling-import rules),
and it runs as its own process, never inside a frozen loop.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import re
import signal
import subprocess
import sys
import time
from dataclasses import dataclass, field
from typing import List, Optional

from lib_py.agents.input_region import extract_input_region
from lib_py.agents.registry import agent_of_row, get_adapter
from lib_py.agents.sanitize import sanitize_herdr_agent_name
from lib_py.delegation.engine import TERMINAL_STATUSES
from lib_py.loop import watcher

log = logging.getLogger("stall_watch")

RC_TERMINAL = 0
RC_HERDR_UNAVAILABLE = 2

NOTIFY_BLOCKED = "notify_blocked"
NOTIFY_IDLE_AND_NUDGE = "notify_idle_and_nudge"
IDLE_STATES = ("idle", "done")
HEARTBEAT_MARK = "lib_py.delegation.heartbeat"
SKILLS_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
NUDGE = ("Job {job}: you look idle with no background work. If the job is done, publish "
         "completed or error with the command in your brief; otherwise continue where you left off.")


@dataclass
class Signals:
    agent_status: Optional[str] = None
    state_change_seq: Optional[int] = None
    pane_id: Optional[str] = None
    blocked_text: str = ""
    agent_pid: Optional[int] = None
    background_children: List[str] = field(default_factory=list)
    input_empty: Optional[bool] = None
    job_status: Optional[str] = None
    idle_for: float = 0.0


@dataclass
class WatchState:
    idle_s: float = 600.0
    reliable: bool = True
    nudged: bool = False
    blocked_notified: bool = False


def judge(sig: Signals, st: WatchState) -> Optional[str]:
    """Pure decision over one set of signals. Never an action that closes a job."""
    if sig.job_status in TERMINAL_STATUSES:
        return None
    if sig.agent_status == "blocked":
        return None if st.blocked_notified else NOTIFY_BLOCKED
    if (st.reliable and not st.nudged and sig.agent_status in IDLE_STATES
            and sig.idle_for >= st.idle_s and sig.agent_pid
            and not sig.background_children and sig.input_empty is True):
        return NOTIFY_IDLE_AND_NUDGE
    return None


def _proc_table():
    """pid -> (ppid, name, cmdline) from /proc (Linux); `ps` elsewhere."""
    table = {}
    if os.path.isdir("/proc/self"):
        for d in os.listdir("/proc"):
            if not d.isdigit():
                continue
            try:
                with open("/proc/%s/stat" % d) as f:
                    stat = f.read()
                with open("/proc/%s/cmdline" % d, "rb") as f:
                    argv = [a.decode("utf-8", "replace") for a in f.read().split(b"\0") if a]
            except OSError:
                continue
            name = stat[stat.index("(") + 1:stat.rindex(")")]
            ppid = int(stat[stat.rindex(")") + 2:].split()[1])
            table[int(d)] = (ppid, name, argv)
        return table
    out = subprocess.run(["ps", "-eo", "pid=,ppid=,comm=,args="], capture_output=True, text=True).stdout
    for line in out.splitlines():
        parts = line.split(None, 3)
        if len(parts) >= 3:
            table[int(parts[0])] = (int(parts[1]), os.path.basename(parts[2]), (parts[3] if len(parts) > 3 else parts[2]).split())
    return table


def background_children(agent_pid: int, resident=()) -> List[str]:
    """Descendants of the agent process, minus the adapter's resident children
    and the heartbeat timer (decision 9: it must never count as work)."""
    table = _proc_table()
    kids = {}
    for pid, (ppid, _, _) in table.items():
        kids.setdefault(ppid, []).append(pid)
    found, todo = [], list(kids.get(agent_pid, []))
    while todo:
        pid = todo.pop()
        todo.extend(kids.get(pid, []))
        _, name, argv = table[pid]
        base = os.path.basename(argv[0]) if argv else ""
        if name in resident or base in resident or name[:15] in [r[:15] for r in resident]:
            continue
        if HEARTBEAT_MARK in " ".join(argv):
            continue
        found.append("%d:%s" % (pid, name))
    return found


def resolve_agent(agent, target):
    """Adapter name for the wrapper's --agent (its default is `claude-code`),
    else from the session name; None when no adapter matches."""
    for name in (agent, (agent or "").split("-")[0], agent_of_row({}, session_name=target)):
        if name and get_adapter(name):
            return get_adapter(name).name
    return None


class HerdrError(Exception):
    def __init__(self, code):
        super().__init__(code)
        self.code = code


class HerdrSource(object):
    """Every herdr and OS read the judgement needs, in one place."""

    def __init__(self, target, agent, session=None, name=None):
        self.target, self.session = target, session or None
        # `name` is the name herdr knows the agent by (#66), worked out once
        # by StallWatch (#71); the kind is read off the original name, whose
        # suffix a long name's hash replaces.
        self.name = name or target
        self.agent = resolve_agent(agent, target)
        adapter = get_adapter(self.agent) if self.agent else None
        self.resident = tuple(adapter.resident_children) if adapter else ()

    def _herdr(self, *args):
        argv = ["herdr"] + (["--session", self.session] if self.session and self.session != "default" else []) + list(args)
        try:
            proc = subprocess.run(argv, capture_output=True, text=True, timeout=30,
                                  env=watcher.scoped_env(self.session, argv=argv))
        except Exception as exc:
            raise HerdrError("exec_failed: %s" % exc)
        if proc.returncode != 0:
            try:
                code = json.loads(proc.stderr)["error"]["code"]
            except Exception:
                code = "rc%d" % proc.returncode
            raise HerdrError(code)
        return proc.stdout

    def agent_get(self):
        payload = json.loads(self._herdr("agent", "get", self.name) or "{}")
        return (payload.get("result") or {}).get("agent") or {}

    def agent_pid(self, pane_id):
        info = json.loads(self._herdr("pane", "process-info", "--pane", pane_id) or "{}")
        procs = ((info.get("result") or {}).get("process_info") or {}).get("foreground_processes") or []
        # Not shell_pid: the shell is not the agent; without a pid we withhold.
        return procs[0].get("pid") if procs else None

    def blocked_text(self, pane_id):
        text = self._herdr("pane", "read", pane_id, "--source", "detection", "--lines", "8")
        for line in text.splitlines():
            line = re.sub(r"[^\x20-\x7e]+", " ", line).strip()
            line = re.sub(r"\s+", " ", line)
            if line:
                return line[:160]
        return ""

    def input_empty(self):
        """True/False from the adapter's input-box facts, None when it has none."""
        if not self.agent:
            return None
        try:
            return extract_input_region(self._herdr("agent", "read", self.name), self.agent) == ""
        except ValueError:
            return None

    def prompt(self, text):
        self._herdr("agent", "prompt", self.name, text)

    def collect(self, sig: Signals) -> Signals:
        """Fill the process and input signals of an idle agent."""
        agent = self.agent_get()
        sig.agent_status = agent.get("agent_status", sig.agent_status)
        sig.state_change_seq = agent.get("state_change_seq", sig.state_change_seq)
        sig.pane_id = agent.get("pane_id") or sig.pane_id
        sig.agent_pid = self.agent_pid(sig.pane_id) if sig.pane_id else None
        sig.background_children = background_children(sig.agent_pid, self.resident) if sig.agent_pid else []
        sig.input_empty = self.input_empty()
        return sig


class StallWatch(object):
    def __init__(self, args, source=None):
        self.args = args
        self.herdr_name = sanitize_herdr_agent_name(args.target)  # for herdr only; logs keep args.target
        self.src = source or HerdrSource(args.target, args.agent, args.herdr_session, name=self.herdr_name)
        name = resolve_agent(args.agent, args.target)
        adapter = get_adapter(name) if name else None
        reliable = bool(adapter.agent_status_reliable) if adapter else True
        self.st = WatchState(idle_s=float(os.environ.get("MAM_STALL_IDLE_S", "600")), reliable=reliable)
        self.fail_s = float(os.environ.get("MAM_STALL_HERDR_FAIL_S", "120"))
        self.once = set()

    def _notify(self, detail):
        """The only publisher: a `progress` event, nothing else. False when the
        registry or the wire refused it (a refused notice withholds the nudge)."""
        cmd = [sys.executable, "-m", "lib_py.delegation.publish_event", "--registry-dir", self.args.registry_dir,
               "--job", self.args.job, "--event", "progress", "--detail", detail, "--attempts", "1"]
        if self.args.expect_dispatch_token:
            cmd += ["--expect-dispatch-token", self.args.expect_dispatch_token]
        if self.args.attempt_id is not None:
            cmd += ["--attempt-id", str(self.args.attempt_id)]
        env = dict(os.environ)
        env["PYTHONPATH"] = SKILLS_DIR + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
        # rc 2 = broker unreachable after the local record and events.log were written.
        res = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=120)
        if res.returncode not in (0, 2):
            log.warning("stall_watch: publishing %r failed (rc %d): %s", detail, res.returncode, res.stderr.strip()[-200:])
            return False
        return True

    def _notify_once(self, key, detail, message):
        if key in self.once:
            return
        self.once.add(key)
        log.warning(message)
        self._notify(detail)

    def _job_status(self):
        try:
            with open(os.path.join(self.args.registry_dir, self.args.job + ".json")) as f:
                return json.load(f).get("status")
        except (OSError, ValueError):
            return None

    def _dispatch_open(self, strict=False):
        """Whether the job is still open *and* still the dispatch we were started
        for: a `completed` or a re-dispatch that landed during a blocking herdr
        wait must not be nudged (review R2). A missing record counts as gone;
        an unreadable one as a transient - unless `strict` (the check that
        permits a prompt), where only a readable, open record of this dispatch
        says yes (round-2 review D1)."""
        try:
            with open(os.path.join(self.args.registry_dir, self.args.job + ".json")) as f:
                rec = json.load(f)
        except FileNotFoundError:
            return False
        except (OSError, ValueError):
            return not strict
        if rec.get("status") in TERMINAL_STATUSES:
            return False
        tok = self.args.expect_dispatch_token
        if tok and rec.get("dispatch_token") and rec.get("dispatch_token") != tok:
            return False
        att = self.args.attempt_id
        if att is not None and rec.get("attempt_id") is not None and int(rec["attempt_id"]) != int(att):
            return False
        return True

    def _wrapper_gone(self):
        if self.args.watch_pid and not _alive(self.args.watch_pid):
            log.info("stall_watch: wrapper pid %d is gone; stall watch stops", self.args.watch_pid)
            return True
        return False

    def _may_prompt(self):
        """The last gate before keys reach the pane: the wrapper is alive and the
        record, read now, is this dispatch and open. Anything the notice's
        publish (up to 120 s) let happen - a completed, a re-dispatch, a dead
        wrapper, an unreadable record - withholds the nudge."""
        return not self._wrapper_gone() and self._dispatch_open(strict=True)

    def run(self) -> int:
        st, sig, agent = self.st, Signals(), self.args.agent
        if not st.reliable:
            log.warning("stall_watch: idle judgement off for %s (agent_status unreliable)", agent)
        backoff = watcher.ReconnectBackoff()
        fail_since, idle_since = None, None
        while True:
            sig.job_status = self._job_status()
            if sig.job_status in TERMINAL_STATUSES or not self._dispatch_open():
                return RC_TERMINAL
            if self._wrapper_gone():
                # The wrapper is gone (killed, crashed): nobody will stop us, so
                # stop ourselves - here, after every wait and before a prompt.
                return RC_TERMINAL
            if sig.agent_status is None:
                until, tmo = watcher.VALID_STATES, 1000  # learn the current state
            elif sig.agent_status in IDLE_STATES:
                left = st.idle_s - (time.monotonic() - idle_since)
                # Past S the judgement repeats every 30 s, not on every tick.
                until, tmo = ("working", "blocked"), int(1000 * (min(left, 30.0) if left > 0 else 30.0))
                tmo = max(tmo, 1000)
            elif sig.agent_status == "blocked":
                until, tmo = ("idle", "done", "working"), 60000
            else:
                until, tmo = ("idle", "done", "blocked"), 60000
            out = watcher.wait_once(self.herdr_name, until=until, timeout_ms=tmo, session=self.args.herdr_session)
            if self._wrapper_gone():  # round-2 review D2: before any notice
                return RC_TERMINAL
            if out.kind in (watcher.TRANSPORT_ERROR, watcher.AGENT_NOT_FOUND):
                code = "agent_not_found" if out.kind == watcher.AGENT_NOT_FOUND else (out.error_code or "transport_error")
                now = time.monotonic()
                if fail_since is None:
                    fail_since = now
                    log.warning("stall_watch: herdr unavailable: %s %s", code, out.detail)
                    self._notify("stall_watch:agent_not_found" if code == "agent_not_found"
                                 else "stall_watch:herdr_unavailable:%s" % code)
                if now - fail_since >= self.fail_s:
                    log.warning("stall_watch: herdr unavailable for %ds; stall watch stops (%s)", int(now - fail_since), code)
                    return RC_HERDR_UNAVAILABLE
                time.sleep(min(backoff.record_failure(), max(0.1, self.fail_s - (now - fail_since))))
                sig.agent_status = None  # re-observe after a gap
                continue
            backoff.record_success()
            fail_since = None
            if not self._dispatch_open():
                log.info("stall_watch: job %s closed or re-dispatched during the wait; stall watch stops", self.args.job)
                return RC_TERMINAL
            if out.kind == watcher.MATCHED:
                new = out.agent.get("agent_status")
                sig.pane_id = out.agent.get("pane_id") or sig.pane_id
                if new in IDLE_STATES and sig.agent_status not in IDLE_STATES:
                    idle_since, sig.state_change_seq = time.monotonic(), out.agent.get("state_change_seq")
                if new != "blocked":
                    st.blocked_notified = False
                sig.agent_status = new
            try:
                idle_since = self._act(sig, idle_since)
            except HerdrError as exc:
                log.warning("stall_watch: herdr call failed while judging: %s", exc.code)

    def _act(self, sig, idle_since):
        """Notify/nudge per `judge`; returns the idle start to keep (restarted
        when a transition was missed between the wait and the judgement)."""
        st, agent = self.st, self.args.agent
        if sig.agent_status == "blocked" and not st.blocked_notified:
            sig.blocked_text = self.src.blocked_text(sig.pane_id or self.src.agent_get().get("pane_id", ""))
            if judge(sig, st) == NOTIFY_BLOCKED:
                log.warning("stall_watch: agent %s is blocked: %s", self.args.target, sig.blocked_text)
                self._notify(("stalled:blocked " + sig.blocked_text).strip())
                st.blocked_notified = True
            return idle_since
        if sig.agent_status not in IDLE_STATES or not st.reliable or st.nudged:
            return idle_since
        sig.idle_for = time.monotonic() - idle_since
        if sig.idle_for < st.idle_s:
            return idle_since
        seq = sig.state_change_seq
        self.src.collect(sig)
        if seq is not None and sig.state_change_seq is not None and sig.state_change_seq != seq:
            # Review R4: a transition we did not see. Decision 9 wants S seconds
            # of *continuous* idle, so the window restarts from this observation
            # (collect() already took the new seq as the baseline).
            log.info("stall_watch: state_change_seq %s -> %s since the idle start; idle window restarts",
                     seq, sig.state_change_seq)
            return time.monotonic()
        if sig.background_children:
            log.debug("stall_watch: idle but background children %s", sig.background_children)
        elif sig.input_empty is None:
            self._notify_once("input_unavailable", "stall_watch:input_check_unavailable:%s" % agent,
                              "stall_watch: input check unavailable for %s; idle nudge withheld" % agent)
        elif sig.input_empty is False:
            log.info("stall_watch: input box of %s is not empty; nudge withheld", self.args.target)
            if sig.idle_for >= 3 * st.idle_s:
                self._notify_once("never_empty", "stall_watch:input_never_empty:%s" % agent,
                                  "stall_watch: input box of %s never looked empty in %ds" % (agent, int(sig.idle_for)))
        if judge(sig, st) == NOTIFY_IDLE_AND_NUDGE:
            if not self._dispatch_open():  # collect() took time; the job may have closed meanwhile
                log.info("stall_watch: job %s closed or re-dispatched before the nudge; nothing sent", self.args.job)
                return idle_since
            log.warning("stall_watch: agent %s idle for %ds with no background work; nudging once",
                        self.args.target, int(sig.idle_for))
            st.nudged = True
            if not self._notify("stalled:idle"):
                log.warning("stall_watch: the stalled:idle notice was refused; nudge withheld")
                return idle_since
            if not self._may_prompt():  # the publish can take long; re-read before the keys
                log.info("stall_watch: job %s closed, re-dispatched or unreadable, or the wrapper left, "
                         "during the notice; nudge withheld", self.args.job)
                return idle_since
            self.src.prompt(NUDGE.format(job=self.args.job))
        return idle_since


def _alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    return True


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description="notify-only stall watcher for one local herdr job (#30)")
    p.add_argument("--job", required=True)
    p.add_argument("--registry-dir", required=True)
    p.add_argument("--target", required=True, help="session name of the job (herdr is called with its sanitized form)")
    p.add_argument("--agent", required=True, help="agent CLI (adapter name)")
    p.add_argument("--herdr-session", default="")
    p.add_argument("--expect-dispatch-token", default=None)
    p.add_argument("--attempt-id", type=int, default=None)
    p.add_argument("--watch-pid", type=int, default=None, help="exit once this process (the wrapper) is gone")
    args = p.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s stall_watch %(levelname)s %(message)s")
    # The wrapper stops us with SIGKILL once its subscriber ends; a TERM from anyone else: leave quietly.
    signal.signal(signal.SIGTERM, lambda *_: os._exit(RC_TERMINAL))
    return StallWatch(args).run()


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