#!/usr/bin/env python3
"""Unknown-call discovery over firmware control surfaces (H4, "Sandsifter for RM").

The problem this tool attacks: NVIDIA ships new GSP images on a hidden
schedule and new RM messages/controls/services appear without documentation.
Instead of waiting for a new control to be noticed, sweep the input space and
ask the firmware itself which calls exist, which code path each one takes, and
what it changes.

Each probe uses a fresh QEMU session, settles the firmware through its
self-test, hooks the command dispatcher (and every handler the control block
declares), applies exactly one command with a deterministic argument pattern,
runs a bounded window and classifies the outcome:

    known-handler  the documented/id-table handler for this id ran
    injected-slot  the id routed through the injected-service handler table
    default-path   the dispatcher's catch-all ran (undiscovered id)
    trap           the id reached an ``ebreak`` (debugger trap)
    fault          the id produced a load/store/fetch/illegal exception
    no-progress    the mailbox was never consumed within the budget

Findings are written to ``findings.jsonl`` and a markdown report groups ids by
class, listing unknown/undiscovered ids explicitly.

Surfaces (``--surface``):

* ``demo`` -- the bundled uproc mailbox command space (works today on QEMU);
* ``rm``   -- an RM control/method entry point for the future simulator or a
  live session.  It is implemented behind the same :class:`Surface`
  abstraction (a probe function + an effect observer), so a real RM surface is
  a new class, not a rewrite.  When no entry point is exposed the tool prints
  a documented message and exits cleanly.

Safety: the bundled demo (or an explicitly allowed foreign ELF) under QEMU,
fresh session per probe, bounded budgets, host writes limited to ``--out``.
Never hardware, a guest, an existing gdbstub or the native simulator by
default.  See docs/19-CONTROL-DISCOVERY.md.

Run from the tool root:

    python3 harness/control_discovery.py --surface demo --out dumps/discovery
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import shutil
import struct
import sys
import time
from collections import Counter
from dataclasses import dataclass, field

try:
    from .common import (
        DEMO_ELF,
        MAX_LOG_BYTES,
        ensure_dir,
        load_target,
        qemu_available,
        read_faults_window,
        sha256_file,
    )
except ImportError:  # direct script execution (python3 harness/control_discovery.py)
    from common import (  # type: ignore[no-redef]
        DEMO_ELF,
        MAX_LOG_BYTES,
        ensure_dir,
        load_target,
        qemu_available,
        read_faults_window,
        sha256_file,
    )

from greeneye import launch  # noqa: E402
from greeneye.errors import GreenEyeError, RspTimeout  # noqa: E402

FORMAT = "greeneye-control-discovery/1"
MASK64 = (1 << 64) - 1

#: gsp_ctrl field offsets (firmware/uproc_demo/gsp_ctrl.schema.json).
CTRL_MAGIC = 0
CTRL_STATE = 8
CTRL_RPC_COUNT = 16
CTRL_LAST_RPC = 24
CTRL_MAILBOX = 32
CTRL_ARG0 = 40
CTRL_ARG1 = 48
CTRL_ARG2 = 56
CTRL_RESULT = 64
CTRL_BUF = 72
CTRL_BUF_SIZE = 256
CTRL_HANDLER_COUNT = 328
CTRL_HANDLERS = 336
CTRL_SIZE = 400

#: The demo dispatcher's documented switch cases.
DEMO_CASE_HANDLERS = {
    1: "rpc_ping",
    2: "rpc_echo",
    3: "rpc_fill",
    4: "rpc_sum",
    5: "rpc_state",
    6: "rpc_fault",
}
#: Mailbox commands 16..23 select handler-table slots 0..7 (docs/12 section 4).
INJECT_CMD_BASE = 16
INJECT_CMD_END = 24

SETTLE_MARKER = "SELFTEST DONE"
SETTLE_SLICE = 0.02
RUN_SLICE = 0.02
TRAP_GRACE = 0.05
#: Bound on instruction-driven single-steps used to retire a dispatcher's
#: store sequence after the wall-clock observation window expires mid-epilogue
#: (keeps findings byte-reproducible instead of timer-dependent).
EPILOGUE_MAX_STEPS = 64
#: gdb signals that carry a hazard (illegal, trap, bus, segv); SIGINT from the
#: budget interrupt is wall-clock dependent and must never enter a finding.
HAZARD_SIGNALS = (4, 5, 7, 11)

PATTERNS = ("zero", "sentinel", "boundary")
OUTCOMES = (
    "known-handler",
    "injected-slot",
    "method-return",
    "default-path",
    "trap",
    "fault",
    "no-progress",
    "settle-failed",
    "error",
)
DEFAULT_RECORD_CLASSES = ("injected-slot", "method-return", "trap", "fault", "default-path", "settle-failed", "error")


# ---------------------------------------------------------------------------
# deterministic probe generation
# ---------------------------------------------------------------------------


@dataclass
class Probe:
    index: int
    cmd: int
    pattern: str
    args: tuple[int, int, int]

    @property
    def name(self) -> str:
        return f"{self.index:04d}-id{self.cmd:03d}-{self.pattern}"


def pattern_args(pattern: str, seed: int, cmd: int) -> tuple[int, int, int]:
    """Deterministic argument vector per (pattern, seed, cmd).

    The per-command RNG stream keeps sentinels distinct across ids while a
    fixed seed makes the whole sweep replayable.
    """
    if pattern == "zero":
        return (0, 0, 0)
    rng = random.Random(seed * 1_000_003 + cmd * 97 + 0x5EED)
    if pattern == "sentinel":
        return tuple(rng.getrandbits(64) for _ in range(3))  # type: ignore[return-value]
    if pattern == "boundary":
        return (MASK64, 1, 0x7FFF_FFFF_FFFF_FFFF)
    raise GreenEyeError(f"unknown pattern {pattern!r} (expected one of {', '.join(PATTERNS)})")


def generate_probes(max_id: int, patterns: list[str], seed: int, iterations: int = 0) -> list[Probe]:
    probes: list[Probe] = []
    for cmd in range(1, max_id + 1):
        for pattern in patterns:
            probes.append(Probe(index=len(probes), cmd=cmd, pattern=pattern, args=pattern_args(pattern, seed, cmd)))
    if iterations > 0:
        probes = probes[:iterations]
    return probes


# ---------------------------------------------------------------------------
# effect observer
# ---------------------------------------------------------------------------


@dataclass
class WatchedField:
    name: str
    offset: int
    size: int
    fmt: str  # u8 | u32 | u64 | bytes


@dataclass
class EffectObserver:
    """Captures the effect half of a surface: field snapshots + store PCs."""

    watched: tuple[WatchedField, ...] = ()
    hits: list[dict] = field(default_factory=list)

    def reset(self) -> None:
        self.hits.clear()

    def arm(self, session, base: int) -> None:
        """Install write watchpoints for every watched field."""
        for watched in self.watched:
            if watched.fmt == "bytes":
                continue  # byte-range watchpoints are not portable; hash at snapshot time
            session.watch(base + watched.offset, size=watched.size, kind="w", callback=self._make_callback(watched))

    def _make_callback(self, watched: WatchedField):
        def callback(hit) -> None:
            self.hits.append(
                {
                    "field": watched.name,
                    "pc": hit.pc,
                    "symbol": hit.symbol.name if hit.symbol else None,
                    "addr": hit.addr,
                }
            )

        return callback

    def snapshot(self, session, base: int) -> dict:
        values: dict = {}
        for watched in self.watched:
            raw = session.memory.read(base + watched.offset, watched.size)
            if watched.fmt == "bytes":
                values[watched.name] = {
                    "sha256": hashlib.sha256(raw).hexdigest()[:16],
                    "len": len(raw),
                    "preview": raw[:16].hex(),
                }
            elif watched.fmt == "u8":
                values[watched.name] = raw[0]
            elif watched.fmt == "u32":
                values[watched.name] = struct.unpack_from("<I", raw, 0)[0]
            else:
                values[watched.name] = struct.unpack_from("<Q", raw, 0)[0]
        return values

    def delta(self, before: dict, after: dict) -> dict:
        changed = [
            {"field": name, "before": before.get(name), "after": after.get(name)}
            for name in before
            if before.get(name) != after.get(name)
        ]
        return {"changed": changed, "watch": list(self.hits)}


# ---------------------------------------------------------------------------
# surfaces
# ---------------------------------------------------------------------------


class Surface:
    """A sweepable input surface: command space + probe action + effects.

    A future surface is a subclass; the sweep, classification and reporting
    machinery is shared.  ``probe`` applies one command, ``observer`` returns
    the effect model, and ``handler_map``/``command_for_addr`` model the
    entry points so ids can be attributed to code.
    """

    name = "abstract"

    def describe(self) -> str:
        return self.name

    def available(self) -> tuple[bool, str | None]:
        return True, None

    def documented_handlers(self, session) -> dict[int, str]:
        return {}

    def handler_map(self, session) -> dict[int, tuple[int, str | None]]:
        """Map command id -> (entry address, symbol name) for declared handlers."""
        return {}

    def command_hook(self, session) -> str | None:
        return None

    def is_injected(self, cmd: int) -> bool:
        return False

    def probe(self, session, cmd: int, args: tuple[int, int, int], *, timeout: float) -> int | None:
        """Apply one command. Returns a direct return value when the surface is
        a call entry point, else None (effects carry the reply)."""
        raise NotImplementedError

    def observer(self) -> EffectObserver:
        return EffectObserver()

    def command_for_addr(self, session, addr: int) -> list[dict]:
        """Ids whose declared handler resolves to ``addr`` (trigger discovery)."""
        out: list[dict] = []
        for cmd, (entry, name) in sorted(self.handler_map(session).items()):
            if entry == addr:
                out.append({"id": cmd, "handler": name})
        return out


class DemoSurface(Surface):
    """The uproc demo mailbox surface (gsp_ctrl.handlers + 1..6 switch)."""

    name = "demo"

    def documented_handlers(self, session) -> dict[int, str]:
        return {
            cmd: name for cmd, name in DEMO_CASE_HANDLERS.items() if session.symbols.get(name) is not None
        }

    def handler_map(self, session) -> dict[int, tuple[int, str | None]]:
        mapping: dict[int, tuple[int, str | None]] = {}
        for cmd, name in self.documented_handlers(session).items():
            mapping[cmd] = (session.symbols.lookup(name), name)
        ctrl = session.symbols.get("gsp_ctrl")
        if ctrl is not None:
            for slot in range(8):
                entry = session.memory.read_u64(ctrl.addr + CTRL_HANDLERS + slot * 8)
                if entry:
                    symbol = session.symbols.resolve(entry)
                    mapping[INJECT_CMD_BASE + slot] = (entry, symbol.name if symbol else None)
        return mapping

    def command_hook(self, session) -> str | None:
        return "rpc_dispatch" if session.symbols.get("rpc_dispatch") is not None else None

    def is_injected(self, cmd: int) -> bool:
        return INJECT_CMD_BASE <= cmd < INJECT_CMD_END

    def probe(self, session, cmd: int, args: tuple[int, int, int], *, timeout: float) -> int | None:
        ctrl = session.symbols.lookup("gsp_ctrl")
        # Args first: the firmware polls the mailbox last and consumes all four.
        session.memory.write_u64(ctrl + CTRL_ARG0, args[0])
        session.memory.write_u64(ctrl + CTRL_ARG1, args[1])
        session.memory.write_u64(ctrl + CTRL_ARG2, args[2])
        session.memory.write_u64(ctrl + CTRL_MAILBOX, cmd)
        return None

    def observer(self) -> EffectObserver:
        return EffectObserver(
            watched=(
                WatchedField("state", CTRL_STATE, 4, "u32"),
                WatchedField("rpc_count", CTRL_RPC_COUNT, 8, "u64"),
                WatchedField("last_rpc", CTRL_LAST_RPC, 8, "u64"),
                WatchedField("result", CTRL_RESULT, 8, "u64"),
                WatchedField("buf", CTRL_BUF, CTRL_BUF_SIZE, "bytes"),
            )
        )

    def snapshot_extra(self, session, base: int) -> dict:
        raw = session.memory.read(base, CTRL_SIZE)
        return {
            "magic": struct.unpack_from("<I", raw, CTRL_MAGIC)[0],
            "handler_count": struct.unpack_from("<Q", raw, CTRL_HANDLER_COUNT)[0],
        }


class RMSurface(Surface):
    """Future RM control/method entry point (roadmap H1.3/H4).

    On a simulator or live session that exposes an RM method dispatcher, the
    surface is: call the method entry with ``(control, a0, a1, a2)`` and read
    the reply buffer at ``reply_addr``.  Today no such entry is exposed, so
    ``available()`` reports why and the CLI exits cleanly.
    """

    name = "rm"

    def __init__(self, entry: str | None = None, reply_addr: int | None = None, reply_size: int = 0x100):
        self.entry = entry
        self.reply_addr = reply_addr
        self.reply_size = reply_size

    def describe(self) -> str:
        bits = [f"entry={self.entry or '<none>'}"]
        bits.append(f"reply={self.reply_addr:#x}/{self.reply_size:#x}" if self.reply_addr else "reply=<none>")
        return "rm(" + ", ".join(bits) + ")"

    def available(self) -> tuple[bool, str | None]:
        if not self.entry:
            return False, (
                "surface 'rm' unavailable: no RM control/method entry point is exposed by the "
                "simulator or a live session yet (roadmap H1.3). Pass --rm-entry SYMBOL once the "
                "RM method dispatcher exists; the demo surface works today with --surface demo."
            )
        return True, None

    def is_injected(self, cmd: int) -> bool:
        return False

    def probe(self, session, cmd: int, args: tuple[int, int, int], *, timeout: float) -> int | None:
        return session.call(self.entry, (cmd, args[0], args[1], args[2]), timeout=timeout)

    def observer(self) -> EffectObserver:
        if not self.reply_addr:
            return EffectObserver()
        return EffectObserver(
            watched=(WatchedField("reply", self.reply_addr, self.reply_size, "bytes"),)
        )

    def snapshot_extra(self, session, base: int) -> dict:
        return {}


def make_surface(args) -> Surface:
    if args.surface == "demo":
        return DemoSurface()
    if args.surface == "rm":
        return RMSurface(entry=args.rm_entry, reply_addr=args.rm_reply, reply_size=args.rm_reply_size)
    raise GreenEyeError(f"unknown surface {args.surface!r}")


# ---------------------------------------------------------------------------
# one probe
# ---------------------------------------------------------------------------


@dataclass
class Finding:
    probe: int
    cmd: int
    pattern: str
    args: list[int]
    outcome: str = "error"
    handler: str | None = None
    handler_entry: int | None = None
    documented_handler: str | None = None
    dispatch_seen: bool = False
    dispatch_pc: int | None = None
    dispatch_args: list[int] | None = None
    call_retval: int | None = None
    rpc_count_delta: int = 0
    effects: dict = field(default_factory=dict)
    fault: dict | None = None
    signals: list[dict] = field(default_factory=list)
    settle_ok: bool = True
    session: str | None = None
    reason: str | None = None

    def to_json(self) -> dict:
        return {
            "probe": self.probe,
            "cmd": self.cmd,
            "pattern": self.pattern,
            "args": self.args,
            "outcome": self.outcome,
            "handler": self.handler,
            "handler_entry": self.handler_entry,
            "documented_handler": self.documented_handler,
            "dispatch_seen": self.dispatch_seen,
            "dispatch_pc": self.dispatch_pc,
            "dispatch_args": self.dispatch_args,
            "call_retval": self.call_retval,
            "rpc_count_delta": self.rpc_count_delta,
            "effects": self.effects,
            "fault": self.fault,
            "signals": self.signals,
            "settle_ok": self.settle_ok,
            "session": self.session,
            "reason": self.reason,
        }


def _compact_fault(fault: dict) -> dict:
    return {
        "cause": fault.get("cause"),
        "desc": fault.get("desc"),
        "epc": fault.get("epc"),
        "tval": fault.get("tval"),
        "access": fault.get("access"),
    }


def classify_outcome(
    *,
    cmd: int,
    is_injected: bool,
    handler: str | None,
    dispatch_seen: bool,
    call_retval: int | None,
    faults: list[dict],
    signals: list[dict],
    settle_ok: bool,
) -> str:
    """Priority-ordered classification of one probe result (pure function)."""
    if not settle_ok:
        return "settle-failed"
    causes = [int(fault.get("cause", 0)) for fault in faults]
    if any(cause == 3 for cause in causes) or any(sig.get("signal") == 5 for sig in signals):
        return "trap"
    if causes or any(sig.get("signal") in (4, 7, 11) for sig in signals):
        return "fault"
    if call_retval is not None:
        return "method-return"
    if not dispatch_seen:
        return "no-progress"
    if is_injected and handler:
        return "injected-slot"
    if handler:
        return "known-handler"
    return "default-path"


class ProbeRunner:
    """Runs one probe in a fresh session and builds the finding."""

    def __init__(self, surface: Surface, options: "DiscoveryOptions", image, kernel_sha: str):
        self.surface = surface
        self.options = options
        self.image = image
        self.kernel_sha = kernel_sha

    def run(self, probe: Probe) -> Finding:
        finding = Finding(probe=probe.index, cmd=probe.cmd, pattern=probe.pattern, args=list(probe.args))
        staging = os.path.join(self.options.out, ".staging", probe.name)
        ensure_dir(staging)
        qemu_log = os.path.join(staging, "qemu-faults.log")
        # Record into staging unconditionally: the classification decides whether
        # the session is evidence worth keeping, and re-running a probe to
        # record it later would double the sweep cost.
        record_path = staging
        session = None
        try:
            extra_args = None
            if self.options.backend == "qemu":
                extra_args = ["-d", "guest_errors,unimp,int", "-D", qemu_log]
            session = launch(
                kernel=self.options.kernel,
                record=record_path,
                name=f"discovery-{probe.name}",
                backend=self.options.backend,
                extra_args=extra_args,
            )
            finding.settle_ok = self._settle(session, self.options.settle)
            if not finding.settle_ok:
                finding.reason = f"{SETTLE_MARKER!r} not observed within the settle budget"
                return self._finish(session, probe, finding)

            observer = self.surface.observer()
            before = observer.snapshot(session, self._base(session))
            before.update(self._extra(self.surface, session))
            observer.reset()
            observer.arm(session, self._base(session))

            dispatches: list[dict] = []
            handlers: list[dict] = []
            signals: list[dict] = []
            self._install_hooks(session, dispatches, handlers, signals)
            observer.reset()

            log_offset = os.path.getsize(qemu_log) if os.path.exists(qemu_log) else 0
            call_retval: int | None = None
            try:
                call_retval = self.surface.probe(session, probe.cmd, probe.args, timeout=self.options.timeout)
            except RspTimeout as exc:
                finding.reason = f"probe call timed out: {exc}"
            self._run_window(session, dispatches, handlers, signals, observer)

            faults: list[dict] = []
            if self.options.backend == "qemu":
                faults, _capped = read_faults_window(qemu_log, log_offset, self.options.max_log_bytes)

            after = observer.snapshot(session, self._base(session))
            after.update(self._extra(self.surface, session))
            effects = observer.delta(before, after)
            effects["before"] = before
            effects["after"] = after

            dispatch = dispatches[0] if dispatches else None
            handler = handlers[0] if handlers else None
            mapped = self.surface.handler_map(session)
            finding.documented_handler = self.surface.documented_handlers(session).get(probe.cmd)
            finding.dispatch_seen = dispatch is not None
            finding.dispatch_pc = dispatch.get("addr") if dispatch else None
            finding.dispatch_args = dispatch.get("args") if dispatch else None
            finding.handler = handler.get("symbol") if handler else None
            finding.handler_entry = handler.get("addr") if handler else None
            if handler is None and probe.cmd in mapped:
                entry, name = mapped[probe.cmd]
                finding.handler = name
                finding.handler_entry = entry
            finding.call_retval = call_retval
            finding.rpc_count_delta = int(effects.get("after", {}).get("rpc_count", 0)) - int(
                effects.get("before", {}).get("rpc_count", 0)
            )
            finding.effects = effects
            finding.fault = _compact_fault(faults[0]) if faults else None
            finding.signals = signals
            finding.outcome = classify_outcome(
                cmd=probe.cmd,
                is_injected=self.surface.is_injected(probe.cmd),
                handler=finding.handler,
                dispatch_seen=dispatch is not None,
                call_retval=call_retval,
                faults=faults,
                signals=signals,
                settle_ok=True,
            )
            return self._finish(session, probe, finding)
        except Exception as exc:  # one bad probe must not end the sweep
            finding.outcome = "error"
            finding.reason = f"{type(exc).__name__}: {exc}"
            return self._finish(session, probe, finding)

    # -- helpers -----------------------------------------------------------

    def _base(self, session) -> int:
        return session.symbols.lookup("gsp_ctrl") if session.symbols.get("gsp_ctrl") else 0

    def _extra(self, surface: Surface, session) -> dict:
        extra = getattr(surface, "snapshot_extra", None)
        if extra is None:
            return {}
        return extra(session, self._base(session))

    def _install_hooks(self, session, dispatches: list[dict], handlers: list[dict], signals: list[dict]) -> None:
        def on_signal(hit) -> None:
            stop = hit.stop
            if stop is not None and stop.signal in HAZARD_SIGNALS:
                signals.append({"signal": stop.signal, "pc": hit.pc})

        session.on_event("stop", on_signal)
        hook_name = self.surface.command_hook(session)
        if hook_name:
            session.hook(
                hook_name,
                on_enter=lambda hit: dispatches.append(
                    {
                        "addr": hit.addr,
                        "symbol": hook_name,
                        "args": [hit.arg(0), hit.arg(1), hit.arg(2), hit.arg(3)],
                        "caller": hit.regs.get("ra", 0),
                    }
                ),
                listener=f"discovery:{hook_name}",
            )
        for cmd, (entry, name) in sorted(self.surface.handler_map(session).items()):
            session.hook(
                entry,
                on_enter=lambda hit, name=name: handlers.append(
                    {
                        "addr": hit.addr,
                        "symbol": hit.symbol.name if hit.symbol else name,
                        "args": [hit.arg(0), hit.arg(1), hit.arg(2)],
                        "caller": hit.regs.get("ra", 0),
                    }
                ),
                listener=f"discovery:handler-{entry:#x}",
            )

    def _settle(self, session, budget: float) -> bool:
        deadline = time.monotonic() + budget
        while time.monotonic() < deadline:
            if SETTLE_MARKER in session.target.read_uart():
                return True
            session.run(duration=SETTLE_SLICE, max_stops=1)
        return SETTLE_MARKER in session.target.read_uart()

    def _run_window(self, session, dispatches, handlers, signals, observer) -> None:
        """Two-phase bounded window: wait for dispatch, then for completion.

        Completion is a ``result`` store (the dispatcher writes it after every
        handler) or a trap/grace for ids that never return to the dispatcher.
        The wall-clock loop is the budget for the no-progress/hung cases; when
        a dispatch has been seen the store sequence is then retired with
        bounded single-steps so the recorded effects cannot depend on where
        the timer happened to expire.
        """
        deadline = time.monotonic() + self.options.timeout
        handler_seen_at: float | None = None
        while time.monotonic() < deadline:
            if not session.target.alive():
                break
            remaining = max(0.005, min(RUN_SLICE, deadline - time.monotonic()))
            session.run(duration=remaining, max_stops=64)
            if handlers and handler_seen_at is None:
                handler_seen_at = time.monotonic()
            result_written = any(hit.get("field") == "result" for hit in observer.hits)
            trapped = any(sig.get("signal") == 5 for sig in signals) or session.pc < 0x100
            if dispatches and (result_written or trapped):
                break
            if handler_seen_at is not None and time.monotonic() - handler_seen_at >= TRAP_GRACE:
                # Slow/hung handler: one grace window after the entry, then stop.
                break
        self._retire_dispatch(session, dispatches, signals, observer)

    def _retire_dispatch(self, session, dispatches, signals, observer) -> None:
        """Finish a dispatched command's store sequence with single-steps.

        The demo dispatcher stores ``result``, ``last_rpc`` and finally
        ``rpc_count``; a wall-clock observation window can expire between
        those stores, which made findings timing-dependent.  Once a dispatch
        has been seen, step (bounded and instruction-driven) until the final
        watched store retires, a trap is seen, or the budget is exhausted.
        """
        if not dispatches or not any(watched.name == "rpc_count" for watched in observer.watched):
            return
        if any(hit.get("field") == "rpc_count" for hit in observer.hits):
            return
        if any(sig.get("signal") == 5 for sig in signals):
            return
        for _ in range(EPILOGUE_MAX_STEPS):
            if not session.target.alive():
                return
            try:
                session.step()
            except (RspTimeout, GreenEyeError):
                return
            if any(hit.get("field") == "rpc_count" for hit in observer.hits):
                return
            if any(sig.get("signal") == 5 for sig in signals):
                return

    def _keep(self, outcome: str) -> bool:
        return self.options.record_any or outcome in self.options.record_classes

    def _finish(self, session, probe: Probe, finding: Finding) -> Finding:
        if session is not None:
            try:
                session.close()
            except Exception:
                pass
        staging = os.path.join(self.options.out, ".staging", probe.name)
        if self._keep(finding.outcome) and os.path.isdir(staging):
            final = os.path.join(self.options.out, "sessions", probe.name)
            ensure_dir(os.path.dirname(final))
            if os.path.isdir(final):
                shutil.rmtree(final, ignore_errors=True)
            os.replace(staging, final)
            finding.session = os.path.join("sessions", probe.name)
        else:
            shutil.rmtree(staging, ignore_errors=True)
        return finding


# ---------------------------------------------------------------------------
# report
# ---------------------------------------------------------------------------


def aggregate(findings: list[Finding]) -> dict:
    per_id: dict[int, dict] = {}
    for finding in findings:
        row = per_id.setdefault(
            finding.cmd,
            {
                "id": finding.cmd,
                "patterns": [],
                "outcomes": [],
                "handlers": [],
                "documented_handler": finding.documented_handler,
                "results": [],
                "rpc_count_delta": 0,
                "sessions": [],
                "faults": [],
            },
        )
        row["patterns"].append(finding.pattern)
        row["outcomes"].append(finding.outcome)
        if finding.handler:
            row["handlers"].append(finding.handler)
        result = (finding.effects or {}).get("after", {}).get("result")
        if result is not None and finding.dispatch_seen:
            row["results"].append(result)
        row["rpc_count_delta"] += finding.rpc_count_delta
        if finding.session:
            row["sessions"].append(finding.session)
        if finding.fault:
            row["faults"].append(finding.fault)
    for row in per_id.values():
        row["outcomes"] = sorted(set(row["outcomes"]))
        row["handlers"] = sorted(set(row["handlers"]))
        row["results"] = sorted(set(row["results"]))
    return dict(sorted(per_id.items()))


def primary_class(row: dict) -> str:
    """One label per id for grouping: handle, injected, trap, fault, unknown."""
    outcomes = set(row["outcomes"])
    if "trap" in outcomes:
        return "trap/fault"
    if "fault" in outcomes:
        return "trap/fault"
    if "known-handler" in outcomes and row["documented_handler"]:
        return "known-handler"
    if "injected-slot" in outcomes:
        return "injected-slot"
    if "method-return" in outcomes:
        return "method-return"
    if outcomes & {"default-path", "no-progress"}:
        return "unknown/undiscovered"
    return "unknown/undiscovered"


def render_markdown(report: dict, options: "DiscoveryOptions") -> str:
    per_id = report["ids"]
    outcome_counts = Counter(finding["outcome"] for finding in report["findings"])
    grouped: dict[str, list[int]] = {}
    for row in per_id.values():
        grouped.setdefault(primary_class(row), []).append(row["id"])
    lines = [
        "# risc-v-gsp-rm-debugger control discovery report",
        "",
        "> Status: implemented, documented and tested (audit: docs/27-TOOL-AUDIT.md; value demos: docs/25-VALUE-VALIDATION.md). Generated by `harness/control_discovery.py`.",
        "",
        f"* format: `{FORMAT}`",
        f"* surface: `{report['surface']}`",
        f"* kernel: `{report['kernel']}`",
        f"* kernel sha256: `{report['kernel_sha256']}`",
        f"* backend: `{report['backend']}`",
        f"* ids swept: 1..{report['max_id']} ({report['probes_run']} probes), "
        f"seed={report['seed']}, patterns={','.join(report['patterns'])}",
        f"* per-probe timeout: {report['timeout']}s; settle budget: {report['settle']}s",
        "",
        "## Outcome counts",
        "",
    ]
    lines += _table(
        ["outcome", "count"],
        [[outcome, outcome_counts[outcome]] for outcome in sorted(outcome_counts)],
    )
    lines += ["", "## Ids by class", ""]
    for klass in ("known-handler", "injected-slot", "method-return", "trap/fault", "unknown/undiscovered"):
        ids = grouped.get(klass, [])
        lines += [f"* **{klass}**: " + (", ".join(str(i) for i in ids) if ids else "_none_")]
    lines += ["", "## Known handler ids", ""]
    known = [row for row in per_id.values() if row["documented_handler"]]
    lines += _render_rows(known)
    lines += ["", "## Injected-slot dispatch ids", ""]
    injected = [row for row in per_id.values() if "injected-slot" in row["outcomes"]]
    lines += _render_rows(injected)
    lines += ["", "## Unknown / undiscovered ids", ""]
    unknown = [
        row
        for row in per_id.values()
        if not row["documented_handler"] and not (set(row["outcomes"]) & {"trap", "fault"})
    ]
    if unknown:
        lines += _render_rows(unknown)
        lines += [
            "",
            "These ids reached the dispatcher's catch-all (result ``~id``) and no handler "
            "table entry; they are the candidate space for new/unmapped controls.",
        ]
    else:
        lines += ["_none_"]
    lines += ["", "## Traps and faults", ""]
    hazards = [row for row in per_id.values() if set(row["outcomes"]) & {"trap", "fault"}]
    lines += _render_rows(hazards)
    lines += ["", "## All findings by id", ""]
    lines += _render_rows(list(per_id.values()))
    lines += ["", "## Reproduce", "", "```bash", report["reproduce"], "```", ""]
    return "\n".join(lines) + "\n"


def _render_rows(rows: list[dict]) -> list[str]:
    if not rows:
        return ["_none_"]
    return _table(
        ["id", "documented handler", "observed handler", "outcomes", "results", "rpc delta", "sessions"],
        [
            [
                row["id"],
                row["documented_handler"] or "-",
                ", ".join(row["handlers"]) or "-",
                ", ".join(row["outcomes"]),
                ", ".join(f"{value:#x}" for value in row["results"][:4]) or "-",
                row["rpc_count_delta"],
                ", ".join(f"`{path}`" for path in row["sessions"]) or "-",
            ]
            for row in rows
        ],
    )


def _table(headers: list[str], rows: list[list]) -> list[str]:
    lines = ["| " + " | ".join(headers) + " |", "|" + "|".join("---" for _ in headers) + "|"]
    for row in rows:
        lines.append("| " + " | ".join(str(cell) for cell in row) + " |")
    return lines


# ---------------------------------------------------------------------------
# orchestration
# ---------------------------------------------------------------------------


@dataclass
class DiscoveryOptions:
    surface: str
    out: str
    kernel: str = DEMO_ELF
    max_id: int = 63
    patterns: list[str] = field(default_factory=lambda: list(PATTERNS))
    seed: int = 1
    iterations: int = 0
    timeout: float = 0.25
    settle: float = 2.0
    record_all: bool = False
    record_classes: tuple[str, ...] = DEFAULT_RECORD_CLASSES
    backend: str = "qemu"
    allow_foreign_elf: bool = False
    max_log_bytes: int = MAX_LOG_BYTES
    rm_entry: str | None = None
    rm_reply: int | None = None
    rm_reply_size: int = 0x100

    @property
    def record_any(self) -> bool:
        return self.record_all


def run_discovery(options: DiscoveryOptions) -> dict:
    if options.backend == "qemu" and not qemu_available():
        raise GreenEyeError("qemu-system-riscv64 not found; install qemu-system-misc or set GREENEYE_QEMU")
    if not os.path.isfile(options.kernel):
        raise GreenEyeError(f"kernel not found: {options.kernel}")
    if not options.allow_foreign_elf and os.path.abspath(options.kernel) != os.path.abspath(DEMO_ELF):
        raise GreenEyeError(
            "control discovery is restricted to the bundled demo firmware; "
            "pass --allow-foreign-elf to sweep another ELF under QEMU"
        )
    image = load_target(options.kernel)
    kernel_sha = sha256_file(options.kernel)
    surface = make_surface(options)
    ensure_dir(options.out)
    probes = generate_probes(options.max_id, options.patterns, options.seed, options.iterations)
    runner = ProbeRunner(surface, options, image, kernel_sha)
    findings: list[Finding] = []
    findings_path = os.path.join(options.out, "findings.jsonl")
    with open(findings_path, "w", encoding="utf-8") as fh:
        for probe in probes:
            finding = runner.run(probe)
            findings.append(finding)
            fh.write(json.dumps(finding.to_json(), sort_keys=True) + "\n")
            fh.flush()
            print(
                f"[discovery] {finding.probe + 1}/{len(probes)} id={finding.cmd} {finding.pattern} "
                f"-> {finding.outcome}" + (f" ({finding.handler})" if finding.handler else ""),
                file=sys.stderr,
            )
    shutil.rmtree(os.path.join(options.out, ".staging"), ignore_errors=True)
    report = {
        "format": FORMAT,
        "tool": "harness/control_discovery.py",
        "surface": surface.describe(),
        "kernel": options.kernel,
        "kernel_sha256": kernel_sha,
        "backend": options.backend,
        "max_id": options.max_id,
        "patterns": options.patterns,
        "seed": options.seed,
        "timeout": options.timeout,
        "settle": options.settle,
        "probes_run": len(findings),
        "ids": aggregate(findings),
        "findings": [finding.to_json() for finding in findings],
        "reproduce": reproduce_command(options),
    }
    write_json(os.path.join(options.out, "discovery-report.json"), report)
    with open(os.path.join(options.out, "control-discovery.md"), "w", encoding="utf-8") as fh:
        fh.write(render_markdown(report, options))
    return report


def reproduce_command(options: DiscoveryOptions) -> str:
    foreign = " --allow-foreign-elf" if options.allow_foreign_elf else ""
    record = " --record-all" if options.record_all else ""
    rm = ""
    if options.surface == "rm":
        rm = f" --rm-entry {options.rm_entry}" if options.rm_entry else ""
        if options.rm_reply:
            rm += f" --rm-reply {options.rm_reply:#x}"
    return (
        f"python3 harness/control_discovery.py --surface {options.surface} --kernel {options.kernel} "
        f"--out {options.out} --max-id {options.max_id} --patterns {','.join(options.patterns)} "
        f"--seed {options.seed} --timeout {options.timeout} --settle {options.settle}"
        f"{record}{foreign}{rm}"
    )


def write_json(path: str, payload: dict) -> str:
    ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(payload, fh, indent=2, sort_keys=True)
        fh.write("\n")
    return path


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def parse_patterns(text: str) -> list[str]:
    patterns = [part.strip() for part in text.split(",") if part.strip()]
    unknown = [pattern for pattern in patterns if pattern not in PATTERNS]
    if unknown:
        raise argparse.ArgumentTypeError(f"unknown pattern(s) {unknown}; expected {', '.join(PATTERNS)}")
    return patterns


def parse_record_classes(text: str) -> tuple[str, ...]:
    classes = tuple(part.strip() for part in text.split(",") if part.strip())
    unknown = [klass for klass in classes if klass not in OUTCOMES]
    if unknown:
        raise argparse.ArgumentTypeError(f"unknown outcome class(es) {unknown}; expected {', '.join(OUTCOMES)}")
    return classes


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="control_discovery", description=__doc__.splitlines()[0])
    parser.add_argument("--surface", choices=("demo", "rm"), default="demo", help="input surface to sweep")
    parser.add_argument("--kernel", default=DEMO_ELF, help="demo/re-hosted ELF (default: bundled demo)")
    parser.add_argument("--out", required=True, help="output directory for findings and reports")
    parser.add_argument("--max-id", type=int, default=63, help="highest command id to sweep (0..4096)")
    parser.add_argument("--patterns", type=parse_patterns, default=list(PATTERNS), help="comma list: zero,sentinel,boundary")
    parser.add_argument("--seed", type=int, default=1, help="deterministic pattern seed")
    parser.add_argument("--iterations", type=int, default=0, help="cap total probes (0 = all)")
    parser.add_argument("--timeout", type=float, default=0.25, help="bounded observation window per probe (seconds)")
    parser.add_argument("--settle", type=float, default=2.0, help="budget to observe the self-test before probing")
    parser.add_argument("--record-all", action="store_true", help="record every probe (default: interesting classes only)")
    parser.add_argument(
        "--record-classes",
        type=parse_record_classes,
        default=DEFAULT_RECORD_CLASSES,
        help="outcome classes whose sessions are kept",
    )
    parser.add_argument("--backend", choices=("qemu", "native"), default="qemu", help="execution backend (qemu is the supported default)")
    parser.add_argument("--allow-foreign-elf", action="store_true", help="permit a non-demo ELF under QEMU")
    parser.add_argument("--max-log-bytes", type=int, default=MAX_LOG_BYTES, help="QEMU fault-log window per probe")
    parser.add_argument("--rm-entry", help="RM surface: method/control entry symbol (future, H1.3)")
    parser.add_argument("--rm-reply", type=lambda value: int(value, 0), help="RM surface: reply buffer address")
    parser.add_argument("--rm-reply-size", type=lambda value: int(value, 0), default=0x100, help="RM surface: reply buffer size")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not 1 <= args.max_id <= 4096:
        parser.error("--max-id must be in 1..4096")
    if args.iterations < 0:
        parser.error("--iterations must be >= 0")
    if not 0.01 <= args.timeout <= 30.0:
        parser.error("--timeout must be in 0.01..30.0")
    options = DiscoveryOptions(
        surface=args.surface,
        out=os.path.abspath(args.out),
        kernel=os.path.abspath(args.kernel),
        max_id=args.max_id,
        patterns=args.patterns,
        seed=args.seed,
        iterations=args.iterations,
        timeout=args.timeout,
        settle=args.settle,
        record_all=args.record_all,
        record_classes=args.record_classes,
        backend=args.backend,
        allow_foreign_elf=args.allow_foreign_elf,
        max_log_bytes=args.max_log_bytes,
        rm_entry=args.rm_entry,
        rm_reply=args.rm_reply,
        rm_reply_size=args.rm_reply_size,
    )
    surface = make_surface(args)
    available, reason = surface.available()
    if not available:
        print(f"[discovery] skipped: {reason}")
        return 0
    try:
        report = run_discovery(options)
    except GreenEyeError as exc:
        print(f"control discovery error: {exc}", file=sys.stderr)
        return 2
    counts = Counter(finding["outcome"] for finding in report["findings"])
    print(f"[discovery] {report['probes_run']} probe(s) on surface={report['surface']}, out={options.out}")
    for outcome in sorted(counts):
        print(f"[discovery] {outcome}: {counts[outcome]}")
    print(f"[discovery] wrote {os.path.join(options.out, 'findings.jsonl')}")
    print(f"[discovery] wrote {os.path.join(options.out, 'control-discovery.md')}")
    return 0


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())