#!/usr/bin/env python3
"""Tiny QMP client for the qemu/ device-sim tests.

    qmp.py <socket> screendump <out.ppm>
    qmp.py <socket> tap <x> <y>          # absolute 0..32767 (virtio-tablet)
    qmp.py <socket> quit
    qmp.py <socket> drive <script> <outdir> [--size N]

`drive` runs a whole interaction on ONE connection, so the VM boots once and
the scenario costs a couple of seconds per step instead of a full boot. Its
script is one command per line, '#' comments and blank lines ignored, and
coordinates are PANEL PIXELS (0..size-1, default 720) rather than the
0..32767 absolute axis the raw `tap` takes -- the point of the mode is to write
"tap the gear at 47,676" straight off a screenshot.

    shot NAME              screendump to <outdir>/NAME.ppm
    tap X Y                press, hold past several LVGL poll periods, release
    swipe X1 Y1 X2 Y2 [MS] drag, interpolated so LVGL sees real motion
    fling X1 Y1 X2 Y2      swipe fast enough to leave momentum behind
    sleep SECONDS          let animations settle or timers fire
    echo TEXT              progress marker in the driver's own output

With `--ctl SOCK` (the unix socket run.sh --ctl creates; ui-drive.sh passes
it) the script can also ask the UI itself, through the debug channel in
warden_debug.c. The same channel is what tools/warden-ctl reaches over SSH on
a real panel, so these verbs mean the same thing on the rig and on hardware:

    nav MENU[/TAB]         jump to a page; fails if the page is unknown
    scroll X Y DY          scroll the scrollable under pixel X,Y by DY
                           pixels (positive reveals content below), no
                           animation: the deterministic stand-in for a
                           swipe that only brings a control into view
    home                   close every popup and go to Dashboard/Dashboard,
                           the state a fresh boot starts in
    wake                   wake the UI as if touched, minus the swallow
                           (sleep.c dims and then blanks an idle panel and
                           eats the touch that ends either); every compiled
                           flow script opens with it
    page | stats | hit X Y print the reply, judge nothing
    rs485 silence|restore ADDR
                           take a simulated bus device off the bus (holds its
                           address, answers nothing) or put it back; needs
                           --rs485-control (ui-drive.sh --rs485-devices)
    ctl WORDS...           raw passthrough for anything the channel grows
    assert_page MENU/TAB   the active page is exactly this
    assert_hit X Y CLASS [box=X1,Y1,WxH] [TEXT...]
                           a real tap at X,Y lands on an object of CLASS,
                           with that exact box when one is given, whose
                           caption contains TEXT: the pre-tap self-check that
                           reports a MOVED button as such, not as broken
                           behaviour
    wait_hit X Y CLASS TIMEOUT_S [box=X1,Y1,WxH] [TEXT...]
                           assert_hit, polled every 0.5s (poll_until) until it
                           holds or TIMEOUT_S elapses -- for a page still
                           building when the check runs

Two more channels reach behavioural state instead of the widget tree.
`--refs FILE` (default <outdir>/refs.json) is the reference store the pixel
ones read and write; imgtools.py in this directory does the actual image
math, this file only drives it.

    assert_json PATH OP VALUE
    wait_json PATH OP VALUE TIMEOUT_S
                           PATH is dotted with optional [N] indexes (e.g.
                           network.scan.ranges[0]) into the guest's
                           /tmp/warden-web-status.json, fetched fresh on
                           every check with `@cat` over the SAME control
                           channel as nav/hit (see the bridge added to
                           rootfs/sbin/init). OP is one of eq, ne, contains,
                           len_eq, len_ge, gt, lt. `wait_json` polls every
                           0.5s until OP holds or TIMEOUT_S elapses and
                           `assert_json` reads once. A PATH the document
                           doesn't have is a `fail`, not a crash -- that
                           document is a live snapshot and the key might
                           simply not be built yet.
    assert_stat FIELD OP VALUE
                           FIELD is cpu, fps, render or idle (0 awake,
                           1 dimmed, 2 asleep), read off a fresh `stats`
                           reply (warden_debug.c). Same OP vocabulary as
                           assert_json.
    capture_region NAME X Y W H TOLERANCE
                           screendump now, crop to X,Y,WxH, and (over)write
                           NAME in the refs file with both a phash and a
                           structural hash plus TOLERANCE (exact, loose or
                           structural). Run this once, by hand, to author a
                           reference -- the same way targets are authored,
                           see flows/README.md.
    assert_region NAME [TOLERANCE]
                           fresh screendump, crop to NAME's stored box,
                           compare against its stored hash at its stored
                           tolerance. TOLERANCE, when given, must be the one
                           NAME was captured under (flowc.py always emits it).
    wait_region NAME TOLERANCE TIMEOUT_S
                           assert_region, polled every 0.5s (poll_until) until
                           it holds or TIMEOUT_S elapses
    assert_ocr NAME REGEX  same crop, run through `tesseract`, REGEX searched
                           against the extracted text (spaces in REGEX are
                           fine, same as assert_hit's TEXT).

assert_region/assert_ocr against a NAME with no captured reference, and
assert_ocr with no `tesseract` on PATH, are FATAL for that one step: a
silently-skipped check reads as coverage that was never actually there. So
are a TOLERANCE that differs from the captured one and a stored box that does
not fit the screendump: both are the reference's fault, not the UI's, and one
bad reference must not take the rest of the script with it.

Every step is recorded to <outdir>/results.jsonl as ok / fail / fatal. An
assertion mismatch is a `fail` and the script continues, so one run reports
every broken expectation. With `--console LOG` the console is checked after
EVERY step for the stage-2 init's "warden-ui EXITED" line, so a crash is a
`fatal` pinned to the step that caused it rather than a stale picture found at
the end. The exit status is non-zero if anything failed.

A step that names an unknown command is a FATAL error rather than a skip: a
silently-ignored line in a scenario is a test that proves nothing.
"""
import json
import math
import os
import re
import shutil
import socket
import subprocess
import sys
import time

# imgtools.py lives beside this file and does the actual pixel math (phash,
# structural hash, crop, compare) for the region/ocr verbs -- and imports
# Pillow at its OWN module scope to do it. Importing it here at qmp.py's
# module scope would drag that cost onto every subprocess invocation of this
# file, including screendump/tap/quit, which never touch a pixel:
# ui-drive.sh's own boot-wait loop calls screendump specifically to avoid a
# Pillow dependency (see its comment beside the plain byte-loop colour
# count), so the other half of that same loop must not quietly re-add one.
# Deferred to need_imgtools() at first actual use instead; a missing or
# broken imgtools still must not take down the verbs that don't need it.
imgtools = None
_imgtools_import_error = None


def _load_imgtools():
    global imgtools, _imgtools_import_error
    if imgtools is None and _imgtools_import_error is None:
        try:
            import imgtools  # noqa: F811 -- binds the module-level name above
        except ImportError as e:
            _imgtools_import_error = e
    return imgtools

AXIS_MAX = 32767

# Bounds every blocking read on the QMP socket -- the greeting banner, the
# qmp_capabilities handshake, and every screendump/tap/quit round trip --
# the same way Ctl.__init__ already bounds the control channel (see its
# `timeout` default below). A socket with no timeout blocks forever on a
# wedged VM (a TCG stall or a kernel panic loop), and ui-drive.sh's own
# cleanup() calls `quit` on this socket before it ever reaches
# `reap "$QEMU_PID"` -- the bounded kill that is supposed to guarantee a
# wedged qemu-system-arm cannot outlive the script.
QMP_TIMEOUT_S = 20.0

# Where webstatus.c (ui-src/src/warden/webstatus.c) publishes its atomic
# snapshot inside the guest. assert_json/wait_json `@cat` this path over the
# control channel; see the bridge in rootfs/sbin/init.
STATUS_JSON_PATH = "/tmp/warden-web-status.json"


def rpc(sock, sock_file, obj):
    sock.sendall((json.dumps(obj) + "\n").encode())
    while True:
        line = sock_file.readline()
        if not line:
            raise RuntimeError("QMP connection closed")
        msg = json.loads(line)
        if "return" in msg:
            return msg["return"]
        if "error" in msg:
            raise RuntimeError(f"QMP error: {msg['error']}")
        # asynchronous events are interleaved; skip them


def send_events(s, f, events):
    rpc(s, f, {"execute": "input-send-event", "arguments": {"events": events}})


def abs_ev(axis, value):
    return {"type": "abs", "data": {"axis": axis, "value": value}}


def btn_ev(down):
    return {"type": "btn", "data": {"down": down, "button": "left"}}


def to_axis(px, size):
    """Panel pixels -> the tablet's absolute axis, clamped to the panel.

    LVGL's evdev driver maps the axis back with integer truncation
    (lv_evdev.c _evdev_calibrate: px = axis * (width - 1) / AXIS_MAX), so the
    axis value must be the SMALLEST one that truncates to px, i.e. rounded up.
    Truncating here as well composed two floors and landed one pixel short for
    many values (130 -> 5924 -> 129), which is how a tap could miss the control
    that `hit` at the same pixel had just confirmed (flare-edge #148 triage).
    tools/touch-inject writes the pixel itself on a panel, so hardware never
    had this seam."""
    px = max(0, min(size - 1, int(px)))
    return min(AXIS_MAX, math.ceil(px * AXIS_MAX / (size - 1)))


def do_tap(s, f, ax, ay, hold=0.2):
    send_events(s, f, [abs_ev("x", ax), abs_ev("y", ay), btn_ev(True)])
    # Hold the press across several LVGL indev poll periods (33 ms each): an
    # instantaneous press+release lands inside one poll and no click is ever
    # registered.
    time.sleep(hold)
    send_events(s, f, [btn_ev(False)])


def do_swipe(s, f, x1, y1, x2, y2, size, ms=400, steps=None):
    """Drag with interpolated motion.

    LVGL decides a gesture from the movement BETWEEN indev polls, so a press at
    the start and a release at the end (with nothing in between) reads as a
    click on whatever was under the finger, not a scroll. The interpolation
    below is the whole reason a scroll can be tested at all; the step count is
    derived from the duration so the poll period always sees a few pixels of
    travel.
    """
    ax1, ay1 = to_axis(x1, size), to_axis(y1, size)
    ax2, ay2 = to_axis(x2, size), to_axis(y2, size)
    if steps is None:
        steps = max(6, int(ms / 25))
    send_events(s, f, [abs_ev("x", ax1), abs_ev("y", ay1), btn_ev(True)])
    time.sleep(0.05)
    for i in range(1, steps + 1):
        t = i / steps
        send_events(s, f, [
            abs_ev("x", int(ax1 + (ax2 - ax1) * t)),
            abs_ev("y", int(ay1 + (ay2 - ay1) * t)),
        ])
        time.sleep(ms / 1000.0 / steps)
    send_events(s, f, [btn_ev(False)])


class Ctl:
    """The UI's debug channel, reached through run.sh --ctl.

    Inside the VM that channel is a FIFO polled by warden-ui (warden_debug.c);
    rootfs/sbin/init bridges it to a pci-serial port that run.sh exposes as a
    unix socket. One command line in, the reply out, then a sentinel line. On
    real hardware flare-edge tools/warden-ctl reaches the identical FIFO over
    SSH, so the vocabulary here is the vocabulary there: a flow script that
    runs against this rig runs against a panel.

    A missing or silent channel is an infrastructure fault and raises; it is
    never reported as an assertion failure, because a rig that cannot ask is
    not evidence about the UI either way.
    """

    # Must match, byte for byte, the `echo "<<END>>" >&3` in
    # rootfs/sbin/init's ctl bridge (the FIFO-to-socket relay this class
    # talks to). Nothing enforces the two staying in sync -- a sentinel
    # changed on one side and not the other means send() below blocks until
    # its own 15s timeout on every ctl-dependent step, with nothing at that
    # point pointing back at this mismatch as the cause.
    SENTINEL = "<<END>>"

    # A real reply (a hit's box, the status JSON) is at most a few KB. The
    # 15s socket timeout below bounds each individual recv(), not the total
    # bytes accepted, so a peer that keeps streaming data fast enough to beat
    # that per-call timeout, but never emits a newline or SENTINEL, would
    # otherwise grow self.buf without limit and exhaust host memory before
    # anything fails. Fail closed instead, the same way a channel that goes
    # silent or closes outright already does.
    MAX_BUF = 256 * 1024

    def __init__(self, path, timeout=15.0):
        self.sock = socket.socket(socket.AF_UNIX)
        self.sock.settimeout(timeout)
        self.sock.connect(path)
        self.buf = b""

    def send(self, cmd):
        self.sock.sendall((cmd + "\n").encode())
        lines = []
        while True:
            nl = self.buf.find(b"\n")
            if nl < 0:
                chunk = self.sock.recv(4096)
                if not chunk:
                    raise RuntimeError("control channel closed")
                self.buf += chunk
                if len(self.buf) > self.MAX_BUF:
                    raise RuntimeError(
                        f"control channel reply too large (over {self.MAX_BUF} "
                        f"bytes with no newline or {self.SENTINEL!r} seen)")
                continue
            line = self.buf[:nl].decode("utf-8", "replace").rstrip("\r")
            self.buf = self.buf[nl + 1:]
            if line == self.SENTINEL:
                break
            # The tty is in cooked mode, so the command comes back echoed.
            if line.strip() == cmd.strip():
                continue
            lines.append(line)
        return "\n".join(lines).strip()


HIT_RE = re.compile(r'^hit \d+,\d+: (\S+)(?: text="(.*)" box=(-?\d+),(-?\d+),(\d+)x(\d+))?$')


def judge_hit(ctl, x, y, cls, rest):
    """One `hit` round trip judged the way assert_hit documents: -> (ok, detail).
    REST is the optional box=X1,Y1,WxH followed by the caption words."""
    rest = list(rest)
    want_box = None
    if rest and rest[0].startswith("box="):
        want_box = rest[0][4:]
        rest = rest[1:]
    want_text = " ".join(rest)
    got = parse_hit(ctl.send(f"hit {x} {y}"))
    if got is None:
        return False, "nothing clickable there"
    if got[0] != cls:
        return False, f"{got[0]} text={got[1]!r} box={got[2]}"
    if want_box and got[2] != want_box:
        return False, f"{cls} moved: box={got[2]}"
    if want_text and want_text not in got[1]:
        return False, f"{cls} text={got[1]!r}"
    return True, ""


def poll_until(check, timeout_s, period=0.5):
    """Run CHECK (-> (ok, detail)) until it holds or TIMEOUT_S passes;
    -> (ok, detail, seconds waited, fatal). The wait_* verbs share this so
    they all mean the same thing by a timeout.

    A CHECK that raises RuntimeError or OSError -- the control channel
    dying mid-poll, e.g. a socket.timeout past Ctl's own 15s per-call
    timeout -- stops polling right there instead of retrying a channel
    that is probably gone, and comes back with FATAL true so the caller
    records a `fatal` row (the same idea as ConsoleWatch.check() turning a
    console crash into a labeled row) rather than the exception escaping
    drive() and silently dropping every step after it."""
    start = time.monotonic()
    deadline = start + timeout_s
    while True:
        try:
            ok, detail = check()
        except (RuntimeError, OSError) as e:
            return False, str(e), time.monotonic() - start, True
        if ok or time.monotonic() >= deadline:
            return ok, detail, time.monotonic() - start, False
        time.sleep(period)


def parse_hit(reply):
    """-> (cls, text, box) with box as 'x1,y1,WxH', or None for 'nothing';
    raises on an unparseable reply."""
    m = HIT_RE.match(reply)
    if not m:
        raise RuntimeError(f"unparseable hit reply: {reply!r}")
    if m.group(1) == "nothing":
        return None
    box = f"{m.group(3)},{m.group(4)},{m.group(5)}x{m.group(6)}" if m.group(3) else ""
    return m.group(1), m.group(2) or "", box


class ConsoleWatch:
    """Incremental check for the stage-2 init's console crash marker.

    The framebuffer just keeps the last frame on a crash; only the console
    announces it, and drive() checks after EVERY step so a crash is pinned
    to the step that caused it, not discovered at the end. But the console
    log is not a static boot log -- rootfs/sbin/init pipes warden-ui's own
    stdout through `tee` into the same stream the host captures, so it grows
    for the life of the run. A fresh open()+read() of the WHOLE file on
    every one of those per-step checks costs O(steps * final_size) for
    nothing: a marker absent from bytes already scanned cannot retroactively
    appear there. This instead keeps the byte offset already scanned and
    reads only what was appended since the last check, with a short tail
    kept across calls so a marker split across two reads is still caught.
    One instance per drive() run."""

    MARKER = b"warden-ui EXITED"

    def __init__(self, console_path):
        self.path = console_path
        self.offset = 0
        self.tail = b""
        self.found = None

    def check(self):
        """-> the matched line (decoded), once and cached forever after
        (nothing later needs another read); None while nothing has matched
        yet, or PATH is unset or unreadable."""
        if not self.path or self.found is not None:
            return self.found
        try:
            with open(self.path, "rb") as fh:
                fh.seek(self.offset)
                chunk = fh.read()
                self.offset = fh.tell()
        except OSError:
            return None
        data = self.tail + chunk
        i = data.find(self.MARKER)
        if i >= 0:
            self.found = data[i:i + 80].split(b"\n", 1)[0].decode("utf-8", "replace")
            return self.found
        # Keep enough of the tail that a marker whose first byte landed in
        # this read, but the rest in the next one, is still caught.
        self.tail = data[-(len(self.MARKER) - 1):]
        return None


PATH_TOKEN_RE = re.compile(r'^([^\[\]]*)((?:\[\d+\])*)$')
IDX_RE = re.compile(r'\[(\d+)\]')


def resolve_path(doc, path):
    """Walk a dotted PATH (keys and [N] indexes, e.g. 'a.b[0].c') over an
    already-parsed JSON document. -> (value, None) or (None, 'no such path').

    Pure and I/O-free on purpose (contract: testable on its own) -- fetching
    the document is a separate step (fetch_status_json) so this function can
    be unit-tested against a plain dict/list literal with nothing running.
    """
    cur = doc
    for part in path.split("."):
        m = PATH_TOKEN_RE.match(part)
        if not m:
            return None, "no such path"
        key, idxs = m.group(1), m.group(2)
        if key:
            if not isinstance(cur, dict) or key not in cur:
                return None, "no such path"
            cur = cur[key]
        for idx_s in IDX_RE.findall(idxs):
            idx = int(idx_s)
            if not isinstance(cur, list) or idx >= len(cur) or idx < 0:
                return None, "no such path"
            cur = cur[idx]
    return cur, None


def resolve_token(raw):
    """A verb's literal OP argument -> a comparable Python value. A JSON
    literal (a number, true/false/null, or a quoted string) parses as
    itself; anything else -- most values in practice, e.g. `connected` --
    compares as the raw string. This is what lets both
    `assert_json state.net eq connected` and `assert_json state.count eq 3`
    work from one unquoted token, with no escaping convention of its own."""
    try:
        return json.loads(raw)
    except ValueError:
        return raw


def _len_of(v):
    try:
        return len(v)
    except TypeError:
        raise TypeError(f"{v!r} has no length") from None


def apply_op(op, actual, raw_value):
    """The comparison vocabulary shared by assert_json/wait_json/assert_stat:
    eq, ne, contains, len_eq, len_ge, gt, lt. Raises on a combination that
    cannot be judged (contains on a number, gt across incompatible types)
    so the caller turns that into a `fail` detail instead of a driver crash --
    a spec that asks a nonsensical question should be reported, not silently
    True or False."""
    value = resolve_token(raw_value)
    if op == "eq":
        return actual == value
    if op == "ne":
        return actual != value
    if op == "contains":
        if isinstance(actual, str):
            return str(value) in actual
        if isinstance(actual, (list, tuple, dict)):
            return value in actual
        raise TypeError(f"contains: {actual!r} is not a string or list")
    if op == "len_eq":
        return _len_of(actual) == int(value)
    if op == "len_ge":
        return _len_of(actual) >= int(value)
    if op == "gt":
        return actual > value
    if op == "lt":
        return actual < value
    raise ValueError(f"unknown op {op!r}")


def fetch_status_json(ctl):
    """`@cat` the guest's webstatus.c snapshot over the control channel and
    parse it. -> (doc, None) or (None, detail) for a torn read or a
    not-yet-written file -- an ordinary retry/fail for the caller, not a
    crash of the whole driver. ctl.send() itself is deliberately NOT guarded
    here: a dead control channel (RuntimeError on EOF, OSError/socket.timeout
    on a stalled read) is left to propagate, so wait_json's poll_until (and
    drive()'s own top-level guard for assert_json) can tell a dead channel
    apart from a value the document simply doesn't have yet, and report the
    former as a fast fatal rather than retrying it out to TIMEOUT_S."""
    reply = ctl.send(f"@cat {STATUS_JSON_PATH}")
    if reply.startswith("bridge: no such file"):
        return None, reply
    try:
        return json.loads(reply), None
    except ValueError as e:
        return None, f"status json unparsable: {e}"


def eval_json(ctl, path, op, value):
    """One fetch + path-resolve + op-apply round for assert_json/wait_json.
    -> (ok, detail); detail is empty on success, otherwise the reason."""
    doc, err = fetch_status_json(ctl)
    if err is not None:
        return False, err
    actual, perr = resolve_path(doc, path)
    if perr is not None:
        return False, perr
    try:
        ok = apply_op(op, actual, value)
    except (TypeError, ValueError) as e:
        return False, f"{op} {value!r} vs {actual!r}: {e}"
    return ok, "" if ok else f"{path} is {actual!r}"


STATS_FIELD_RE = {
    "cpu": re.compile(r'^cpu:\s*(-?\d+(?:\.\d+)?)%?\s*$'),
    "fps": re.compile(r'^fps:\s*(-?\d+(?:\.\d+)?)\s*$'),
    "render": re.compile(r'^render:\s*(-?\d+(?:\.\d+)?)\s*ms/frame\s*$'),
    "idle": re.compile(r'^idle:\s*(\d+)\s*$'),
}


def parse_stats(reply):
    """The `stats` reply (warden_debug.c: 'page: X\\ncpu: N%\\nfps: N\\n
    render: A.BB ms/frame\\nrga: N%\\nidle: N') -> {'cpu'|'fps'|'render'|'idle': float} for
    whichever lines are present. A field the reply lacks is simply absent
    from the result -- the caller reports that as 'no such field', the same
    shape as resolve_path's 'no such path' for the json channel."""
    out = {}
    for ln in reply.splitlines():
        ln = ln.strip()
        for field, rx in STATS_FIELD_RE.items():
            m = rx.match(ln)
            if m:
                out[field] = float(m.group(1))
    return out


def load_refs(path):
    """The region/ocr reference store (see qmp.py's own docstring for the
    schema). Missing file -> {}, same as an empty COVERAGE.yaml: capture_region
    is how one gets created, not a prerequisite to have one already."""
    if not path or not os.path.exists(path):
        return {}
    with open(path) as fh:
        return json.load(fh)


def save_refs(path, refs):
    # Written after every capture_region, not batched at the end: a script
    # that dies three steps later must not lose a reference that was already
    # good.
    with open(path, "w") as fh:
        json.dump(refs, fh, indent=2, sort_keys=True)
        fh.write("\n")


def write_png(img, path):
    """imgtools' (w, h, rgb_bytes) tuple -> a PNG file, because tesseract has
    no raw-RGB input mode. Only assert_ocr needs this; assert_region and
    capture_region work on imgtools' own tuples end to end. Pillow is
    imported here, not at module scope, for the same reason imgtools.py's
    own import is deferred above: a bare screendump/tap/quit subprocess must
    never pay for it. By the time assert_ocr reaches this call,
    need_imgtools() has already required imgtools -- which itself imports
    Pillow -- so in practice this import is a cache hit, not a fresh cost."""
    try:
        from PIL import Image
    except ImportError as e:
        raise RuntimeError(f"Pillow (PIL) is not installed: {e}") from e
    w, h, data = img
    Image.frombytes("RGB", (w, h), data).save(path)


def safe_out_path(outdir, name, prefix="", suffix=""):
    """<outdir>/PREFIXnameSUFFIX, or (None, detail) for a NAME that could
    write outside outdir. flowc.py does not restrict shot/region/ocr names
    beyond forbidding whitespace (see its own `shot` and `one_token`
    comments), so a hand-edited or buggy *.drive.txt line -- or a refs.json
    entry with a name capture_region itself would never have written -- must
    not be allowed to place NAME ahead of outdir in the joined path. An
    absolute NAME (e.g. `shot /etc/cron.d/x`) makes os.path.join() discard
    outdir entirely and return NAME verbatim; a NAME with any other '/' can
    still walk out of outdir with enough '..' segments. Mirrors
    flow-run-hw.sh's shot_out_path() for the identical class of bug on the
    hardware runner."""
    if not name or "/" in name:
        return None, f"{name!r} must not contain '/' (would escape outdir)"
    return os.path.join(os.path.abspath(outdir), f"{prefix}{name}{suffix}"), None


def rs485_send(control_path, line):
    """One command to mbsim.py's control socket (flare-edge tools/modbus-sim,
    --control): -> its one-line reply, `ok ...` or `error ...`."""
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as c:
        c.settimeout(5)
        c.connect(control_path)
        c.sendall((line + "\n").encode())
        return c.makefile("r").readline().strip()


class Ctx:
    """State one drive() run threads through every verb handler below,
    replacing the closures drive() used to build fresh on each call (record
    stays in drive() itself -- it owns the results.jsonl handle -- but
    need_ctl, need_imgtools and fresh_region move here since verb handlers,
    not drive(), are what call them now). One instance per run."""

    def __init__(self, s, f, size, ctl, refs, refs_path, rs485_control, outdir, script_path):
        self.s = s
        self.f = f
        self.size = size
        self.ctl = ctl
        self.refs = refs
        self.refs_path = refs_path
        self.rs485_control = rs485_control
        self.outdir = outdir
        self.script_path = script_path

    def need_ctl(self, lineno, cmd):
        if self.ctl is None:
            sys.exit(f"FATAL: {self.script_path}:{lineno}: '{cmd}' needs the control "
                     f"channel; run with --ctl (ui-drive.sh passes it)")

    def need_imgtools(self, lineno, cmd):
        # A whole-script infrastructure gap (imgtools.py absent), same class
        # as a missing --ctl: no region/ocr verb can do anything without it,
        # so this exits the process rather than recording a per-step fail.
        # This is also the first point that actually needs imgtools loaded
        # (see _load_imgtools above), so it is where the deferred import
        # happens, not module load.
        if _load_imgtools() is None:
            sys.exit(f"FATAL: {self.script_path}:{lineno}: '{cmd}' needs imgtools.py "
                     f"next to qmp.py (vendor/warden-sdk/qemu/tests/imgtools.py)")

    def fresh_region(self, name, tag):
        """screendump NOW (never reuse an earlier `shot`) and crop to NAME's
        stored box. Fails closed on an uncaptured NAME: assert_region and
        assert_ocr both need this, capture_region does not."""
        ref = self.refs.get(name)
        if ref is None:
            return None, None, f"no reference for {name!r} (run capture_region first)"
        path, err = safe_out_path(self.outdir, name, prefix=f"{tag}_", suffix=".ppm")
        if err:
            return None, None, err
        rpc(self.s, self.f, {"execute": "screendump", "arguments": {"filename": path}})
        try:
            img = imgtools.load_ppm(path)
            cropped = imgtools.crop(img, ref["x"], ref["y"], ref["w"], ref["h"])
        except (ValueError, OSError, KeyError, TypeError) as e:
            # A box outside this screendump (captured at another --size, or
            # mistyped in the store) is a per-step fatal, never a traceback
            # that ends the run: every later step still gets judged.
            return None, None, f"reference {name!r} unusable: {e}"
        return ref, cropped, None


# One function per script verb, VERBS-dispatched below instead of a single
# growing if/elif chain: each takes the run's shared Ctx plus the parsed
# line and returns either None (no results.jsonl row -- only echo does
# this) or (status, detail) for drive() to record. Keeping one function per
# verb, the way flowc.py tables its own CHANNELS/ACTIONS vocabulary, is
# what lets a verb be read, tested or reused on its own instead of only as
# one arm of drive().

def verb_shot(ctx, lineno, cmd, args, line):
    name = args[0]
    path, err = safe_out_path(ctx.outdir, name, suffix=".ppm")
    if err:
        return "fail", f"shot name {err}"
    rpc(ctx.s, ctx.f, {"execute": "screendump", "arguments": {"filename": path}})
    return "ok", ""


def verb_tap(ctx, lineno, cmd, args, line):
    x, y = int(args[0]), int(args[1])
    do_tap(ctx.s, ctx.f, to_axis(x, ctx.size), to_axis(y, ctx.size))
    return "ok", ""


def verb_swipe(ctx, lineno, cmd, args, line):
    ms = int(args[4]) if len(args) > 4 else 400
    do_swipe(ctx.s, ctx.f, int(args[0]), int(args[1]), int(args[2]), int(args[3]), ctx.size, ms)
    return "ok", ""


def verb_fling(ctx, lineno, cmd, args, line):
    do_swipe(ctx.s, ctx.f, int(args[0]), int(args[1]), int(args[2]), int(args[3]),
             ctx.size, ms=120)
    return "ok", ""


def verb_sleep(ctx, lineno, cmd, args, line):
    time.sleep(float(args[0]))
    return "ok", ""


def verb_echo(ctx, lineno, cmd, args, line):
    print(f"  [{lineno}] {' '.join(args)}", flush=True)
    return None


def verb_nav(ctx, lineno, cmd, args, line):
    ctx.need_ctl(lineno, cmd)
    reply = ctx.ctl.send("nav " + " ".join(args))
    return ("ok" if reply.endswith(": ok") else "fail"), reply


def verb_wake(ctx, lineno, cmd, args, line):
    # Judged, unlike the query verbs: a UI that does not answer the wake is
    # one whose next tap may be swallowed, and that must not read as a
    # passing step.
    ctx.need_ctl(lineno, cmd)
    reply = ctx.ctl.send("wake")
    return ("ok" if reply == "wake: ok" else "fail"), reply


def verb_scroll_home(ctx, lineno, cmd, args, line):
    # scroll X Y DY: the scrollable under the pixel moves DY pixels with no
    # animation (warden_debug.c), the deterministic stand-in for a swipe
    # whose only job was to bring a control into view.
    # home: close every popup and return to Dashboard/Dashboard, the state a
    # fresh boot starts in, so a live panel can run one flow after another.
    # Both are judged by the channel's own ack.
    ctx.need_ctl(lineno, cmd)
    reply = ctx.ctl.send(line)
    return ("ok" if reply.startswith(f"{cmd}: ok") else "fail"), reply


def verb_query(ctx, lineno, cmd, args, line):
    # page | hit | stats | ctl: print the reply, never judge it. `ctl` is a
    # raw passthrough for anything the channel grows later.
    ctx.need_ctl(lineno, cmd)
    reply = ctx.ctl.send(line if cmd != "ctl" else " ".join(args))
    return "ok", reply


def verb_rs485(ctx, lineno, cmd, args, line):
    # rs485 silence|restore ADDR: take a simulated device off the bus (it
    # holds its address and answers nothing, exactly an absent unit) or put
    # it back, while the guest keeps polling. This is the only runtime lever
    # on the roster ui-drive.sh --rs485-devices fixed at boot, and it exists
    # so a flow can prove a screen noticing an adopted device go quiet
    # (flare-edge #184). No simulated bus in this run is fatal: the step
    # could not check anything, and a silent skip would read as a pass.
    if ctx.rs485_control is None:
        return "fatal", ("no simulated RS485 bus in this run "
                          "(ui-drive.sh --rs485-devices)")
    if len(args) != 2 or args[0] not in ("silence", "restore"):
        return "fatal", "expected: rs485 silence|restore ADDR"
    try:
        reply = rs485_send(ctx.rs485_control, f"{args[0]} {args[1]}")
    except OSError as e:
        reply = f"error control socket: {e}"
    return ("ok" if reply.startswith("ok") else "fail"), reply


def verb_assert_page(ctx, lineno, cmd, args, line):
    ctx.need_ctl(lineno, cmd)
    want = " ".join(args)
    got = ctx.ctl.send("page")
    return ("ok" if got == want else "fail"), (f"page is {got!r}" if got != want else "")


def verb_assert_hit(ctx, lineno, cmd, args, line):
    # assert_hit X Y CLASS [box=X1,Y1,WxH] [TEXT...]: a real tap at X,Y
    # would land on an object of CLASS, with that exact bounding box if one
    # is given, whose caption contains TEXT. This is the pre-tap self-check:
    # "the button moved" fails HERE, by name, so it is never mistaken for
    # the behaviour behind the button.
    #
    # The box is the strong identity. Two rows of a list share a class, and
    # an icon's caption is a glyph nobody wants in a spec; the geometry the
    # UI itself reports is what tells one instance from another.
    ctx.need_ctl(lineno, cmd)
    ok, detail = judge_hit(ctx.ctl, args[0], args[1], args[2], args[3:])
    return ("ok" if ok else "fail"), detail


def verb_wait_hit(ctx, lineno, cmd, args, line):
    # wait_hit X Y CLASS TIMEOUT_S [box=X1,Y1,WxH] [TEXT...]: assert_hit
    # polled every 0.5s until it holds or TIMEOUT_S passes. A page still
    # building after `nav` answers `hit` with whatever is there at that
    # instant, so a check one second in raced the layout (flare-edge #175).
    # The timeout is the author's bound on how long settling may take, not a
    # sleep: a check that holds early returns early, and the detail says how
    # long it took.
    ctx.need_ctl(lineno, cmd)
    x, y, cls, timeout_s = args[0], args[1], args[2], float(args[3])
    rest = args[4:]
    ok, detail, waited, fatal = poll_until(
        lambda: judge_hit(ctx.ctl, x, y, cls, rest), timeout_s)
    if fatal:
        return "fatal", detail
    return ("ok" if ok else "fail"), (f"waited {waited:.1f}s" if ok else detail)


def verb_assert_json(ctx, lineno, cmd, args, line):
    # assert_json PATH OP VALUE: one read of the guest's status JSON (see
    # eval_json). PATH not (yet) in the doc is a `fail`, since the doc is a
    # live snapshot, not a schema.
    ctx.need_ctl(lineno, cmd)
    ok, detail = eval_json(ctx.ctl, args[0], args[1], args[2])
    return ("ok" if ok else "fail"), detail


def verb_wait_json(ctx, lineno, cmd, args, line):
    # wait_json PATH OP VALUE TIMEOUT_S: eval_json polled every 0.5s
    # (poll_until) until it holds or the deadline passes. The status file is
    # written by a 2s guest timer (webstatus.c), so this exists because a
    # value the doc will reach shortly is not the same fact as a value it
    # never reaches -- assert_json alone cannot tell those apart.
    ctx.need_ctl(lineno, cmd)
    timeout_s = float(args[3])
    ok, detail, waited, fatal = poll_until(
        lambda: eval_json(ctx.ctl, args[0], args[1], args[2]), timeout_s)
    if fatal:
        return "fatal", detail
    return ("ok" if ok else "fail"), (f"waited {waited:.1f}s" if ok else detail)


def verb_assert_stat(ctx, lineno, cmd, args, line):
    # assert_stat FIELD OP VALUE: FIELD off a fresh `stats` reply.
    ctx.need_ctl(lineno, cmd)
    field, op, value = args[0], args[1], args[2]
    if field not in ("cpu", "fps", "render", "idle"):
        return "fail", (f"unknown stat field {field!r} "
                         f"(known: cpu, fps, render, idle)")
    stats = parse_stats(ctx.ctl.send("stats"))
    if field not in stats:
        return "fail", "no such field"
    try:
        ok = apply_op(op, stats[field], value)
    except (TypeError, ValueError) as e:
        return "fail", str(e)
    return ("ok" if ok else "fail"), ("" if ok else f"{field} is {stats[field]}")


def verb_capture_region(ctx, lineno, cmd, args, line):
    # capture_region NAME X Y W H TOLERANCE: (re)writes NAME in the refs
    # file from a fresh screendump. Both hashes are stored regardless of
    # TOLERANCE so a later spec edit can change tolerance class without a
    # recapture.
    ctx.need_imgtools(lineno, cmd)
    name = args[0]
    x, y, w, h = int(args[1]), int(args[2]), int(args[3]), int(args[4])
    tolerance = args[5]
    if tolerance not in ("exact", "loose", "structural"):
        return "fail", (f"unknown tolerance {tolerance!r} "
                         f"(known: exact, loose, structural)")
    path, err = safe_out_path(ctx.outdir, name, prefix="capture_", suffix=".ppm")
    if err:
        return "fail", f"region name {err}"
    rpc(ctx.s, ctx.f, {"execute": "screendump", "arguments": {"filename": path}})
    try:
        cropped = imgtools.crop(imgtools.load_ppm(path), x, y, w, h)
    except (ValueError, OSError) as e:
        return "fatal", f"cannot capture {name!r}: {e}"
    ctx.refs[name] = {
        "x": x, "y": y, "w": w, "h": h, "tolerance": tolerance,
        "phash": f"{imgtools.phash(cropped):016x}",
        "structural": imgtools.structural(cropped).hex(),
    }
    save_refs(ctx.refs_path, ctx.refs)
    return "ok", f"captured box={x},{y},{w}x{h}"


def verb_assert_region(ctx, lineno, cmd, args, line):
    # assert_region NAME [TOLERANCE]: fresh screendump, crop to NAME's
    # stored box, compare at NAME's stored tolerance. A missing NAME is
    # FATAL for this step: an uncaptured reference is a spec/authoring gap,
    # not a UI defect the run should merely `fail` on. So is a TOLERANCE
    # other than the one NAME was captured under: the script and the
    # reference would be two claims about the same pixels, and judging by
    # either alone would hide that.
    ctx.need_imgtools(lineno, cmd)
    name = args[0]
    want_tol = args[1] if len(args) > 1 else None
    ref = ctx.refs.get(name)
    if ref is not None and want_tol and want_tol != ref.get("tolerance"):
        return "fatal", (f"reference {name!r} was captured as {ref.get('tolerance')}, "
                          f"script expects {want_tol}: recapture or fix the spec")
    ref, cropped, err = ctx.fresh_region(name, "assert")
    if err:
        return "fatal", err
    ok, detail = imgtools.compare(ref, cropped, ref["tolerance"])
    return ("ok" if ok else "fail"), detail


def verb_wait_region(ctx, lineno, cmd, args, line):
    # wait_region NAME TOLERANCE TIMEOUT_S: assert_region polled every 0.5s,
    # a fresh screendump each time, until the crop matches or the deadline
    # passes (flare-edge #175). A missing reference or a foreign tolerance
    # is fatal exactly as for assert_region: waiting cannot fix either, so
    # polling stops on the first such answer.
    ctx.need_imgtools(lineno, cmd)
    name, want_tol, timeout_s = args[0], args[1], float(args[2])
    ref = ctx.refs.get(name)
    if ref is not None and want_tol != ref.get("tolerance"):
        return "fatal", (f"reference {name!r} was captured as {ref.get('tolerance')}, "
                          f"script expects {want_tol}: recapture or fix the spec")
    fault = []

    def check_region():
        r, cropped, err = ctx.fresh_region(name, "assert")
        if err:
            fault.append(err)
            return True, err
        return imgtools.compare(r, cropped, r["tolerance"])

    ok, detail, waited, poll_fatal = poll_until(check_region, timeout_s)
    if fault:
        return "fatal", fault[0]
    if poll_fatal:
        return "fatal", detail
    return ("ok" if ok else "fail"), (f"waited {waited:.1f}s" if ok else detail)


def verb_assert_ocr(ctx, lineno, cmd, args, line):
    # assert_ocr NAME REGEX: same crop as assert_region, OCR'd through
    # `tesseract`, REGEX searched (re.search, spaces allowed like
    # assert_hit's TEXT) against the extracted text. No tesseract on PATH,
    # or no reference for NAME, is FATAL: pixels were never actually
    # checked, so this must never look like a skipped-but-passing step.
    ctx.need_imgtools(lineno, cmd)
    name, pattern = args[0], " ".join(args[1:])
    if shutil.which("tesseract") is None:
        return "fatal", "tesseract not installed"
    ref, cropped, err = ctx.fresh_region(name, "ocr")
    if err:
        return "fatal", err
    # fresh_region() above already ran NAME through safe_out_path() and
    # bailed on `err` if it could escape outdir, so it is safe to build this
    # second path from NAME directly.
    png_path = os.path.join(os.path.abspath(ctx.outdir), f"ocr_{name}.png")
    try:
        write_png(cropped, png_path)
        text = subprocess.run(
            ["tesseract", png_path, "stdout"],
            capture_output=True, text=True, timeout=20, check=True,
        ).stdout
    except Exception as e:
        return "fatal", f"ocr failed: {e}"
    if re.search(pattern, text):
        return "ok", ""
    return "fail", f"text was {text.strip()!r}"


# verb NAME -> handler. Several names share one handler (scroll/home;
# page/hit/stats/ctl) the same way they shared one elif arm before; the
# handler still gets the matched name as `cmd` for messages that name it.
VERBS = {
    "shot": verb_shot,
    "tap": verb_tap,
    "swipe": verb_swipe,
    "fling": verb_fling,
    "sleep": verb_sleep,
    "echo": verb_echo,
    "nav": verb_nav,
    "wake": verb_wake,
    "scroll": verb_scroll_home,
    "home": verb_scroll_home,
    "page": verb_query,
    "hit": verb_query,
    "stats": verb_query,
    "ctl": verb_query,
    "rs485": verb_rs485,
    "assert_page": verb_assert_page,
    "assert_hit": verb_assert_hit,
    "wait_hit": verb_wait_hit,
    "assert_json": verb_assert_json,
    "wait_json": verb_wait_json,
    "assert_stat": verb_assert_stat,
    "capture_region": verb_capture_region,
    "assert_region": verb_assert_region,
    "wait_region": verb_wait_region,
    "assert_ocr": verb_assert_ocr,
}


def drive(s, f, script_path, outdir, size, ctl_path=None, console_path=None, refs_path=None,
          rs485_control=None):
    os.makedirs(outdir, exist_ok=True)
    with open(script_path) as fh:
        lines = fh.readlines()

    ctl = Ctl(ctl_path) if ctl_path else None
    results = open(os.path.join(outdir, "results.jsonl"), "w")
    counts = {"ok": 0, "fail": 0, "fatal": 0}

    def record(lineno, cmd, status, detail=""):
        counts[status] += 1
        results.write(json.dumps({"line": lineno, "cmd": cmd, "status": status,
                                  "detail": detail}) + "\n")
        results.flush()
        tag = {"ok": "", "fail": "FAIL ", "fatal": "FATAL "}[status]
        print(f"  [{lineno}] {tag}{cmd}{(' -- ' + detail) if detail else ''}", flush=True)

    refs_path = refs_path or os.path.join(outdir, "refs.json")
    ctx = Ctx(s, f, size, ctl, load_refs(refs_path), refs_path, rs485_control, outdir, script_path)
    console = ConsoleWatch(console_path)

    for lineno, raw in enumerate(lines, 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        parts = line.split()
        cmd, args = parts[0], parts[1:]

        handler = VERBS.get(cmd)
        if handler is None:
            results.close()  # the rows before this one are still evidence
            sys.exit(f"FATAL: {script_path}:{lineno}: unknown command '{cmd}'")

        try:
            result = handler(ctx, lineno, cmd, args, line)
        except (RuntimeError, OSError, ValueError, IndexError, TypeError) as e:
            # poll_until already turns this into a fatal row for the wait_*
            # verbs; every other verb reaches ctl.send() (nav, wake,
            # scroll/home, page/hit/stats/ctl, assert_page, assert_hit,
            # assert_json, assert_stat) or the QMP socket via rpc()
            # (tap/swipe/fling/shot, the region verbs) with no guard of its
            # own, so without this a dying channel or a gone QMP socket ends
            # the whole run as an unhandled traceback instead of one fatal
            # step. ValueError/IndexError/TypeError are caught for the same
            # reason: several verbs parse their own arguments with bare
            # int()/float()/positional indexing before any handler-local
            # guard (tap, swipe, fling, sleep, wait_hit, wait_json,
            # capture_region, wait_region), and a malformed or missing
            # argument -- a typo'd coordinate, a hand-edited *.txt script, a
            # future flowc.py bug -- used to raise straight out of drive()
            # and silently drop every row from that line onward, the exact
            # truncated-run failure mode this file exists to rule out
            # (flare-edge #244).
            # Recorded the same as any other fatal row: the run keeps going
            # past it.
            result = ("fatal", str(e))
        if result is not None:
            record(lineno, line, *result)

        crashed = console.check()
        if crashed:
            record(lineno, line, "fatal", crashed)
            break

    results.close()
    print(f"== {counts['ok']} ok, {counts['fail']} failed, {counts['fatal']} fatal "
          f"-> {os.path.join(outdir, 'results.jsonl')}", flush=True)
    if counts["fail"] or counts["fatal"]:
        sys.exit(1)


def flag(argv, name, default=None):
    """One optional `NAME VALUE` pair out of argv, or DEFAULT when NAME is
    absent. Factored so a sixth optional flag is a one-line call instead of
    another hand-rolled argv.index() lookup -- and so a copy-pasted lookup
    can no longer search for one flag while reporting a different one."""
    return argv[argv.index(name) + 1] if name in argv else default


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    path, cmd = sys.argv[1], sys.argv[2]
    need = {"screendump": 4, "tap": 5, "quit": 3, "drive": 5}
    if cmd not in need:
        sys.exit(f"unknown command {cmd}\n{__doc__}")
    if len(sys.argv) < need[cmd]:
        sys.exit(f"{cmd}: missing argument(s)\n{__doc__}")

    size = int(flag(sys.argv, "--size", 720))
    ctl_path = flag(sys.argv, "--ctl")
    console_path = flag(sys.argv, "--console")
    refs_path = flag(sys.argv, "--refs")
    rs485_control = flag(sys.argv, "--rs485-control")

    s = socket.socket(socket.AF_UNIX)
    s.settimeout(QMP_TIMEOUT_S)
    s.connect(path)
    f = s.makefile("r")
    f.readline()  # greeting banner
    rpc(s, f, {"execute": "qmp_capabilities"})

    if cmd == "screendump":
        rpc(s, f, {"execute": "screendump", "arguments": {"filename": sys.argv[3]}})
    elif cmd == "tap":
        do_tap(s, f, int(sys.argv[3]), int(sys.argv[4]))
    elif cmd == "drive":
        drive(s, f, sys.argv[3], sys.argv[4], size, ctl_path, console_path, refs_path,
              rs485_control)
    elif cmd == "quit":
        s.sendall(b'{"execute":"quit"}\n')


if __name__ == "__main__":
    main()