"""Wraps `herdr agent wait`. Observes; never concludes.

Live contract (herdr 0.8.2, re-confirmed while planning):

    herdr agent wait <TARGET> --until <STATUS>... --timeout <MS>

  * `--timeout` is MILLISECONDS. The parameter here is named `timeout_ms` so the
    unit cannot be misread at a call site; there is deliberately no seconds form.
  * `--until` is always passed explicitly. Omitting it matches idle/done/blocked,
    silently excluding `working`, which is a state the watcher cares about.
  * success -> exit 0, snapshot JSON on STDOUT.
  * timeout and unknown target -> exit 1, error JSON on STDERR, stdout empty.
    The stream split is the contract: an error is never read from stdout.

A single wait is one synchronous barrier returning one final snapshot, so
continuous observation is built by repeating it, not by streaming.
"""
from __future__ import annotations

import json
import logging
import os
import subprocess
from typing import Any, Dict, Optional, Sequence

from lib_py.loop.evidence import Observation

logger = logging.getLogger("loop.watcher")

DEFAULT_UNTIL = ("idle", "done", "blocked")
VALID_STATES = ("idle", "working", "blocked", "done", "unknown")

# Outcome kinds. `timeout` is an ordinary no-event result, not a failure.
MATCHED = "matched"
TIMEOUT = "timeout"
AGENT_NOT_FOUND = "agent_not_found"
TRANSPORT_ERROR = "transport_error"

BACKOFF_START_S = 1.0
BACKOFF_MAX_S = 16.0

# ISSUE-23: ids a herdr pane exports about ITS OWN server. Passing them to a
# call aimed at a different server makes herdr scope the lookup to a workspace
# that does not exist there, and the target is reported as missing. The shell
# side drops them in mam_herdr; this is the one python path that runs herdr
# without going through it.
CALLER_SCOPE_VARS = ("HERDR_WORKSPACE_ID", "HERDR_PANE_ID", "HERDR_TAB_ID")


def caller_server(environ=None):
    """The server whose pane this process runs in, "" when not in a pane.

    HERDR_SESSION when present; otherwise the shape of HERDR_SOCKET_PATH, which
    herdr exports to every pane (HERDR_SESSION is absent in a default-server
    pane on herdr 0.8.2 - the orchestrator's usual home, and where the #23 leak
    comes from). Same convention as lib.sh's _herdr_caller_server:

        .../.config/herdr/herdr.sock                 -> "default"
        .../.config/herdr/sessions/<name>/herdr.sock -> <name>

    Matched by shape, not by $HOME. Anything else is "" (not in a pane).
    """
    src = os.environ if environ is None else environ
    sess = (src.get("HERDR_SESSION") or "").strip()
    if sess:
        return sess
    sock = (src.get("HERDR_SOCKET_PATH") or "").strip()
    parts = sock.split("/")
    if len(parts) >= 4 and parts[-1] == "herdr.sock":
        if parts[-3:-1] == [".config", "herdr"]:
            return "default"
        if len(parts) >= 5 and parts[-4:-2] == ["herdr", "sessions"] and parts[-5] == ".config":
            return parts[-2]
    return ""


def scoped_env(session=None, environ=None):
    """The environment for a wait aimed at `session`.

    Unchanged unless this process runs inside a pane of another server: only
    then are the caller's own workspace/pane/tab ids meaningless for the call
    and dropped. No pane and same-server keep every value.
    """
    src = os.environ if environ is None else environ
    caller = caller_server(src)
    target = (session or "").strip()
    # `default` is the same sentinel build_argv reads: no --session flag, so
    # herdr answers on the ambient server - the caller's own.
    if not caller or not target or target == "default" or caller == target:
        return dict(src)
    return dict((k, v) for k, v in src.items() if k not in CALLER_SCOPE_VARS)


class WaitOutcome(object):
    __slots__ = ("kind", "agent", "error_code", "detail", "raw")

    def __init__(self, kind, agent=None, error_code=None, detail="", raw=None):
        self.kind = kind
        self.agent = agent
        self.error_code = error_code
        self.detail = detail
        self.raw = raw

    @property
    def healthy(self) -> bool:
        """A timeout is healthy: the socket answered, the state just did not change."""
        return self.kind in (MATCHED, TIMEOUT, AGENT_NOT_FOUND)

    def __repr__(self):  # pragma: no cover - debugging aid
        return "WaitOutcome({!r}, error_code={!r})".format(self.kind, self.error_code)


def build_argv(target: str, until: Sequence[str] = DEFAULT_UNTIL,
               timeout_ms: Optional[int] = None,
               session: Optional[str] = None) -> list:
    if not target:
        raise ValueError("wait target is required")
    if session and session != "default":
        argv = ["herdr", "--session", session, "agent", "wait", target]
    else:
        argv = ["herdr", "agent", "wait", target]
    for state in until:
        if state not in VALID_STATES:
            raise ValueError("unsupported --until state: {!r}".format(state))
        argv += ["--until", state]
    if timeout_ms is not None:
        if int(timeout_ms) <= 0:
            raise ValueError("timeout_ms must be positive milliseconds")
        argv += ["--timeout", str(int(timeout_ms))]
    return argv


def _parse(stream: str) -> Optional[Dict[str, Any]]:
    try:
        payload = json.loads(stream)
        return payload if isinstance(payload, dict) else None
    except Exception:
        return None


def wait_once(target: str, until: Sequence[str] = DEFAULT_UNTIL,
              timeout_ms: Optional[int] = None, session: Optional[str] = None,
              run=subprocess.run) -> WaitOutcome:
    """One barrier call. `run` is injectable so tests never touch a real herdr."""
    argv = build_argv(target, until, timeout_ms, session=session)
    try:
        proc = run(argv, capture_output=True, text=True, env=scoped_env(session))
    except Exception as exc:
        return WaitOutcome(TRANSPORT_ERROR, detail="herdr could not be executed: {}".format(exc))

    if proc.returncode == 0:
        payload = _parse(proc.stdout or "")
        agent = ((payload or {}).get("result") or {}).get("agent")
        if not isinstance(agent, dict):
            return WaitOutcome(TRANSPORT_ERROR,
                               detail="exit 0 without a parseable agent snapshot",
                               raw=proc.stdout)
        return WaitOutcome(MATCHED, agent=agent, raw=payload)

    # Non-zero: the payload is on stderr by contract, and stdout is empty.
    payload = _parse(proc.stderr or "")
    code = ((payload or {}).get("error") or {}).get("code")
    if code == "timeout":
        return WaitOutcome(TIMEOUT, error_code=code, raw=payload)
    if code == "agent_not_found":
        return WaitOutcome(AGENT_NOT_FOUND, error_code=code, raw=payload)
    return WaitOutcome(TRANSPORT_ERROR, error_code=code,
                       detail="unrecognised failure (rc={})".format(proc.returncode),
                       raw=proc.stderr)


def observation_from(outcome: WaitOutcome, capture_ref=None) -> Optional[Observation]:
    """Turn a matched snapshot into evidence. Non-matches carry no agent state."""
    if outcome.kind != MATCHED or not outcome.agent:
        return None
    agent = outcome.agent
    return Observation(
        "herdr",
        {
            "agent_status": agent.get("agent_status"),
            "interactive_ready": agent.get("interactive_ready"),
            "pane_id": agent.get("pane_id"),
            "state_change_seq": agent.get("state_change_seq"),
        },
        capture_ref=capture_ref,
    )


class ReconnectBackoff(object):
    """Doubling backoff for transport failures, reset by any healthy answer.

    The snapshot cache is dropped on failure on purpose: after a socket gap the
    watcher must re-observe rather than reason from a state that may have moved
    while it was blind. That is what makes `hold_reconnect` the honest decision.
    """

    def __init__(self, start=BACKOFF_START_S, maximum=BACKOFF_MAX_S):
        self.start = start
        self.maximum = maximum
        self.current = 0.0
        self.consecutive_failures = 0
        self.cached_snapshot = None

    def record_success(self, snapshot=None) -> None:
        self.consecutive_failures = 0
        self.current = 0.0
        if snapshot is not None:
            self.cached_snapshot = snapshot

    def record_failure(self) -> float:
        self.consecutive_failures += 1
        self.cached_snapshot = None  # stale evidence must not survive a gap
        self.current = self.start if self.current == 0.0 else min(self.current * 2.0, self.maximum)
        return self.current

    def sequence(self, n: int) -> list:
        """The first n delays, for documentation and for tests to pin."""
        saved = (self.current, self.consecutive_failures, self.cached_snapshot)
        try:
            self.current, self.consecutive_failures = 0.0, 0
            return [self.record_failure() for _ in range(n)]
        finally:
            self.current, self.consecutive_failures, self.cached_snapshot = saved


def main(argv: Optional[Sequence[str]] = None) -> int:
    import argparse
    import os
    import sys

    parser = argparse.ArgumentParser(description="watcher CLI for herdr agent wait")
    sub = parser.add_subparsers(dest="command", required=True)
    p_wait = sub.add_parser("wait", help="wait for agent state")
    p_wait.add_argument("target", help="target agent name")
    p_wait.add_argument("--session", default=None, help="herdr session name")
    p_wait.add_argument("--until", action="append", default=[], help="until state(s)")
    p_wait.add_argument("--timeout-ms", type=int, default=None, help="timeout in ms")

    args = parser.parse_args(argv)
    session = args.session or os.environ.get("HERDR_SESSION_NAME") or None
    until = args.until or DEFAULT_UNTIL
    outcome = wait_once(args.target, until=until, timeout_ms=args.timeout_ms, session=session)
    agent_st = None
    if outcome.agent:
        agent_st = outcome.agent.get("agent_status")
    print(json.dumps({
        "kind": outcome.kind,
        "agent_status": agent_st,
        "error_code": outcome.error_code,
    }))
    rc_map = {
        MATCHED: 0,
        TIMEOUT: 2,
        AGENT_NOT_FOUND: 3,
        TRANSPORT_ERROR: 4,
    }
    return rc_map.get(outcome.kind, 4)


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