"""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, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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 any(path.is_file() and path.stat().st_mtime >= dispatched_at.timestamp()
           for path in (log, status)):
        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
                and path.stat().st_mtime >= target.dispatched_at.timestamp())


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 _dispatch_worker_id(record: Mapping[str, Any]) -> str:
    """이 디스패치 행이 말하는 워커 id.

    execution-identity v2 행에는 `workerId` 가 실리지 않는다
    (`dispatch_core._dispatch_record`). 그 경우 `assignmentRef` 의 마지막 마디가
    워커 id 이므로 — `initial/codex-verifier`, `critic/acceptance` — 거기서 읽는다.
    """
    worker_id = record.get("workerId")
    if isinstance(worker_id, str) and worker_id.strip():
        return worker_id.strip()
    assignment_ref = record.get("assignmentRef")
    if isinstance(assignment_ref, str) and assignment_ref.strip():
        return assignment_ref.strip().rsplit("/", 1)[-1]
    return ""


def _dispatch_fallback(state: Mapping[str, Any], worker_id: str, dispatch_id: str = "") -> dict:
    """명시한 배정 또는 하나로 확정되는 구형 배정만 선택한다."""
    records = state.get("workerDispatches") or []
    matches = [
        row for row in records if isinstance(row, Mapping)
        and (row.get("dispatchId") == dispatch_id if dispatch_id
             else _dispatch_worker_id(row) == worker_id)
    ]
    if len(matches) > 1:
        raise DispatchError(f"multiple dispatches for {worker_id or dispatch_id}; pass --dispatch-id")
    if dispatch_id and not matches:
        raise DispatchError(f"team-state has no dispatchId={dispatch_id}")
    return dict(matches[0]) if matches else {}


def _worker_row(team_state_path: Path, worker_id: str, dispatch_id: str = "") -> dict:
    """새 배정의 필드는 원자적으로 읽고 구형 단일 기록만 보완한다."""
    state = load_json_object(team_state_path, "team-state")
    dispatch = _dispatch_fallback(state, worker_id, dispatch_id)
    if dispatch.get("dispatchId") and dispatch.get("startedAt"):
        result = dispatch.get("resultPath")
        root = _project_root_for_team_state(team_state_path)
        if result and any(
            isinstance(row, Mapping) and row.get("dispatchId") != dispatch["dispatchId"]
            and row.get("status") not in {"completed", "error", "timeout", "not-run"}
            and row.get("resultPath")
            and (root / row["resultPath"]).resolve() == (root / result).resolve()
            for row in state.get("workerDispatches", [])
        ):
            raise DispatchError("resultPath is shared by dispatches; use an attempt-specific result path")
        return dispatch
    workers = state.get("workers")
    if not isinstance(workers, list):
        raise DispatchError(f"team-state workers must be an array: {team_state_path}")
    selected_worker = _dispatch_worker_id(dispatch) if dispatch_id else worker_id
    worker = next((row for row in workers if isinstance(row, dict)
                   and row.get("workerId") == selected_worker), None)
    if worker is None:
        raise DispatchError(f"team-state has no workerId={selected_worker}: {team_state_path}")
    if dispatch and worker.get("promptPath") and dispatch.get("promptPath"):
        root = _project_root_for_team_state(team_state_path)
        if (root / worker["promptPath"]).resolve() != (root / dispatch["promptPath"]).resolve():
            raise DispatchError("dispatch has no startedAt and roster promptPath differs; repair the recorded attempt")
    merged = dict(worker)
    for key, value in dispatch.items():
        current = merged.get(key)
        if key not in merged or (isinstance(current, str) and not current.strip()):
            merged[key] = value
    return merged


def probe_target(team_state_value: str, worker_id: str, *, dispatch_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, dispatch_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)


_CLI_EPILOG = r"""Usage:
  okstra worker-liveness [--team-state <path> --worker <worker-id>]...
                         [--max-idle <seconds>] [--launch-grace <seconds>] [--json]
  okstra worker-liveness --wait [--team-state <path> --worker <worker-id>]...
                         [--interval <seconds>] [--timeout <seconds>]

--wait polls until every named worker's persisted resultPath lands (exit 0), one
          worker probes unhealthy (exit 1), or --timeout passes (exit 2).
          Defaults: --interval 20, --timeout 2400. Run it as a background
          command and branch on the exit code. Use this instead of writing a
          poll loop: both the "is it done" test (the persisted resultPath, not a
          filename you assemble) and the "is it dead" test live in here, and a
          hand-written loop re-derives them and gets one wrong. A worker row
          without a resultPath is refused rather than waited on by guess.

--team-state and --worker select one pending worker. The worker row's
          livenessMode picks the probe: audit-heartbeat reads the in-process
          worker audit sidecar path and reports stalled past the heartbeat
          cadence; wrapper-status reports did-not-launch when no prompt sibling
          .log or .status.json exists past the launch grace. Both graces begin
          at workers[].startedAt, never at an artifact mtime — the audit sidecar
          is reused on re-dispatch, so the previous attempt's last heartbeat is
          not this dispatch's signal.

Output: JSON { ok, checkedAt, probes[], unhealthy[] }, plus { outcome, pending[],
waitedSeconds } under --wait. Exit 1 when any worker is stalled or did not
launch, so a caller can branch on the exit code. Read-only: it reports, it never
kills or re-dispatches.
"""


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        epilog=_CLI_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        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("--dispatch-id", action="append", default=[],
                        help="exact dispatch id paired with --team-state; do not mix with --worker")
    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 args.worker and args.dispatch_id:
        parser.error("use --worker or --dispatch-id, not both")
    selectors = args.dispatch_id or args.worker
    if len(args.team_state) != len(selectors):
        parser.error("each --team-state must have one paired --worker or --dispatch-id")
    if not args.team_state:
        parser.error("pass at least one --team-state/--worker pair")

    try:
        targets = [
            probe_target(team_state, "" if args.dispatch_id else worker,
                         dispatch_id=worker if args.dispatch_id else "")
            for team_state, worker in zip(args.team_state, selectors, 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, selectors, 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:]))