#!/usr/bin/env python3
"""Drop-in replacement for the `expect` tool used by the low-level proof
scripts, using only the Python standard library (no external deps).

Drives a QEMU guest over its serial-on-stdio chardev: spawns QEMU with
stdin/stdout pipes, watches the guest output for expected patterns, sends
replies, and reports pass/fail. This removes the host `expect` package
dependency (which is documented but not always installed, and cannot be
installed without sudo).

Spec model: an ordered list of steps. Each step is `expect:<text>` (wait for
the literal text in guest output) or `send:<text>` (write text + CR to the
guest). After all steps, wait for `--pass <text>` to succeed, or `--fail
<text>` / timeout to fail.
"""
import argparse
import os
import select
import subprocess
import sys
import time


def main() -> int:
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument("--timeout", type=int, default=240)
    p.add_argument("--log", required=True)
    p.add_argument("--step", action="append", default=[])
    p.add_argument("--pass_marker", action="append", default=[])
    p.add_argument("--fail_marker", action="append", default=[])
    p.add_argument("qemu", nargs=argparse.REMAINDER)
    args = p.parse_args()

    qemu_cmd = args.qemu
    if qemu_cmd and qemu_cmd[0] == "--":
        qemu_cmd = qemu_cmd[1:]
    if not qemu_cmd:
        print("ERROR: no qemu command given", file=sys.stderr)
        return 2

    log_f = open(args.log, "wb", buffering=0)
    try:
        proc = subprocess.Popen(
            qemu_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            bufsize=0,
        )
    except FileNotFoundError as e:
        print(f"ERROR: cannot spawn qemu: {e}", file=sys.stderr)
        return 2

    out = proc.stdout
    assert out is not None, "qemu stdout pipe is required for serial interaction"

    deadline = time.monotonic() + args.timeout
    buf = bytearray()
    state = {"failed": False, "passes": set(), "died_early": False}

    def record(chunk: bytes) -> str:
        buf.extend(chunk)
        log_f.write(chunk)
        text = buf.decode("utf-8", errors="replace")
        for fm in args.fail_marker:
            if fm in text and fm not in state["passes"]:
                state["failed"] = True
        for pm in args.pass_marker:
            if pm in text:
                state["passes"].add(pm)
        return text

    def send(text: str) -> None:
        if proc.stdin:
            try:
                proc.stdin.write(text.encode() + b"\r")
                proc.stdin.flush()
            except (BrokenPipeError, OSError):
                pass

    steps = list(args.step)
    i = 0
    # Phase 1: walk expect/send steps in order. The deadline resets on every
    # step advance so each expect step gets its own full timeout window,
    # matching the per-pattern semantics of the expect tool's `set timeout`.
    while i < len(steps):
        if time.monotonic() > deadline or state["failed"]:
            break
        step = steps[i]
        if step.startswith("send:"):
            send(step[len("send:"):])
            i += 1
            deadline = time.monotonic() + args.timeout
            continue
        if step.startswith("expect:"):
            want = step[len("expect:"):]
            if want in buf.decode("utf-8", errors="replace"):
                i += 1
                deadline = time.monotonic() + args.timeout
                continue
        r, _, _ = select.select([out], [], [], 0.5)
        if r:
            chunk = os.read(out.fileno(), 65536)
            if not chunk:
                state["died_early"] = True
                break
            record(chunk)
        if proc.poll() is not None and not r:
            state["died_early"] = True
            break

    # Phase 2: wait for pass markers (or fail/timeout). Fresh window: markers
    # that were also expect steps are usually already recorded by phase 1.
    # Skipped when phase 1 broke early: the guest check that prints pass
    # markers is only sent by a completed step sequence, so waiting would
    # just double the failure time.
    if i >= len(steps):
        deadline = time.monotonic() + args.timeout
        while not state["failed"]:
            if state["passes"] and len(state["passes"]) == len(args.pass_marker):
                break
            if not args.pass_marker:
                break
            if time.monotonic() > deadline:
                break
            r, _, _ = select.select([out], [], [], 0.5)
            if r:
                chunk = os.read(out.fileno(), 65536)
                if not chunk:
                    state["died_early"] = True
                    break
                record(chunk)
            if proc.poll() is not None and not r:
                state["died_early"] = True
                break

    try:
        proc.terminate()
        try:
            proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            proc.kill()
    except OSError:
        pass
    log_f.close()

    if state["failed"]:
        print(f"ERROR: a --fail marker appeared; see {args.log}", file=sys.stderr)
        return 1
    if args.pass_marker and len(state["passes"]) != len(args.pass_marker):
        if state["died_early"]:
            print(f"ERROR: qemu exited before the check completed; see {args.log}", file=sys.stderr)
            return 3
        print(f"ERROR: expected pass markers not all seen; see {args.log}", file=sys.stderr)
        return 1
    return 0


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