#!/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
      sudo: false             # default false
      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"


@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,
        )
    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 = u.get("mode")
        upload = UploadSpec(
            src=_as_str(u["src"], f"{where} upload.src"),
            dest=_as_str(u["dest"], f"{where} upload.dest"),
            mode=str(mode) if mode is not None else None,
        )
    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? `sudo -n -v` exits 0 if so.
    p = subprocess.run(
        ctx.ssh_cmd("sudo -n -v"),
        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 "/"
    mode_clause = f" && chmod {up.mode} {shlex.quote(up.dest)}" if up.mode else ""

    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 {up.mode} {shlex.quote(up.dest)}"), None, prefix
            )
        return rc

    # Sudo path: probe lazily, then scp to /tmp as user, then sudo mv + chmod server-side.
    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
    script = (
        f"mkdir -p {shlex.quote(parent)} "
        f"&& mv -f {shlex.quote(staging)} {shlex.quote(up.dest)}"
        f"{mode_clause}"
    )
    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) -> int:
    """Run a shell expression on the remote under bash -c, discarding output."""
    # Wrapping in bash -c gives `!`, `[[`, pipes, etc. consistent semantics
    # across hosts whose default login shell might be dash (Debian default) or
    # something else weird.
    wrapped = "bash -c " + shlex.quote(remote_cmd)
    return subprocess.run(
        ctx.ssh_cmd(wrapped),
        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)
        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)}")
        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)}")
        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)
    # 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 ""
                print(f"  {DIM('dry-run:')} upload: {u.src} → {u.dest}{mode}")
            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


# ─── 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(
        "--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.",
    )
    return p.parse_args(argv)


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"
            )
        step = Step(
            name="ad-hoc upload",
            upload=UploadSpec(src=src, dest=dest, mode=mode),
            sudo=args.sudo,
        )
    return Playbook(vars={}, steps=[step], verify=[])


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))

    # 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 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())