#!/usr/bin/python3
# Pinned to /usr/bin/python3 rather than `env python3` so the script is
# immune to an active user virtualenv shadowing system site-packages
# (where PyYAML lives on Debian — `apt install python3-yaml`).
"""
elway — a mini playbook runner over SSH.

Named for John Elway: quarterbacks run plays. You hand it a play (a single
shell command or an upload), or a playbook (a YAML list of steps + an
optional verify phase), and it shoots them across the ssh link with live
output, structured reporting, and stop-on-error semantics.

Conventions:
  * Single host per run. Fleet-wide = shell loop outside.
  * Leverages your ~/.ssh/config aliases. No inventory file.
  * SSH connection reuse via ControlMaster for multi-step speed.
  * Sudo password (if needed) prompted once at start, reused via `sudo -S`,
    never written to disk, never logged.
  * Simple `{{ var }}` substitution. No Jinja, no filters, no loops.

Quick starts:
  elway irv-ml1 --shell 'docker compose ls'
  elway irv-ml1 --shell 'mkdir -p /worktank/foo' --sudo
  elway irv-ml1 --upload stacks/comfyui/compose.yaml:/opt/docker/compose/comfyui/compose.yaml
  elway irv-ml1 --playbook playbooks/deploy-stack.yaml --var stack=comfyui

Playbook schema (YAML):
  vars:                       # optional inline defaults; --var CLI overrides
    stack: comfyui
  steps:                      # required list
    - name: <label>           # required
      shell: <command>        # exactly one of shell / upload
      # or
      upload:
        src: <local path>
        dest: <remote path>
        mode: "0644"          # optional
        owner: "user:group"   # optional, sudo only; default root:root
      sudo: false             # default false
      # A sudo upload lands root:root unless `owner:` says otherwise. (Before
      # 2026-09-23 it kept the SSH user's ownership, because the staged file is
      # scp'd as that user and then `sudo mv`'d into place.) Without sudo the
      # file is the SSH user's by construction, and `owner:` is refused.
      stop_on_fail: true      # optional per-step override of global flag
      # Idempotency (tier 1): if ANY of these conditions says 'skip',
      # the step is marked `skipped`, not run, not counted as failure.
      when: "<remote shell expr>"     # skip unless expr exits 0
      creates: <remote path>          # skip if this path already exists
      removes: <remote path>          # skip if this path is already absent
      # Idempotency (tier 2): runs on the remote AFTER a successful step.
      # If exit 0, step is `changed`; exit != 0, `ok` (meaning no-op).
      # Without this, successful steps default to `changed` (Ansible's
      # default for shell/command modules).
      changed_when: "<remote shell expr>"
  verify:                     # optional; always runs after steps (even on fail)
    - name: Port responds
      shell: curl -sf http://localhost:8188/ >/dev/null

Per-step state machine:
  ok       — ran successfully, changed_when reported no change
  changed  — ran successfully (default for shell/upload without changed_when)
  skipped  — conditions said skip, OR halted after a prior failure, OR dry-run
  failed   — exit code != 0

Overall outcome / exit code:
  0 if OK (all ok) or CHANGED (some changed, none failed)
  1 if FAILED (any failed step, non-ignored)
  2 — usage / config error before execution began
"""

from __future__ import annotations

import argparse
import dataclasses
import datetime
import getpass
import json
import os
import re
import shlex
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import Optional

try:
    import yaml
except ImportError:
    sys.stderr.write(
        "elway: PyYAML is required. Install with:\n"
        "    sudo apt install python3-yaml\n"
        "or, inside a venv:\n"
        "    pip install pyyaml\n"
    )
    sys.exit(2)


# ─── Run logging ──────────────────────────────────────────────────────────
#
# Every run tees output to two files in addition to the terminal so the
# operator (or a sibling tool) can re-read what happened without scrolling
# back. Two locations:
#
#   /tmp/elway-last.log                — overwritten per run; the easy
#                                        "what just happened" target
#   ~/.cache/elway/runs/<ts>-<host>-<playbook>.log
#                                      — permanent timestamped record;
#                                        accumulates over time
#
# ANSI color codes are kept in the file (modern less/tail handle them);
# strip with `sed 's/\x1b\[[0-9;]*m//g'` if pasting somewhere that doesn't.

LAST_LOG_PATH = Path("/tmp/elway-last.log")
LOG_ROOT = Path.home() / ".cache" / "elway" / "runs"


class _Tee:
    """File-like that fans writes out to one or more streams."""

    def __init__(self, *streams) -> None:
        self._streams = [s for s in streams if s is not None]

    def write(self, data) -> int:
        for s in self._streams:
            try:
                s.write(data)
            except (BrokenPipeError, ValueError):
                pass
        return len(data)

    def flush(self) -> None:
        for s in self._streams:
            try:
                s.flush()
            except (BrokenPipeError, ValueError):
                pass

    def isatty(self) -> bool:
        # Preserve color decisions: if the *first* stream (the real
        # terminal) is a tty, the Tee reports as one. The on-disk files
        # get the same colored output, which is fine.
        return bool(self._streams) and getattr(self._streams[0], "isatty", lambda: False)()


def _open_run_logs(host: str, playbook: Optional[str], explicit: Optional[Path]) -> tuple[Optional[Path], list]:
    """Open log files for this run. Returns (permanent_log_path, [open_handles])."""
    handles = []
    perm_path = None

    safe_host = re.sub(r"[^a-zA-Z0-9._-]", "_", host)
    pb_label = "ad-hoc" if not playbook else re.sub(r"[^a-zA-Z0-9._-]", "_", Path(playbook).stem)

    if explicit is not None:
        # Caller forced a specific path — only open that one (no /tmp tee).
        try:
            explicit.parent.mkdir(parents=True, exist_ok=True)
            handles.append(open(explicit, "w", encoding="utf-8"))
            perm_path = explicit
        except OSError as exc:
            sys.stderr.write(f"[elway] could not open --log {explicit}: {exc}\n")
        return perm_path, handles

    # Default: tee both /tmp/elway-last.log and a timestamped permanent log.
    try:
        handles.append(open(LAST_LOG_PATH, "w", encoding="utf-8"))
    except OSError as exc:
        sys.stderr.write(f"[elway] could not open {LAST_LOG_PATH}: {exc}\n")

    try:
        LOG_ROOT.mkdir(parents=True, exist_ok=True)
        ts = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        perm_path = LOG_ROOT / f"{ts}-{safe_host}-{pb_label}.log"
        handles.append(open(perm_path, "w", encoding="utf-8"))
    except OSError as exc:
        sys.stderr.write(f"[elway] could not open permanent log under {LOG_ROOT}: {exc}\n")
        perm_path = None

    return perm_path, handles


# ─── Terminal output helpers ──────────────────────────────────────────────

IS_TTY = sys.stdout.isatty() and os.environ.get("NO_COLOR") != "1"


def c(code: str, s: str) -> str:
    return f"\033[{code}m{s}\033[0m" if IS_TTY else s


GREEN = lambda s: c("32", s)
RED = lambda s: c("31", s)
YELLOW = lambda s: c("33", s)
CYAN = lambda s: c("36", s)
DIM = lambda s: c("2", s)
BOLD = lambda s: c("1", s)


# ─── Data model ────────────────────────────────────────────────────────────


@dataclasses.dataclass
class UploadSpec:
    src: str
    dest: str
    mode: Optional[str] = None  # octal string like "0644"
    owner: Optional[str] = None  # "user[:group]"; sudo only; None = root:root


# user or user:group — names or numeric ids. Deliberately narrow: the value is
# spliced into a remote shell line, and anything outside this set is a typo or
# an injection, not an owner.
OWNER_RE = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_.-]*(:[A-Za-z0-9_][A-Za-z0-9_.-]*)?$")
DEFAULT_SUDO_OWNER = "root:root"


# Octal only. Like owner, mode is spliced into a remote (often root) shell
# line, so anything else is refused rather than quoted-and-hoped.
MODE_RE = re.compile(r"^[0-7]{3,4}$")


def check_mode(mode: Optional[str], where: str) -> None:
    if mode is not None and not MODE_RE.match(mode):
        raise SystemExit(f"elway: {where} mode {mode!r} is not an octal mode like \"0644\"")


def _yaml_str_field(u: dict, key: str, where: str) -> Optional[str]:
    """A field that must be a quoted YAML string. YAML 1.1 turns a bare 0644
    into the int 420 and a bare 1000:0 into 60000 (sexagesimal), and yes/no
    into booleans; str() of any of those is a DIFFERENT, valid-looking value.
    Refuse instead of guessing."""
    v = u.get(key)
    if v is None:
        return None
    if not isinstance(v, str):
        raise SystemExit(
            f"elway: {where} {key} must be a quoted string (got {type(v).__name__} "
            f"{v!r} — YAML reinterpreted it; write {key}: \"...\")"
        )
    return v


def check_owner(owner: Optional[str], sudo: bool, where: str) -> None:
    """Refuse an owner that is malformed, or set on a non-sudo upload."""
    if owner is None:
        return
    if not sudo:
        raise SystemExit(
            f"elway: {where} owner: needs sudo — without it the file is the "
            f"SSH user's and cannot be chowned to anyone else"
        )
    if not OWNER_RE.match(owner):
        raise SystemExit(f"elway: {where} owner {owner!r} is not user or user:group")


@dataclasses.dataclass
class Step:
    name: str
    shell: Optional[str] = None
    upload: Optional[UploadSpec] = None
    sudo: bool = False
    stop_on_fail: Optional[bool] = None  # None = inherit global
    # Tier 1 — conditional skip. Any of these evaluating to "skip" means the
    # step is not run; it counts as `skipped`, not `failed`.
    #   when:     remote shell expr; skip if exit != 0 (run when it returns 0)
    #   creates:  remote path; skip if it already exists
    #   removes:  remote path; skip if it does NOT exist
    when: Optional[str] = None
    creates: Optional[str] = None
    removes: Optional[str] = None
    # Tier 2 — change detection. Evaluated on the remote after a successful
    # run; exit 0 means "something changed", exit != 0 means "no change".
    # `changed_when: "false"` forces ok; `changed_when: "true"` forces changed.
    # Without this field, a successful step is `changed` by default.
    changed_when: Optional[str] = None

    def kind(self) -> str:
        return "shell" if self.shell is not None else "upload"


@dataclasses.dataclass
class Playbook:
    vars: dict
    steps: list[Step]
    verify: list[Step]


# Step result state machine:
#   ok       — ran successfully, changed_when said "no change"
#   changed  — ran successfully (default when changed_when not set)
#   failed   — exit code != 0 (and stop_on_fail / effective flag engaged)
#   skipped  — skipped due to when/creates/removes OR halted after prior fail
STATES = ("ok", "changed", "failed", "skipped")


@dataclasses.dataclass
class StepResult:
    step: Step
    phase: str  # "steps" or "verify"
    index: int
    total: int
    state: str  # one of STATES
    exit_code: int = 0
    duration_s: float = 0.0
    skip_reason: Optional[str] = None

    @property
    def succeeded(self) -> bool:
        # "succeeded" = didn't fail. Skipped and ok and changed all count.
        return self.state != "failed"


# ─── Template substitution ─────────────────────────────────────────────────

VAR_RE = re.compile(r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}")


def substitute(text: str, vars: dict, where: str) -> str:
    def sub(m: re.Match) -> str:
        key = m.group(1)
        if key not in vars:
            raise SystemExit(
                f"elway: undefined variable {{{{ {key} }}}} in {where} "
                f"(known: {sorted(vars.keys()) or 'none'})"
            )
        return str(vars[key])

    return VAR_RE.sub(sub, text)


def substitute_step(step: Step, vars: dict, where: str) -> Step:
    new = dataclasses.replace(step)
    new.name = substitute(step.name, vars, f"{where} name")
    if step.shell is not None:
        new.shell = substitute(step.shell, vars, f"{where} shell")
    if step.upload is not None:
        new.upload = UploadSpec(
            src=substitute(step.upload.src, vars, f"{where} upload.src"),
            dest=substitute(step.upload.dest, vars, f"{where} upload.dest"),
            mode=step.upload.mode,
            owner=(substitute(step.upload.owner, vars, f"{where} upload.owner")
                   if step.upload.owner is not None else None),
        )
        check_owner(new.upload.owner, step.sudo, f"{where} upload")
    for field in ("when", "creates", "removes", "changed_when"):
        val = getattr(step, field)
        if val is not None:
            setattr(new, field, substitute(val, vars, f"{where} {field}"))
    return new


# ─── Playbook loading ──────────────────────────────────────────────────────


def _as_str(value, where: str) -> str:
    # YAML parses `false` / `true` / `null` / numbers as typed scalars. For fields
    # that are meant to be shell text or paths, always present them as strings so
    # the templater (and eventually the shell) sees what the author intended.
    if value is None:
        raise SystemExit(f"elway: {where} is null; expected a string")
    return str(value)


def _load_step(raw: dict, where: str) -> Step:
    if not isinstance(raw, dict):
        raise SystemExit(f"elway: {where} must be a mapping, got {type(raw).__name__}")
    if "name" not in raw:
        raise SystemExit(f"elway: {where} missing required 'name'")
    has_shell = "shell" in raw
    has_upload = "upload" in raw
    if has_shell == has_upload:
        raise SystemExit(f"elway: {where} must have exactly one of 'shell' or 'upload'")
    upload = None
    shell = None
    if has_shell:
        shell = _as_str(raw["shell"], f"{where} shell")
    if has_upload:
        u = raw["upload"]
        if not isinstance(u, dict) or "src" not in u or "dest" not in u:
            raise SystemExit(f"elway: {where} upload must have src and dest")
        mode = _yaml_str_field(u, "mode", f"{where} upload")
        owner = _yaml_str_field(u, "owner", f"{where} upload")
        check_mode(mode, f"{where} upload")
        upload = UploadSpec(
            src=_as_str(u["src"], f"{where} upload.src"),
            dest=_as_str(u["dest"], f"{where} upload.dest"),
            mode=mode,
            owner=owner,
        )
        if upload.owner is not None and not bool(raw.get("sudo", False)):
            check_owner(upload.owner, False, f"{where} upload")
        if upload.owner is not None and not VAR_RE.search(upload.owner):
            check_owner(upload.owner, True, f"{where} upload")
    def opt_str(field: str) -> Optional[str]:
        v = raw.get(field)
        return None if v is None else _as_str(v, f"{where} {field}")

    return Step(
        name=_as_str(raw["name"], f"{where} name"),
        shell=shell,
        upload=upload,
        sudo=bool(raw.get("sudo", False)),
        stop_on_fail=raw.get("stop_on_fail"),
        when=opt_str("when"),
        creates=opt_str("creates"),
        removes=opt_str("removes"),
        changed_when=opt_str("changed_when"),
    )


def load_playbook(path: Path) -> Playbook:
    try:
        doc = yaml.safe_load(path.read_text())
    except yaml.YAMLError as e:
        raise SystemExit(f"elway: failed to parse {path}: {e}")
    if doc is None:
        raise SystemExit(f"elway: {path} is empty")
    if not isinstance(doc, dict):
        raise SystemExit(f"elway: {path} top-level must be a mapping with 'steps'")

    vars_ = doc.get("vars") or {}
    if not isinstance(vars_, dict):
        raise SystemExit(f"elway: 'vars' in {path} must be a mapping")

    steps_raw = doc.get("steps")
    if not steps_raw:
        raise SystemExit(f"elway: {path} has no 'steps'")
    steps = [_load_step(s, f"{path}:steps[{i}]") for i, s in enumerate(steps_raw)]

    verify_raw = doc.get("verify") or []
    verify = [_load_step(s, f"{path}:verify[{i}]") for i, s in enumerate(verify_raw)]

    return Playbook(vars=vars_, steps=steps, verify=verify)


# ─── SSH / SCP wrappers with ControlMaster ─────────────────────────────────


@dataclasses.dataclass
class SSHContext:
    host: str
    control_path: str
    verbose: bool
    sudo_password: Optional[str] = None
    sudo_probed: bool = False   # so we probe at most once per run, lazily

    def _base_ssh_opts(self) -> list[str]:
        return [
            "-o",
            "ControlMaster=auto",
            "-o",
            f"ControlPath={self.control_path}",
            "-o",
            "ControlPersist=60s",
            "-o",
            "BatchMode=no",
        ]

    def ssh_cmd(self, remote_cmd: str) -> list[str]:
        return ["ssh", *self._base_ssh_opts(), self.host, remote_cmd]

    def scp_cmd(self, src: str, dest: str) -> list[str]:
        return ["scp", *self._base_ssh_opts(), src, f"{self.host}:{dest}"]

    def close_master(self) -> None:
        # Best-effort: terminate the control socket so we don't leak processes.
        try:
            subprocess.run(
                ["ssh", "-O", "exit", "-o", f"ControlPath={self.control_path}", self.host],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
        except Exception:
            pass


def ensure_sudo(ctx: SSHContext) -> None:
    """Idempotent lazy sudo probe. Called before the first sudo step that will
    actually run — not at startup — so playbooks whose sudo steps all skip
    never prompt for a password.

    Sets ctx.sudo_password to None if NOPASSWD sudo works, or to the validated
    password the user entered. Marks ctx.sudo_probed=True so subsequent sudo
    steps reuse the already-acquired credential.
    """
    if ctx.sudo_probed:
        return
    ctx.sudo_probed = True
    # Is sudo configured NOPASSWD for this user?
    #
    # ⚠ This MUST be `sudo -n true`, never `sudo -n -v`. `-v` refreshes the
    # auth timestamp, and a NOPASSWD-only rule creates no timestamp to refresh,
    # so on **sudo >= 1.9.15** `-n -v` returns non-zero even though every real
    # command runs passwordless. Measured 2026-09-01:
    #
    #   pfi-gx10    sudo 1.9.15p5   sudo -n -v rc=1   sudo -n true rc=0
    #   nh3-docker  sudo 1.9.13p3   sudo -n -v rc=0   sudo -n true rc=0
    #   ana-docker  sudo 1.9.13p3   sudo -n -v rc=0   sudo -n true rc=0
    #   irv-ml1     sudo 1.9.13p3   sudo -n -v rc=0   sudo -n true rc=0
    #
    # The old probe cost a whole session on gx10: elway prompted for a password
    # on a host that had working NOPASSWD sudo the entire time, and in a
    # non-interactive run that prompt is an EOFError mid-playbook. Every host
    # will hit this as it moves past sudo 1.9.13. Assert the effective
    # behaviour (can I run a command?), not a proxy for it.
    p = subprocess.run(
        ctx.ssh_cmd("sudo -n true"),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    if p.returncode == 0:
        ctx.sudo_password = None
        return
    pw = getpass.getpass(f"[elway] sudo password for {ctx.host}: ")
    # Validate it works; don't silently carry a wrong password through the whole run.
    validation = subprocess.run(
        ctx.ssh_cmd("sudo -S -p '' -v"),
        input=pw + "\n",
        text=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
    )
    if validation.returncode != 0:
        raise SystemExit("elway: sudo password validation failed; aborting.")
    ctx.sudo_password = pw


# ─── Step execution ────────────────────────────────────────────────────────


def _stream_process(cmd: list[str], stdin_data: Optional[str], prefix: str) -> int:
    """Run `cmd`, streaming stdout+stderr line-by-line with `prefix`. Return exit code."""
    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE if stdin_data is not None else subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # interleave — user confirmed ok
        text=True,
        bufsize=1,
    )
    if stdin_data is not None:
        assert proc.stdin is not None
        try:
            proc.stdin.write(stdin_data)
            proc.stdin.close()
        except BrokenPipeError:
            pass
    assert proc.stdout is not None
    for line in proc.stdout:
        print(f"{prefix}{line.rstrip()}")
    return proc.wait()


def run_shell_step(ctx: SSHContext, step: Step, prefix: str) -> int:
    """Run a shell step on the remote host. Returns exit code."""
    cmd_inner = step.shell or ""
    if step.sudo:
        # Lazy: probe (and prompt if needed) right before the first sudo step
        # actually executes. Skipped sudo steps never trigger a prompt.
        ensure_sudo(ctx)
        # `sudo -S -p '' bash -c <script>` lets sudo consume its password
        # line from stdin, then runs the script. Password comes as one line.
        remote = "sudo -S -p '' bash -c " + shlex.quote(cmd_inner)
        stdin_data = (ctx.sudo_password or "") + "\n"
    else:
        remote = "bash -c " + shlex.quote(cmd_inner)
        stdin_data = None
    cmd = ctx.ssh_cmd(remote)
    if ctx.verbose:
        print(DIM(f"{prefix}$ {' '.join(shlex.quote(c) for c in cmd)}"))
    return _stream_process(cmd, stdin_data, prefix)


def run_upload_step(ctx: SSHContext, step: Step, prefix: str) -> int:
    """Run an upload step. Auto-mkdirs parent dir, preserves/applies mode, sudo-capable."""
    assert step.upload is not None
    up = step.upload
    src = Path(up.src)
    if not src.exists():
        print(f"{prefix}{RED('ERROR')} local file not found: {src}")
        return 2

    parent = os.path.dirname(up.dest) or "/"
    check_mode(up.mode, "upload")

    if not step.sudo:
        # Non-sudo path: mkdir parent, scp direct, chmod.
        mkdir = f"mkdir -p {shlex.quote(parent)}"
        if ctx.verbose:
            print(DIM(f"{prefix}$ ssh {ctx.host} {mkdir!r}"))
        rc = _stream_process(ctx.ssh_cmd(mkdir), None, prefix)
        if rc != 0:
            return rc
        if ctx.verbose:
            print(DIM(f"{prefix}$ scp {up.src} {ctx.host}:{up.dest}"))
        rc = _stream_process(ctx.scp_cmd(str(src), up.dest), None, prefix)
        if rc != 0:
            return rc
        if up.mode:
            if ctx.verbose:
                print(DIM(f"{prefix}$ ssh {ctx.host} chmod {up.mode} {up.dest}"))
            rc = _stream_process(
                ctx.ssh_cmd(f"chmod {shlex.quote(up.mode)} {shlex.quote(up.dest)}"), None, prefix
            )
        return rc

    # Sudo path: probe lazily, then scp to /tmp as user, then sudo mv + chown +
    # chmod server-side. Validate the owner BEFORE anything is copied, so a bad
    # value can never orphan a staged file in the remote /tmp.
    owner = up.owner or DEFAULT_SUDO_OWNER
    check_owner(owner, True, "upload")
    ensure_sudo(ctx)
    staging = f"/tmp/elway.{os.getpid()}.{int(time.time() * 1000)}.{src.name}"
    if ctx.verbose:
        print(DIM(f"{prefix}$ scp {up.src} {ctx.host}:{staging}"))
    rc = _stream_process(ctx.scp_cmd(str(src), staging), None, prefix)
    if rc != 0:
        return rc
    # Prepare the STAGED file completely, then publish it with one mv. A chown
    # that failed after mv would leave the live path owned by the SSH user —
    # the exact state this exists to prevent — so ownership and mode are set
    # before the file is visible at dest. mv is a same-fs rename (or, as root,
    # an ownership-preserving copy), so both survive it.
    #   * chown BEFORE chmod: chown clears setuid/setgid, so a requested 4755
    #     must be applied after it or it silently becomes 0755.
    #   * A directory dest is refused: `mv` would move INTO it (dest=/home/x
    #     would drop the file inside a home dir under a staging name).
    #   * The trap removes the staged file on every exit path; after a
    #     successful mv there is nothing left for it to remove.
    qstage, qdest = shlex.quote(staging), shlex.quote(up.dest)
    mode_clause = f" && chmod {shlex.quote(up.mode)} {qstage}" if up.mode else ""
    script = (
        f"trap {shlex.quote('rm -f -- ' + qstage)} EXIT; "
        f"{{ [ ! -d {qdest} ] || {{ echo 'elway: upload dest is a directory: '{qdest} >&2; exit 2; }}; }} "
        f"&& chown {shlex.quote(owner)} {qstage}"
        f"{mode_clause}"
        f" && mkdir -p {shlex.quote(parent)} "
        f"&& mv -f {qstage} {qdest}"
    )
    remote = "sudo -S -p '' bash -c " + shlex.quote(script)
    if ctx.verbose:
        print(DIM(f"{prefix}$ ssh {ctx.host} {remote}"))
    return _stream_process(
        ctx.ssh_cmd(remote),
        (ctx.sudo_password or "") + "\n",
        prefix,
    )


# ─── Skip + change evaluators (tier 1 + tier 2) ────────────────────────────


def _quiet_rc(ctx: SSHContext, remote_cmd: str, sudo: bool = False) -> int:
    """Run a shell expression on the remote under bash -c, discarding output.

    ⚠ `sudo` MUST track the step's own sudo setting. Conditions used to run
    unprivileged no matter what the step said, and that fails in the dangerous
    direction: a root-requiring `when:` (say `pvesh get ...`, which exits 255
    for a non-root user) returns non-zero, elway reports the step `skipped`,
    and a playbook that never performed its change reports overall OK. Measured
    2026-09-19 on esh-pve — three runs of an exclusion playbook reported
    success while changing nothing, and only the verify phase caught it.

    Sudo-less steps are unaffected: their conditions still evaluate as the
    login user, which is what they mean.
    """
    # Wrapping in bash -c gives `!`, `[[`, pipes, etc. consistent semantics
    # across hosts whose default login shell might be dash (Debian default) or
    # something else weird.
    if sudo:
        ensure_sudo(ctx)
        wrapped = "sudo -S -p '' bash -c " + shlex.quote(remote_cmd)
        stdin_data = (ctx.sudo_password or "") + "\n"
    else:
        wrapped = "bash -c " + shlex.quote(remote_cmd)
        stdin_data = None
    return subprocess.run(
        ctx.ssh_cmd(wrapped),
        input=stdin_data,
        text=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    ).returncode


def evaluate_skip(ctx: SSHContext, step: Step) -> Optional[str]:
    """If any pre-condition says skip, return a human-readable reason. Else None."""
    if step.when is not None:
        rc = _quiet_rc(ctx, step.when, step.sudo)
        if rc != 0:
            return f"when: expr exited {rc} (needed 0 to run)"
    if step.creates is not None:
        rc = _quiet_rc(ctx, f"test -e {shlex.quote(step.creates)}", step.sudo)
        if rc == 0:
            return f"creates: {step.creates} already exists"
    if step.removes is not None:
        rc = _quiet_rc(ctx, f"test -e {shlex.quote(step.removes)}", step.sudo)
        if rc != 0:
            return f"removes: {step.removes} is already absent"
    return None


def evaluate_changed(ctx: SSHContext, step: Step) -> bool:
    """Post-step: should this count as `changed`? Default True (Ansible shell default)."""
    if step.changed_when is None:
        return True
    rc = _quiet_rc(ctx, step.changed_when, step.sudo)
    # exit 0 = the change-detector expression "fired" = step counts as changed.
    return rc == 0


# ─── Driver / reporting ────────────────────────────────────────────────────


STATE_BADGE = {
    "ok":      lambda: GREEN("○ OK"),
    "changed": lambda: YELLOW("● CHANGED"),
    "failed":  lambda: RED("✗ FAILED"),
    "skipped": lambda: DIM("⏭ SKIPPED"),
}


def _print_status(state: str, exit_code: int, elapsed: float, skip_reason: Optional[str]) -> None:
    badge = STATE_BADGE[state]()
    suffix = ""
    if state == "failed":
        suffix = f" (rc={exit_code})"
    elif state == "skipped" and skip_reason:
        suffix = f" — {skip_reason}"
    tail = f"  {DIM(f'{elapsed:.2f}s')}" if state not in ("skipped",) else ""
    print(f"  {badge}{suffix}{tail}")


def run_phase(
    ctx: SSHContext,
    phase: str,
    steps: list[Step],
    vars: dict,
    global_stop_on_fail: bool,
    dry_run: bool,
) -> list[StepResult]:
    results: list[StepResult] = []
    halted = False
    for i, step in enumerate(steps, start=1):
        resolved = substitute_step(step, vars, where=f"{phase}[{i}]")
        header = f"[{ctx.host}] {phase} {i}/{len(steps)}  {BOLD(resolved.name)}"
        if resolved.sudo:
            header += f"  {YELLOW('(sudo)')}"
        print(header)

        if halted:
            _print_status("skipped", 0, 0.0, "earlier step failed")
            results.append(StepResult(resolved, phase, i, len(steps),
                                      state="skipped", skip_reason="earlier step failed"))
            continue

        if dry_run:
            if resolved.kind() == "shell":
                print(f"  {DIM('dry-run:')} shell: {resolved.shell}")
            else:
                u = resolved.upload
                mode = f" mode={u.mode}" if u.mode else ""
                owner = (f" owner={u.owner or DEFAULT_SUDO_OWNER}" if resolved.sudo else "")
                print(f"  {DIM('dry-run:')} upload: {u.src} → {u.dest}{mode}{owner}")
            if resolved.when or resolved.creates or resolved.removes:
                conds = []
                if resolved.when: conds.append(f"when={resolved.when!r}")
                if resolved.creates: conds.append(f"creates={resolved.creates!r}")
                if resolved.removes: conds.append(f"removes={resolved.removes!r}")
                print(f"  {DIM('dry-run:')} conditions: {', '.join(conds)}")
            results.append(StepResult(resolved, phase, i, len(steps), state="skipped",
                                      skip_reason="dry-run"))
            continue

        # Tier 1: conditional skip
        skip_reason = evaluate_skip(ctx, resolved)
        if skip_reason is not None:
            _print_status("skipped", 0, 0.0, skip_reason)
            results.append(StepResult(resolved, phase, i, len(steps),
                                      state="skipped", skip_reason=skip_reason))
            continue

        # Execute
        prefix = f"  {DIM(f'[{phase[0]}{i}]')} "
        t0 = time.monotonic()
        if resolved.kind() == "shell":
            rc = run_shell_step(ctx, resolved, prefix)
        else:
            rc = run_upload_step(ctx, resolved, prefix)
        elapsed = time.monotonic() - t0

        if rc != 0:
            state = "failed"
        else:
            # Tier 2: changed_when decides ok vs changed
            state = "changed" if evaluate_changed(ctx, resolved) else "ok"

        _print_status(state, rc, elapsed, None)
        results.append(StepResult(resolved, phase, i, len(steps),
                                  state=state, exit_code=rc, duration_s=elapsed))

        if state == "failed":
            step_stop = resolved.stop_on_fail
            effective_stop = global_stop_on_fail if step_stop is None else step_stop
            if effective_stop and phase == "steps":
                halted = True
                # Verify still runs in the caller, even after halt.
    return results


def _tally(results: list[StepResult]) -> dict[str, int]:
    return {s: sum(1 for r in results if r.state == s) for s in STATES}


def print_summary(step_results: list[StepResult], verify_results: list[StepResult]) -> int:
    print()
    print(BOLD("── summary ────────────────────────────────────────────"))
    for phase, res in [("steps", step_results), ("verify", verify_results)]:
        if not res:
            continue
        t = _tally(res)
        # Color the ok / changed counts only when there's something to report.
        ok_s     = GREEN(f"{t['ok']} ok")
        chg_s    = (YELLOW(f"{t['changed']} changed") if t["changed"] else DIM("0 changed"))
        fail_s   = (RED(f"{t['failed']} failed") if t["failed"] else DIM("0 failed"))
        skip_s   = DIM(f"{t['skipped']} skipped")
        print(f"  {phase:<8} {ok_s}, {chg_s}, {fail_s}, {skip_s}")

    any_fail = any(r.state == "failed" for r in step_results + verify_results)
    any_changed = any(r.state == "changed" for r in step_results + verify_results)
    if any_fail:
        overall, exit_code = RED("FAILED"), 1
    elif any_changed:
        overall, exit_code = YELLOW("CHANGED"), 0
    else:
        overall, exit_code = GREEN("OK"), 0
    print(f"  overall: {overall}")
    return exit_code


# ─── ops log ───────────────────────────────────────────────────────────────
#
# `infra-ops` and `infra-hermes` are two agents behind one OS identity, so a
# change made over ssh is otherwise fingerprint-less. elway is the sanctioned
# way to CHANGE things on a host, which makes it the right place to record
# them — automatically, because a log you have to remember to write is the
# same class of instrument as a health check that passes in both states.
# See scripts/ops-log and docs/pfi/ops-log.md.

OPS_LOG_BIN = Path(__file__).resolve().parent / "ops-log"


def record_to_ops_log(host: str, playbook: Optional[str], adhoc: Optional[str],
                      step_results: list, verify_results: list) -> None:
    """Best-effort. A failure here must never change elway's exit code or
    output — the audit trail is not allowed to become a new failure mode on
    the path that fixes things."""
    if not OPS_LOG_BIN.is_file():
        return
    tallies = []
    for phase, res in (("steps", step_results), ("verify", verify_results)):
        if not res:
            continue
        t = _tally(res)
        tallies.append(f"{phase} {t['ok']} ok/{t['changed']} changed/"
                       f"{t['failed']} failed/{t['skipped']} skipped")
    all_res = step_results + verify_results
    if any(r.state == "failed" for r in all_res):
        outcome = "failed"
    elif any(r.state == "changed" for r in all_res):
        outcome = "changed"
    else:
        outcome = "ok"
    target = Path(playbook).stem if playbook else "ad-hoc"
    detail = "; ".join(tallies) or "no steps ran"
    if adhoc:
        # Truncated: the ops log is an index of what happened, not a
        # transcript — the full command is in the elway run log.
        detail += f" | {adhoc[:160]}"
    try:
        subprocess.run(
            [str(OPS_LOG_BIN), "record", "--host", host, "--action", "elway",
             "--target", target, "--outcome", outcome, "--detail", detail, "-q"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10)
    except Exception:
        pass


# ─── CLI ───────────────────────────────────────────────────────────────────


def parse_var_args(raw_list: list[str]) -> dict:
    out = {}
    for raw in raw_list:
        if "=" not in raw:
            raise SystemExit(f"elway: --var must be KEY=VALUE, got '{raw}'")
        k, v = raw.split("=", 1)
        out[k.strip()] = v
    return out


def parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="elway",
        description="Run a shell command, file upload, or YAML playbook over SSH.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("host", help="ssh alias or user@host")
    mode = p.add_mutually_exclusive_group(required=True)
    mode.add_argument("--shell", "-s", help="ad-hoc shell command to run on the remote host")
    mode.add_argument(
        "--upload",
        "-u",
        help="ad-hoc upload as LOCAL:REMOTE (optionally :MODE — e.g. compose.yaml:/opt/x/compose.yaml:0644)",
    )
    mode.add_argument("--playbook", "-p", help="path to a YAML playbook")
    p.add_argument("--sudo", action="store_true", help="run ad-hoc --shell/--upload with sudo")
    p.add_argument(
        "--owner",
        default=None,
        help="ad-hoc --upload --sudo only: USER[:GROUP] for the file (default root:root)",
    )
    p.add_argument(
        "--var",
        action="append",
        default=[],
        help="KEY=VALUE (repeatable). Overrides playbook `vars:` defaults.",
    )
    p.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Do not halt on first failure (default: halt).",
    )
    p.add_argument("--dry-run", action="store_true", help="Print the plan without executing.")
    p.add_argument("-v", "--verbose", action="store_true", help="Show ssh/scp invocations.")
    p.add_argument(
        "--log",
        type=Path,
        default=None,
        help="Override log file path. Default: tee to /tmp/elway-last.log + ~/.cache/elway/runs/<ts>-<host>-<playbook>.log",
    )
    p.add_argument(
        "--no-log",
        action="store_true",
        help="Don't write run logs to disk; terminal-only.",
    )
    args = p.parse_args(argv)
    if args.owner is not None and not args.upload:
        # Ignoring it would let the caller believe an ownership was applied.
        p.error("--owner applies only to an ad-hoc --upload (playbooks use upload.owner)")
    return args


def build_adhoc_playbook(args: argparse.Namespace) -> Playbook:
    if args.shell:
        step = Step(name="ad-hoc shell", shell=args.shell, sudo=args.sudo)
    else:
        assert args.upload
        parts = args.upload.split(":")
        if len(parts) == 2:
            src, dest = parts
            mode = None
        elif len(parts) == 3:
            src, dest, mode = parts
        else:
            raise SystemExit(
                "elway: --upload expected LOCAL:REMOTE or LOCAL:REMOTE:MODE"
            )
        check_owner(args.owner, args.sudo, "--upload")
        check_mode(mode, "--upload")
        step = Step(
            name="ad-hoc upload",
            upload=UploadSpec(src=src, dest=dest, mode=mode, owner=args.owner),
            sudo=args.sudo,
        )
    return Playbook(vars={}, steps=[step], verify=[])


def preflight(pb: Playbook, merged_vars: dict) -> None:
    """Resolve every step before anything runs. substitute_step raises on an
    undefined variable or a templated owner that resolves to garbage; doing it
    per-step inside run_phase meant earlier steps had already changed the host
    by the time a later step's config error surfaced."""
    for phase, steps in (("steps", pb.steps), ("verify", pb.verify)):
        for i, step in enumerate(steps, 1):
            substitute_step(step, merged_vars, f"{phase} {i} ({step.name!r})")


def main(argv: Optional[list[str]] = None) -> int:
    args = parse_args(argv if argv is not None else sys.argv[1:])

    if args.playbook:
        pb = load_playbook(Path(args.playbook))
    else:
        pb = build_adhoc_playbook(args)

    # Merge vars: playbook defaults, overridden by CLI.
    merged_vars = dict(pb.vars)
    merged_vars.update(parse_var_args(args.var))
    preflight(pb, merged_vars)   # config errors surface before any remote action

    # Set up run logging — tee stdout to log file(s) so subsequent reads
    # don't have to scroll the terminal. `--no-log` disables.
    log_handles: list = []
    perm_log: Optional[Path] = None
    real_stdout = sys.stdout
    if not args.no_log:
        perm_log, log_handles = _open_run_logs(args.host, args.playbook, args.log)
        if log_handles:
            sys.stdout = _Tee(real_stdout, *log_handles)

    # ControlMaster socket in a temp dir scoped to this run.
    with tempfile.TemporaryDirectory(prefix="elway-") as tmp:
        control_path = os.path.join(tmp, "cm-%r@%h:%p")
        ctx = SSHContext(
            host=args.host,
            control_path=control_path,
            verbose=args.verbose,
        )
        # No upfront sudo probe — lazily triggered from run_shell_step /
        # run_upload_step the first time a sudo step actually runs (i.e.
        # after its when:/creates: skip conditions pass). Playbooks whose
        # sudo steps all skip never prompt.
        try:
            print(BOLD(f"── {args.host} ─ playbook: {args.playbook or 'ad-hoc'}  "
                      f"({len(pb.steps)} steps"
                      + (f", {len(pb.verify)} verify" if pb.verify else "")
                      + ")"))
            if perm_log is not None:
                print(DIM(f"   log: {perm_log}"))
            step_results = run_phase(
                ctx, "steps", pb.steps, merged_vars,
                global_stop_on_fail=not args.continue_on_error,
                dry_run=args.dry_run,
            )
            verify_results = []
            if pb.verify:
                # Verify ALWAYS runs, even if main steps failed — the point
                # of a verify phase is to characterize end state, not to be
                # skipped when things are most interesting.
                verify_results = run_phase(
                    ctx, "verify", pb.verify, merged_vars,
                    global_stop_on_fail=not args.continue_on_error,
                    dry_run=args.dry_run,
                )
            rc = print_summary(step_results, verify_results)
            if not args.dry_run:
                record_to_ops_log(args.host, args.playbook,
                                  args.shell or args.upload, step_results,
                                  verify_results)
            if perm_log is not None:
                print(DIM(f"log saved: {perm_log}"))
            return rc
        finally:
            ctx.close_master()
            sys.stdout = real_stdout
            for h in log_handles:
                try:
                    h.close()
                except Exception:
                    pass


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