"""heartbeat.py — the agent-side automatic heartbeat timer (Issue #18).

The liveness contract obliges an agent to publish `progress` at least every N
seconds while it works. Relying on the brief to say so meant an agent that
forgot, or one whose turn was one long tool call, was closed as unresponsive
while it was fine. This timer publishes the heartbeat for it.

`publish_event.py --event started --heartbeat N` starts one detached timer per
job; it runs in its own session so the agent's shell may return at once. The
timer stops on its own when

  * the job record reaches a terminal status (completed / error / cancelled),
  * the record's dispatch_token or attempt_id changes (the job was
    re-dispatched: this attempt's heartbeats would be refused anyway),
  * the job record disappears, or
  * its lifetime exceeds the job's wall-clock timeout plus one period.

A terminal publish also stops it directly, so no heartbeat can trail a
`completed` (review round 2, codex R3): it raises an attempt-scoped barrier
under the registry lock that allocates seqs (a beat either took its seq before
it or is refused when it asks for one), stops the timer through the pidfile
(`<registry>/<job_id>.heartbeat.pid`), and waits for the timer to drain its
in-flight beats - killing it after `MAM_HEARTBEAT_DRAIN_S` - before it takes
its own seq.

The timer is bound to the agent's life (review B2): it watches one process -
`--watch-pid`, or by default the nearest ancestor of the publishing process
that is an agent CLI (claude, codex, agy, ...; the agent runs the start
command from its shell tool, so that ancestor is the agent itself) - and stops
within one period once that process is gone. A crashed or killed agent
therefore stops heartbeating and is closed after K x N like any silent one;
an agent that is alive but stuck keeps heartbeating, which is the contract's
deliberate trade-off (liveness of the process, not of the work).

Each beat is an ordinary `publish_event` call in this process: it takes the
next seq, is signed with the job's auth_token, carries the dispatch token and
attempt id the timer was started with, and is recorded in events.log before
the wire publish - exactly what the orchestrator's disk and wire liveness
checks read.
"""
from __future__ import annotations

import argparse
import contextlib
import fcntl
import json
import logging
import os
import signal
import subprocess
import sys
import threading
import time
from typing import List, Optional

from lib_py.delegation import mqtt_common
from lib_py.delegation.mqtt_common import default_registry_dir, load_job

logger = logging.getLogger("delegate_job.heartbeat")

TERMINAL_STATUSES = ("completed", "error", "cancelled")
HEARTBEAT_DETAIL = "heartbeat"
MAX_IN_FLIGHT = 8


# The agent CLIs, as lib_py.agents.base names them. Repeated here rather than
# imported: lib_py.delegation may not import a sibling lib_py package (a
# frozen loop binds delegation to a different tree, see __init__).
AGENT_CLI_NAMES = ("claude", "agy", "hermes", "grok", "opencode", "codex")


def _is_agent_cli(command: str) -> bool:
    """Some token before the first option has an agent basename
    (`claude ...`, `node .../claude ...`, `/bin/sh .../agy ...`)."""
    for tok in command.split():
        if tok.startswith("-"):
            break
        if os.path.basename(tok) in AGENT_CLI_NAMES:
            return True
    return False


def _process_table():
    """{pid: (ppid, command)} from ps; empty when ps is unavailable."""
    try:
        out = subprocess.run(["ps", "-axo", "pid=,ppid=,command="], capture_output=True, text=True).stdout
    except OSError:
        return {}
    table = {}
    for line in out.splitlines():
        parts = line.split(None, 2)
        if len(parts) < 2:
            continue
        try:
            pid, ppid = int(parts[0]), int(parts[1])
        except ValueError:
            continue
        table[pid] = (ppid, parts[2] if len(parts) > 2 else "")
    return table


def nearest_agent_cli_ancestor(start_pid: Optional[int] = None) -> Optional[int]:
    """PID of the closest ancestor (start_pid included) that is an agent CLI.

    The agent executes the injected start command through its shell tool, so
    the chain is agent CLI -> shell -> python publish_event; the agent is the
    first agent-CLI process met walking up. None when there is no such
    ancestor (a test, a plain shell, a remote runner)."""
    table = _process_table()
    pid = start_pid if start_pid is not None else os.getpid()
    seen = set()
    while pid and pid > 1 and pid not in seen:
        seen.add(pid)
        entry = table.get(pid)
        if entry is None:
            return None
        ppid, command = entry
        if _is_agent_cli(command):
            return pid
        pid = ppid
    return None


def pidfile_path(job_id: str, registry_dir: str) -> str:
    return os.path.join(registry_dir, "{}.heartbeat.pid".format(job_id))


def _pid_alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    try:
        state = subprocess.run(["ps", "-o", "stat=", "-p", str(pid)],
                               capture_output=True, text=True).stdout.strip()
    except OSError:
        return True
    return bool(state) and not state.startswith("Z")


def running_timer(job_id: str, registry_dir: str) -> Optional[int]:
    """PID of a live timer for this job, or None."""
    try:
        with open(pidfile_path(job_id, registry_dir)) as fh:
            pid = int(fh.read().strip() or 0)
    except (OSError, ValueError):
        return None
    if pid > 0 and _pid_alive(pid):
        return pid
    return None


def _is_timer_process(pid: int, job_id: str) -> bool:
    """The pid really is this job's timer (a stale pidfile after a crash may
    name a recycled pid that now belongs to something else)."""
    entry = _process_table().get(pid)
    if entry is None:
        return False
    command = entry[1]
    return "lib_py.delegation.heartbeat" in command and job_id in command


DRAIN_ENV = "MAM_HEARTBEAT_DRAIN_S"
# Short on purpose: the barrier already keeps any beat from taking a seq after
# the terminal, so the drain only lets an idle timer exit cleanly. A beat still
# on the wire is cancelled with the timer (its seq is older than the terminal's)
# rather than holding back the terminal event - which may be the orchestrator
# closing the job - for seconds of broker round trips.
DEFAULT_DRAIN_S = 1.0


def drain_seconds() -> float:
    """How long a terminal publish waits for the timer to drain (env override)."""
    try:
        return max(0.0, float(os.environ.get(DRAIN_ENV, DEFAULT_DRAIN_S)))
    except ValueError:
        return DEFAULT_DRAIN_S


def lifecycle_lock_path(job_id: str, registry_dir: str) -> str:
    return os.path.join(registry_dir, "{}.heartbeat.lock".format(job_id))


@contextlib.contextmanager
def lifecycle_lock(job_id: str, registry_dir: str, blocking: bool = True):
    """The per-job timer ownership lock (review round 4, codex R3).

    Observing/draining the previous owner, spawning the replacement and
    publishing its pid through the pidfile is one ownership operation, and the
    terminal stop takes part in it, so two concurrent `started --heartbeat`
    calls cannot both arm a timer with only one of them tracked, and a start
    that loses the race against a terminal publish finds the barrier and does
    not arm at all. This is a separate flock from the registry/seq lock on
    purpose: an in-flight beat needs that lock to take its seq, and the drain
    waits for exactly those beats, so holding the seq lock across the drain
    would deadlock. With ``blocking=False`` it yields False instead of waiting
    when the lock is busy - an exiting timer uses that (review round 5, codex F1)."""
    os.makedirs(registry_dir, exist_ok=True)
    with open(lifecycle_lock_path(job_id, registry_dir), "a+") as fh:
        try:
            fcntl.flock(fh.fileno(), fcntl.LOCK_EX if blocking else fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            yield False  # busy and non-blocking: the caller must not touch owned state
            return
        try:
            yield True
        finally:
            fcntl.flock(fh.fileno(), fcntl.LOCK_UN)


def stop_timer(job_id: str, registry_dir: str, drain: Optional[float] = None) -> bool:
    """Terminate the job's timer if one runs. Returns True when one was signalled.

    ``drain`` (seconds): also wait for the timer to exit - on TERM it stops
    scheduling and joins its in-flight beats - and SIGKILL it if it has not
    gone by then, so none of its beats can still publish afterwards. A terminal
    publish uses this (review round 2, codex R3): "TERM sent" is not "timer
    stopped". Serialised with `start_timer` through the lifecycle lock."""
    with lifecycle_lock(job_id, registry_dir):
        return _stop_timer_locked(job_id, registry_dir, drain)


def _stop_timer_locked(job_id: str, registry_dir: str, drain: Optional[float] = None) -> bool:
    """`stop_timer` body; the caller holds the lifecycle lock."""
    pid = running_timer(job_id, registry_dir)
    path = pidfile_path(job_id, registry_dir)
    if pid is None or not _is_timer_process(pid, job_id):
        try:
            os.unlink(path)
        except OSError:
            pass
        return False
    try:
        os.kill(pid, signal.SIGTERM)
    except OSError:
        pass
    if drain is not None:
        deadline = time.monotonic() + drain
        while time.monotonic() < deadline and _pid_alive(pid):
            time.sleep(0.1)
        if _pid_alive(pid):
            logger.warning("heartbeat timer %d for %s did not drain within %.0fs; killing it", pid, job_id, drain)
            try:
                os.kill(pid, signal.SIGKILL)
            except OSError:
                pass
            deadline = time.monotonic() + 2.0
            while time.monotonic() < deadline and _pid_alive(pid):
                time.sleep(0.05)
    try:
        os.unlink(path)
    except OSError:
        pass
    return True


def barrier_path(job_id: str, registry_dir: str) -> str:
    return os.path.join(registry_dir, "{}.heartbeat.barrier".format(job_id))


def raise_barrier(job_id: str, registry_dir: str, dispatch_token=None, attempt_id=None) -> None:
    """Mark this attempt as ending (review round 2, codex R3).

    Written under the registry lock - the lock `next_seq` takes - so a beat
    either took its seq before this point (and the terminal publisher then
    waits for the timer to drain) or meets the barrier when it asks for a seq
    and is refused without consuming one."""
    with mqtt_common.registry_lock(registry_dir):
        with open(barrier_path(job_id, registry_dir), "w") as fh:
            json.dump({"dispatch_token": dispatch_token, "attempt_id": attempt_id}, fh)


def barrier_blocks(job_id: str, registry_dir: str, dispatch_token=None, attempt_id=None) -> bool:
    """True when a terminal barrier was raised for this attempt (a later
    attempt, with its own token, is not blocked). Check it under the registry
    lock when the answer must be atomic with a seq."""
    try:
        with open(barrier_path(job_id, registry_dir)) as fh:
            barrier = json.load(fh)
    except (OSError, ValueError):
        return False
    if not isinstance(barrier, dict):
        return True
    if dispatch_token and barrier.get("dispatch_token"):
        return barrier["dispatch_token"] == dispatch_token
    if attempt_id is not None and barrier.get("attempt_id") is not None:
        return int(barrier["attempt_id"]) == int(attempt_id)
    return True


def start_timer(job_id: str, registry_dir: str, interval: int,
                dispatch_token: Optional[str] = None, attempt_id: Optional[int] = None,
                python: Optional[str] = None, watch_pid: Optional[int] = None) -> int:
    """Spawn the detached timer process and return its pid.

    A timer that already runs for the job is stopped first: a second `started`
    for the same job (a retry of the same command) must not double the rate.
    It is drained as at a terminal event - and killed if it does not exit -
    before its pidfile is handed over (review round 3, codex R3): a beat it had
    already admitted would otherwise still reach the wire after a terminal
    publish that can find only the new timer through the pidfile.
    ``watch_pid`` binds the timer to a process (default: the nearest agent-CLI
    ancestor of this process, if any); 0 disables the watch explicitly.
    Returns 0 without arming when the attempt already ended (terminal record
    or raised barrier): a start that lost the race against the terminal publish
    must not leave a timer behind (review round 4, codex R3).
    """
    if watch_pid is None:
        watch_pid = nearest_agent_cli_ancestor()
    with lifecycle_lock(job_id, registry_dir):
        _stop_timer_locked(job_id, registry_dir, drain=drain_seconds())
        try:
            status = load_job(job_id, registry_dir).get("status")
        except FileNotFoundError:
            status = None
        if status in TERMINAL_STATUSES or barrier_blocks(job_id, registry_dir, dispatch_token, attempt_id):
            logger.info("not arming a heartbeat timer for %s: the attempt already ended", job_id)
            return 0
        return _spawn_timer_locked(job_id, registry_dir, interval, dispatch_token, attempt_id, python, watch_pid)


def _spawn_timer_locked(job_id, registry_dir, interval, dispatch_token, attempt_id, python, watch_pid) -> int:
    """Spawn the timer and publish its pid; the caller holds the lifecycle lock."""
    cmd: List[str] = [python or sys.executable, "-m", "lib_py.delegation.heartbeat",
                      "--job", job_id, "--registry-dir", registry_dir,
                      "--interval", str(int(interval))]
    if dispatch_token:
        cmd += ["--expect-dispatch-token", dispatch_token]
    if attempt_id is not None:
        cmd += ["--attempt-id", str(int(attempt_id))]
    if watch_pid:
        cmd += ["--watch-pid", str(int(watch_pid))]
    env = dict(os.environ)
    # The timer imports lib_py from the same tree as this process.
    pkg_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env["PYTHONPATH"] = pkg_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
    log_path = os.path.join(registry_dir, "{}.heartbeat.log".format(job_id))
    with open(log_path, "ab") as log:
        proc = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT,
                                start_new_session=True, env=env, cwd=os.getcwd())
    pidfile = pidfile_path(job_id, registry_dir)
    try:
        with open(pidfile, "w") as fh:
            fh.write(str(proc.pid))
    except BaseException:
        # Review round 5 (codex F2): a timer whose pid could not be published
        # would be an owner nobody can find or stop. End and reap it here, while
        # the lifecycle lock is still held, and drop only a partial pidfile of
        # this attempted owner (the previous owner's was removed by the stop).
        try:
            proc.kill()
            proc.wait(timeout=5)
        except Exception:
            pass
        try:
            with open(pidfile) as fh:
                partial = fh.read().strip()
            if not partial or str(proc.pid).startswith(partial):
                os.unlink(pidfile)
        except (OSError, ValueError):
            pass
        raise
    logger.info("heartbeat timer pid %d for job %s every %ss (watching pid %s)",
                proc.pid, job_id, interval, watch_pid or "none")
    return proc.pid


def _should_stop(job_id: str, registry_dir: str, dispatch_token, attempt_id,
                 watch_pid: Optional[int] = None) -> Optional[str]:
    if watch_pid and not _pid_alive(watch_pid):
        return "watched agent process {} is gone".format(watch_pid)
    try:
        rec = load_job(job_id, registry_dir)
    except FileNotFoundError:
        return "job record gone"
    if rec.get("status") in TERMINAL_STATUSES:
        return "job is {}".format(rec.get("status"))
    if barrier_blocks(job_id, registry_dir, dispatch_token, attempt_id):
        return "the terminal barrier for this attempt is raised"
    if dispatch_token and rec.get("dispatch_token") != dispatch_token:
        return "dispatch token changed (re-dispatched)"
    if attempt_id is not None and rec.get("attempt_id", rec.get("iteration", 1)) != attempt_id:
        return "attempt changed (re-dispatched)"
    return None


def run(job_id: str, registry_dir: str, interval: int, dispatch_token=None,
        attempt_id=None, max_lifetime: Optional[float] = None,
        watch_pid: Optional[int] = None) -> int:
    """The timer loop. Runs until a stop condition holds."""
    from lib_py.delegation import publish_event  # late import: the CLI imports us

    interval = max(1, int(interval))
    if max_lifetime is None:
        try:
            max_lifetime = float(load_job(job_id, registry_dir).get("timeout_sec", 3600)) + interval
        except FileNotFoundError:
            return 1
    stop = {"flag": False}

    def _on_term(_signum, _frame):
        stop["flag"] = True

    signal.signal(signal.SIGTERM, _on_term)
    signal.signal(signal.SIGINT, _on_term)

    started = time.monotonic()
    own_pidfile = pidfile_path(job_id, registry_dir)
    # Fixed cadence: the next beat is due `interval` after the previous one
    # was *due*, not after its publish returned, so a slow broker (seconds per
    # ACK) does not stretch the period the contract promises. Each beat is
    # recorded on disk before its wire publish, so the on-disk cadence holds
    # even when the wire lags. Sleep first: `started` was just published.
    next_due = time.monotonic() + interval
    in_flight: List[threading.Thread] = []
    while not stop["flag"]:
        while time.monotonic() < next_due and not stop["flag"]:
            time.sleep(min(0.5, max(0.0, next_due - time.monotonic())))
        if stop["flag"]:
            break
        next_due += interval
        if next_due < time.monotonic():
            next_due = time.monotonic() + interval  # fell far behind: resync
        reason = _should_stop(job_id, registry_dir, dispatch_token, attempt_id, watch_pid)
        if reason:
            logger.info("heartbeat timer for %s stopping: %s", job_id, reason)
            break
        if time.monotonic() - started > max_lifetime:
            logger.info("heartbeat timer for %s stopping: lifetime %.0fs exceeded", job_id, max_lifetime)
            break
        argv = ["--registry-dir", registry_dir, "--job", job_id, "--event", "progress",
                "--detail", HEARTBEAT_DETAIL, "--attempts", "2"]
        if dispatch_token:
            argv += ["--expect-dispatch-token", dispatch_token]
        if attempt_id is not None:
            argv += ["--attempt-id", str(attempt_id)]
        # Each beat publishes in its own thread: the local record (seq, HMAC,
        # events.log) is taken at once, and the wire publish - seconds per ACK
        # on a slow broker - must not delay the next beat. In-flight beats are
        # capped so a dead broker cannot pile up threads without bound.
        in_flight[:] = [t for t in in_flight if t.is_alive()]
        if len(in_flight) >= MAX_IN_FLIGHT:
            logger.warning("heartbeat for %s skipped: %d publishes still in flight", job_id, len(in_flight))
            continue

        def _beat(argv=argv):
            try:
                # _beat=True: admitted under the seq lock, never after this
                # attempt's terminal barrier (review round 2, codex R3)
                rc = publish_event.main(argv, _beat=True)
            except Exception as exc:  # pragma: no cover - defensive
                logger.warning("heartbeat publish raised: %s", exc)
                rc = 2
            if rc == 4:
                logger.info("heartbeat publish refused (terminal barrier, terminal record or stale attempt); "
                            "timer for %s will stop", job_id)
                stop["flag"] = True

        t = threading.Thread(target=_beat, daemon=True)
        t.start()
        in_flight.append(t)
    for t in in_flight:
        t.join(timeout=15)
    # Under the lifecycle lock: a replacement may be writing its own pid to
    # this file right now, and that owner's pidfile must not be deleted. Never
    # *wait* for the lock here (review round 5, codex F1): a stopper holds it
    # while it waits for this very exit, so waiting would only end in SIGKILL at
    # its drain deadline. When the lock is busy its holder is an ownership
    # operation, and that operation removes the pidfile itself.
    try:
        with lifecycle_lock(job_id, registry_dir, blocking=False) as owned:
            if owned:
                with open(own_pidfile) as fh:
                    mine = fh.read().strip() == str(os.getpid())
                if mine:
                    os.unlink(own_pidfile)
    except (OSError, ValueError):
        pass
    return 0


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="automatic heartbeat timer for a delegated job")
    parser.add_argument("--job", required=True)
    parser.add_argument("--registry-dir", default=default_registry_dir())
    parser.add_argument("--interval", type=int, required=True, help="heartbeat period N in seconds")
    parser.add_argument("--expect-dispatch-token", default=None)
    parser.add_argument("--attempt-id", type=int, default=None)
    parser.add_argument("--max-lifetime", type=float, default=None)
    parser.add_argument("--watch-pid", type=int, default=None,
                        help="stop once this process (the agent) is gone")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args(argv)
    mqtt_common.ensure_env_loaded()
    mqtt_common.setup_logging(logging.DEBUG if args.verbose else logging.INFO)
    return run(args.job, args.registry_dir, args.interval, args.expect_dispatch_token,
               args.attempt_id, args.max_lifetime, args.watch_pid)


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