"""Mid-run liveness probe for pending workers (`okstra worker-liveness`).

The CLI wrappers reap a silent CLI themselves via their idle watchdog, but the
lead's own wait had no equivalent: it polled Result Paths, which only change at
the very end, so a worker that died at minute 3 was still waited on until its
20–40 minute deadline. Two contracts named that failure without a mechanism —
`adapters/claude-code.md` ("a missing or stale heartbeat consumes the same
one-retry budget") and the absence of any `did-not-launch` status at all.

This module is that mechanism. It reports, not decides: the lead reads the
verdict and spends its existing one-retry budget.

Every probe is selected the same way — ``--team-state`` + ``--worker`` — and the
worker row's ``livenessMode`` decides which of the two artifacts answers:

* ``audit-heartbeat`` — an in-process worker's audit sidecar. Stale past the
  heartbeat cadence (or present with no heartbeat at all) means the worker
  hung. Uses the same line shape and budget the Phase 7 validator applies.
* ``wrapper-status`` — a CLI-wrapper assignment. Neither the worker's
  `<prompt>.log` nor `<prompt>.status.json` past the launch grace means the
  wrapper never ran.

Both graces start at the persisted dispatch timestamp (``workers[].startedAt``),
never at an artifact's mtime. That anchor is why the selector needs team-state:
the audit sidecar is reused when a worker is re-dispatched, so without knowing
when *this* dispatch started, the previous dispatch's last heartbeat reads as
this worker's newest signal and a freshly launched worker probes `stalled`.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

from .wrapper_status import log_path_for_prompt

from okstra_ctl.dispatch_state import (
    DispatchError,
    LIVENESS_AUDIT_HEARTBEAT,
    LIVENESS_WRAPPER_STATUS,
    load_json_object,
)
from okstra_ctl.worker_heartbeat import (
    HEARTBEAT_MAX_GAP_SECONDS,
    Heartbeat,
    latest_heartbeat,
    max_gap_seconds_after,
)

DEFAULT_LAUNCH_GRACE_SECONDS = 60
DEFAULT_POLL_INTERVAL_SECONDS = 20.0
DEFAULT_WAIT_TIMEOUT_SECONDS = 2400.0

# A budget breach is one observation, and "slow" and "dead" are only
# distinguishable across two. On a breach the probe re-reads the sidecar this
# far into the future — a fraction of the stage's own budget, so a stage with a
# longer budget also gets a longer confirmation. Measured false positives this
# absorbs (dev-10400): `analysis` 386s against a 360s budget,
# `data-json-write-start` 1602s against 1260s.
DEFAULT_STALL_CONFIRM_RATIO = 0.5


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _log_path(prompt: Path) -> Path:
    """The wrapper's live log, named as okstra-*-exec.sh names it."""
    return log_path_for_prompt(prompt)


def probe_heartbeat(
    sidecar: Path,
    dispatched_at: datetime,
    now: datetime,
    max_idle: float,
    grace: float,
) -> dict:
    """Liveness of one in-process worker, read from its audit sidecar.

    ``max_idle`` is the floor. A stage whose work is one uninterruptible tool
    call carries its own, longer budget (``max_gap_seconds_after``) — during a
    single large Write the worker cannot append a heartbeat at all.
    """
    probe = {"kind": "heartbeat", "path": str(sidecar)}
    if not sidecar.is_file():
        # The worker writes the sidecar early but not instantly; absence is the
        # launch probe's business, not a hang.
        return {**probe, "state": "pending", "reason": "audit sidecar not written yet"}
    beat = latest_heartbeat(sidecar)
    if beat is None:
        return {
            **probe,
            "state": "stalled",
            "reason": "audit sidecar carries no `- PROGRESS:` heartbeat",
        }
    if beat.at < dispatched_at:
        return _probe_before_first_beat(probe, beat, dispatched_at, now, grace)
    budget = max(max_idle, max_gap_seconds_after(beat.stage))
    idle = (now - beat.at).total_seconds()
    state = "stalled" if idle > budget else "live"
    return {
        **probe,
        "state": state,
        "idleSeconds": int(idle),
        "lastHeartbeat": beat.at.isoformat(),
        "lastStage": beat.stage,
        "budgetSeconds": int(budget),
        "reason": (
            f"no heartbeat for {int(idle)}s after stage `{beat.stage}` "
            f"(budget {int(budget)}s)"
            if state == "stalled"
            else ""
        ),
    }


def _probe_before_first_beat(
    probe: dict,
    beat: Heartbeat,
    dispatched_at: datetime,
    now: datetime,
    grace: float,
) -> dict:
    """Verdict when the newest heartbeat predates this dispatch.

    The audit sidecar is reused on re-dispatch, so that beat belongs to the
    previous attempt and this one has produced no signal yet — the launch
    grace decides, and the idle measure runs from the dispatch, never from a
    heartbeat this worker never wrote.
    """
    waited = (now - dispatched_at).total_seconds()
    if waited <= grace:
        return {
            **probe,
            "state": "pending",
            "waitedSeconds": int(waited),
            "lastStage": beat.stage,
            "reason": "within launch grace",
        }
    return {
        **probe,
        "state": "stalled",
        "waitedSeconds": int(waited),
        "lastStage": beat.stage,
        "reason": (
            f"no heartbeat for this dispatch {int(waited)}s after it started "
            f"(grace {int(grace)}s); newest beat `{beat.stage}` predates it"
        ),
    }


def probe_launch(
    prompt: Path, dispatched_at: datetime, now: datetime, grace: float
) -> dict:
    """Whether the wrapper behind *prompt* ever started."""
    log, status = _log_path(prompt), Path(f"{prompt}.status.json")
    probe = {"kind": "launch", "path": str(prompt)}
    if log.exists() or status.exists():
        return {**probe, "state": "live", "reason": ""}
    waited = (now - dispatched_at).total_seconds()
    if waited <= grace:
        return {
            **probe,
            "state": "pending",
            "waitedSeconds": int(waited),
            "reason": "within launch grace",
        }
    return {
        **probe,
        "state": "did-not-launch",
        "waitedSeconds": int(waited),
        "reason": (
            f"neither {log.name} nor {status.name} exists {int(waited)}s after the "
            f"dispatch started (grace {int(grace)}s)"
        ),
    }


@dataclass(frozen=True)
class ProbeTarget:
    """One pending worker resolved from team-state: which artifact answers for
    it, where that artifact is, when this dispatch started, and the result file
    whose arrival means this worker is done."""
    liveness_mode: str
    artifact: Path
    dispatched_at: datetime
    result_path: Path | None = None


def _confirm_window(probe: dict, stall_confirm: float | None) -> float:
    """Seconds to wait before a budget breach becomes a verdict.

    Only a breach carries ``budgetSeconds``. The other stalled shapes — a
    sidecar with no heartbeat at all, a newest beat that predates this dispatch
    — are not "the worker is mid-tool-call", so waiting tells us nothing new
    about them.
    """
    if "budgetSeconds" not in probe:
        return 0.0
    if stall_confirm is not None:
        return max(0.0, float(stall_confirm))
    return probe["budgetSeconds"] * DEFAULT_STALL_CONFIRM_RATIO


def probe_one(
    target: ProbeTarget,
    *,
    now: datetime,
    max_idle: float,
    launch_grace: float,
    stall_confirm: float | None = None,
    sleep: Callable[[float], None] = time.sleep,
    clock: Callable[[], datetime] = _utc_now,
) -> dict:
    """One worker's verdict, with a budget breach confirmed before it stands.

    The wait this costs is bounded by the confirmation window; what the probe
    exists to avoid is paying ``DEFAULT_WAIT_TIMEOUT_SECONDS`` for a worker that
    died early, and that is still never paid.
    """
    if target.liveness_mode != LIVENESS_AUDIT_HEARTBEAT:
        return probe_launch(target.artifact, target.dispatched_at, now, launch_grace)
    probe = probe_heartbeat(
        target.artifact, target.dispatched_at, now, max_idle, launch_grace
    )
    if probe["state"] != "stalled":
        return probe
    window = _confirm_window(probe, stall_confirm)
    if window <= 0:
        return probe
    sleep(window)
    confirmed = probe_heartbeat(
        target.artifact, target.dispatched_at, clock(), max_idle, launch_grace
    )
    if confirmed.get("lastHeartbeat") == probe.get("lastHeartbeat"):
        return confirmed
    # The question the window asks is whether the heartbeat moved, not whether
    # the re-read is healthy on its own terms. A beat opening a stage with a
    # smaller budget than the window we just slept reads as stale the instant it
    # lands, which would call a worker dead for proving it is alive.
    return {**confirmed, "state": "live", "reason": ""}


def probe_all(
    targets: list[ProbeTarget],
    *,
    now: datetime,
    max_idle: float,
    launch_grace: float,
    stall_confirm: float | None = None,
    sleep: Callable[[float], None] = time.sleep,
    clock: Callable[[], datetime] = _utc_now,
) -> dict:
    probes = [
        probe_one(
            t, now=now, max_idle=max_idle, launch_grace=launch_grace,
            stall_confirm=stall_confirm, sleep=sleep, clock=clock,
        )
        for t in targets
    ]
    unhealthy = [p for p in probes if p["state"] in ("stalled", "did-not-launch")]
    return {"ok": not unhealthy, "checkedAt": now.isoformat(), "probes": probes,
            "unhealthy": unhealthy}


def result_ready(target: ProbeTarget) -> bool:
    """Whether this worker's result file has landed with content in it."""
    path = target.result_path
    return bool(path and path.is_file() and path.stat().st_size > 0)


def wait_for_results(
    targets: list[ProbeTarget],
    *,
    max_idle: float,
    launch_grace: float,
    interval: float,
    timeout: float,
    stall_confirm: float | None = None,
    clock: Callable[[], datetime] = _utc_now,
    sleep: Callable[[float], None] = time.sleep,
) -> dict:
    """Poll until every result file lands, a worker dies, or the deadline passes.

    This is the loop a lead would otherwise hand-write in Bash at each dispatch,
    and every hand-written one has to re-derive the same two facts: what "done"
    means (the persisted ``resultPath``, not a guessed filename) and what
    "dead" means (`probe_all`, whose graces run from ``startedAt`` — never from
    an artifact's mtime, which a re-dispatched worker inherits from its previous
    attempt and reads as instantly stale).

    ``outcome`` is ``completed`` / ``unhealthy`` / ``timeout``.
    """
    started = clock()
    while True:
        now = clock()
        result = probe_all(
            targets, now=now, max_idle=max_idle, launch_grace=launch_grace,
            stall_confirm=stall_confirm, sleep=sleep, clock=clock,
        )
        pending = [
            str(t.result_path) for t in targets if not result_ready(t)
        ]
        waited = (now - started).total_seconds()
        if not pending:
            outcome = "completed"
        elif result["unhealthy"]:
            outcome = "unhealthy"
        elif waited >= timeout:
            outcome = "timeout"
        else:
            sleep(interval)
            continue
        return {
            **result,
            "ok": outcome == "completed",
            "outcome": outcome,
            "pending": pending,
            "waitedSeconds": int(waited),
        }


def _parse_utc(value: object, label: str) -> datetime:
    if not isinstance(value, str) or not value:
        raise DispatchError(f"{label} must be a UTC ISO timestamp")
    try:
        instant = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as exc:
        raise DispatchError(f"{label} must be a UTC ISO timestamp") from exc
    if instant.tzinfo is None or instant.utcoffset() != timezone.utc.utcoffset(instant):
        raise DispatchError(f"{label} must be a UTC ISO timestamp")
    return instant.astimezone(timezone.utc)


def _project_root_for_team_state(team_state_path: Path) -> Path:
    for parent in team_state_path.resolve().parents:
        if parent.name == ".okstra":
            return parent.parent
    raise DispatchError(
        f"team-state is outside a project .okstra directory: {team_state_path}"
    )


_ARTIFACT_FIELD_BY_MODE = {
    LIVENESS_AUDIT_HEARTBEAT: "auditSidecarPath",
    LIVENESS_WRAPPER_STATUS: "promptPath",
}


def _worker_row(team_state_path: Path, worker_id: str) -> dict:
    state = load_json_object(team_state_path, "team-state")
    workers = state.get("workers")
    if not isinstance(workers, list):
        raise DispatchError(f"team-state workers must be an array: {team_state_path}")
    worker = next(
        (
            row for row in workers
            if isinstance(row, dict) and row.get("workerId") == worker_id
        ),
        None,
    )
    if worker is None:
        raise DispatchError(f"team-state has no workerId={worker_id}: {team_state_path}")
    return worker


def probe_target(team_state_value: str, worker_id: str) -> ProbeTarget:
    """Resolve one worker's probe target from team-state.

    ``livenessMode`` is authoritative — never infer the transport from the
    provider or a filename, which is how a lead ends up probing an in-process
    worker for a wrapper log that will never exist.
    """
    team_state_path = Path(team_state_value).resolve()
    worker = _worker_row(team_state_path, worker_id)
    mode = worker.get("livenessMode")
    field = _ARTIFACT_FIELD_BY_MODE.get(mode) if isinstance(mode, str) else None
    if field is None:
        allowed = ", ".join(sorted(_ARTIFACT_FIELD_BY_MODE))
        raise DispatchError(
            f"worker {worker_id} has missing or unknown livenessMode {mode!r}; "
            f"expected one of: {allowed}"
        )
    artifact_value = worker.get(field)
    if not isinstance(artifact_value, str) or not artifact_value.strip():
        raise DispatchError(f"worker {worker_id} has no {field}")
    project_root = _project_root_for_team_state(team_state_path)
    artifact = Path(artifact_value)
    if not artifact.is_absolute():
        artifact = project_root / artifact
    dispatched_at = _parse_utc(worker.get("startedAt"), f"worker {worker_id} startedAt")
    result_value = worker.get("resultPath")
    result_path = None
    if isinstance(result_value, str) and result_value.strip():
        result_path = Path(result_value)
        if not result_path.is_absolute():
            result_path = project_root / result_path
    return ProbeTarget(mode, artifact, dispatched_at, result_path)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="okstra worker-liveness",
        description="Report whether pending workers are still alive (read-only).",
    )
    parser.add_argument("--team-state", action="append", default=[],
                        help="team-state path for a pending worker (repeatable)")
    parser.add_argument("--worker", action="append", default=[],
                        help="worker id paired with --team-state (repeatable)")
    parser.add_argument("--max-idle", type=float, default=HEARTBEAT_MAX_GAP_SECONDS,
                        help="heartbeat staleness budget in seconds")
    parser.add_argument("--launch-grace", type=float, default=DEFAULT_LAUNCH_GRACE_SECONDS,
                        help="seconds a worker may take to write its first artifact")
    parser.add_argument("--json", action="store_true", help="emit JSON (always on)")
    parser.add_argument(
        "--wait", action="store_true",
        help=(
            "poll until every worker's persisted resultPath lands (exit 0), one "
            "worker probes unhealthy (exit 1), or --timeout passes (exit 2). "
            "Use this instead of hand-writing a Bash poll loop."
        ),
    )
    parser.add_argument("--interval", type=float, default=DEFAULT_POLL_INTERVAL_SECONDS,
                        help="--wait poll interval in seconds")
    parser.add_argument("--timeout", type=float, default=DEFAULT_WAIT_TIMEOUT_SECONDS,
                        help="--wait deadline in seconds")
    parser.add_argument(
        "--stall-confirm", type=float, default=None,
        help=(
            "seconds to re-check a heartbeat budget breach before calling it "
            "stalled (default: half that stage's budget; 0 disables). A slow "
            "worker appends its next heartbeat inside this window; a dead one "
            "does not."
        ),
    )
    args = parser.parse_args(argv)

    if len(args.team_state) != len(args.worker):
        parser.error("each --team-state must have one paired --worker")
    if not args.team_state:
        parser.error("pass at least one --team-state/--worker pair")

    try:
        targets = [
            probe_target(team_state, worker)
            for team_state, worker in zip(args.team_state, args.worker, strict=True)
        ]
    except DispatchError as exc:
        parser.error(str(exc))

    if not args.wait:
        result = probe_all(
            targets,
            now=_utc_now(),
            max_idle=args.max_idle,
            launch_grace=args.launch_grace,
            stall_confirm=args.stall_confirm,
        )
        print(json.dumps(result, ensure_ascii=False, indent=2))
        # Non-zero on an unhealthy worker so a poll loop can branch on the exit
        # code without parsing the JSON.
        return 0 if result["ok"] else 1

    unwaitable = [
        worker for target, worker in zip(targets, args.worker, strict=True)
        if target.result_path is None
    ]
    if unwaitable:
        parser.error(
            f"--wait needs a resultPath in team-state for: {', '.join(unwaitable)}. "
            "Waiting on a guessed filename is what --wait exists to prevent."
        )

    result = wait_for_results(
        targets,
        max_idle=args.max_idle,
        launch_grace=args.launch_grace,
        interval=args.interval,
        timeout=args.timeout,
        stall_confirm=args.stall_confirm,
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return {"completed": 0, "unhealthy": 1}.get(result["outcome"], 2)


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))