#!/usr/bin/env python3
"""Hardware acceptance suite for a pureboot deployment.

`tools/check.sh` proves the protocol under simavr on every chip. This proves one
*board*: that the loader actually installed on it answers, that the memories
round-trip over the real link, that the application it flashes runs afterwards,
and that the refusals which keep a 512-byte slot alive still fire. Run it once
when a board is brought up, and again whenever the deployment moves — a new
clock, a new backend, new pins.

Every check derives its bounds from the info block the loader itself reports, so
nothing here is per-chip: the same run covers a 1 KiB tiny whose application
region is 510 usable bytes and a 128 KiB mega whose flash needs a bank in the
selector.

**This overwrites the board's application flash and EEPROM.** Capture them first
with `pbrig.py backup`, which verifies what it captured.

    tools/pbhw.py --programmer atmelice_isp --part t13 --port COM6 \
                  --autobaud --loader build/ab.bin --app build/pbapp.hex \
                  --marker APP
"""

from __future__ import annotations

import argparse
import pathlib
import sys
import tempfile

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import pbrig  # noqa: E402


class Suite:
    def __init__(self, rig: pbrig.Rig, work: pathlib.Path):
        self.rig = rig
        self.work = work
        self.results: list[tuple[str, bool, str]] = []

    def check(self, name: str, ok: bool, detail: str = "") -> bool:
        self.results.append((name, ok, detail))
        print(f"  {'PASS' if ok else 'FAIL'}  {name}" + (f"   {detail}" if detail else ""))
        return ok

    @staticmethod
    def _brief(text: str, limit: int = 78) -> str:
        return " | ".join(l.strip() for l in text.splitlines() if l.strip())[:limit]

    # ----------------------------------------------------------------- checks

    def identity(self) -> object | None:
        """The info block, which every later check takes its bounds from."""
        module = pbrig.load_pureboot(self.rig.d.pureboot)
        self.rig.reset()
        port = self.rig.open_port()  # wrapped for the echo where the line is shared
        try:
            loader = module.Loader(port)
            if self.rig.d.autobaud:
                loader.connect_autobaud(self.rig.d.wait)
            else:
                loader.connect(self.rig.d.wait)
            info = loader.info
            self.check("identity read", True, info.describe())
            return info
        except Exception as error:  # noqa: BLE001 — a dead link is a result
            self.check("identity read", False, str(error)[:70])
            return None
        finally:
            try:
                port.close()
            except Exception:
                pass

    def scan(self) -> None:
        """The --scan walk against real termios and a real oscillator: every
        probe rate must open a port (the off-nominal rates exist only through
        termios2), and one probe must answer — the nominal on a healthy board,
        a neighbor on a drifted one. The rig injects the one reset per probe
        the operator supplies in the field; this is the rate physics the
        simulator cannot arbitrate (a pty carries bytes at any rate), pinned
        on silicon."""
        module = pbrig.load_pureboot(self.rig.d.pureboot)
        found = None
        try:
            for pct in module.scan_ratios():
                rate = module.scan_rate(self.rig.d.baud, pct)
                self.rig.reset()
                try:
                    # Same wrap as identity(): on a shared line an undiscarded
                    # echo answers every rate a scan probes, so the walk would
                    # report the first one it tried.
                    port = self.rig.open_port(rate)
                except module.Error as error:
                    self.check("scan opens every probe rate", False, f"{rate} Bd: {error}")
                    return
                try:
                    module.Loader(port).connect(min(self.rig.d.wait, 6.0))
                    found = pct
                    break
                except module.Error:
                    continue
                finally:
                    port.close()
        except Exception as error:  # noqa: BLE001 — a rig hiccup is a result
            self.check("scan walks the probe ladder", False, str(error)[:70])
            return
        self.check("scan finds the board's rate", found is not None,
                   "no probe answered" if found is None else f"{found:+d} % of {self.rig.d.baud} Bd")

    def eeprom(self, info) -> None:
        size = info.eeprom_size
        if not size:
            print("  skip  EEPROM (this part has none)")
            return
        # A pattern no erase or partial write could produce by accident.
        pattern = bytes((i * 7 + 3) & 0xFF for i in range(size))
        image = self.work / "ee.bin"
        image.write_bytes(pattern)

        rc, out = self.rig.pureboot("--eeprom", str(image), "--verify-eeprom", str(image))
        self.check(f"EEPROM write + verify ({size} B)", rc == 0, self._brief(out))

        back = self.work / "ee-back.bin"
        rc, out = self.rig.pureboot("--read-eeprom", str(back))
        got = back.read_bytes() if back.exists() else b""
        self.check("EEPROM reads back what was written", got == pattern, f"{len(got)} B")

        self.rig.pureboot("--erase-eeprom")
        erased = self.work / "ee-erased.bin"
        self.rig.pureboot("--read-eeprom", str(erased))
        got = erased.read_bytes() if erased.exists() else b""
        self.check("EEPROM erase leaves 0xff", got == b"\xff" * size, f"{len(got)} B")

    def application(self, info, app: pathlib.Path, marker: str,
                    marker_wait: float = 2.5) -> None:
        rc, out = self.rig.pureboot("--flash", str(app), "--verify-flash", str(app))
        self.check(f"application flash + verify ({app.name})", rc == 0, self._brief(out))

        if marker:
            # The tool hands over as it ends its session, so the application is
            # already running — but only on a board whose DTR is unwired, where
            # opening a port simply listens. Where DTR *is* wired to reset (an
            # Arduino, most USB-serial dev boards), this open resets the part
            # and the activation window comes first, so a marker emitted once at
            # startup happens on the far side of a wait this cannot know the
            # length of: the window is a compile-time constant and nothing on
            # the wire reports it. Hence --marker-wait, and a fixture that
            # repeats its banner (PUREBOOT_HEARTBEAT) rather than saying it once.
            data = self.rig.capture(seconds=marker_wait)
            seen = marker.encode() in data
            sample = "".join(chr(b) if 32 <= b < 127 else "." for b in data[:40])
            self.check(f"application runs (emits {marker!r})", seen,
                       f"|{sample}|" if seen or data else
                       f"nothing in {marker_wait:g} s — if this board resets when its port "
                       f"opens, that wait has to outlast the activation window")

        back = self.work / "app-back.bin"
        rc, out = self.rig.pureboot("--read-flash", str(back))
        got = back.read_bytes() if back.exists() else b""
        self.check("application flash reads back", rc == 0 and len(got) == info.base,
                   f"{len(got)} B of {info.base}")

    def _witness(self, info, slot_length: int):
        """Read back the erased region and the loader slot: (erased, slot, how).

        Prefers ISP, because an independent reader is the only one that can
        testify about a loader just asked to erase around itself. Where no
        programmer is attached the link answers instead — which is weaker for
        exactly the reason it is worth having, a destroyed loader being unable
        to report anything at all. The two are never printed under one word:
        an absent probe is a fact about the bench, a wrong byte is a verdict on
        the loader, and a check that conflates them stops being read.
        """
        limit = info.base - 2 if info.patch_vector else info.base
        whole = self.work / "whole.bin"
        if self.rig.read_memory("flash", whole, "r"):
            image = whole.read_bytes()
            image += b"\xff" * (info.flash_size - len(image))
            return image[0:limit], image[info.base:info.base + slot_length], "ISP"

        module = pbrig.load_pureboot(self.rig.d.pureboot)
        port = self.rig.open_port()
        try:
            loader = module.Loader(port)
            if self.rig.d.autobaud:
                loader.connect_autobaud(self.rig.d.wait)
            else:
                loader.connect(self.rig.d.wait)
            return (loader.read_flash(0, limit),
                    loader.read_flash(info.base, slot_length),
                    "the link, no probe attached — the loader's own account")
        except Exception as error:  # noqa: BLE001 — a dead link is a result
            print(f"  skip  slot checks: no programmer, and the link did not "
                  f"answer either ({str(error)[:60]})")
            return None, None, ""
        finally:
            try:
                port.close()
            except Exception:  # noqa: BLE001
                pass

    def erase_and_slot(self, info, loader_image: pathlib.Path | None) -> None:
        rc, out = self.rig.pureboot("--erase-flash")
        self.check("application region erases", rc == 0, self._brief(out))

        want = loader_image.read_bytes() if loader_image and loader_image.exists() else b""
        erased, slot, how = self._witness(info, len(want))
        if erased is None:
            return

        # Erased application flash, up to the trampoline word the host composes
        # on a patched-vector part.
        limit = info.base - 2 if info.patch_vector else info.base
        self.check("erased application region is 0xff",
                   set(erased) <= {0xFF}, f"0x0000..{limit:#06x} via {how}")

        if want:
            self.check("loader slot survives the erase", slot == want,
                       f"{len(want)} B at {info.base:#06x} via {how}")
        else:
            print("  skip  loader slot comparison (pass --loader <image.bin>)")

    def seal(self, info, rounds: int = 1) -> None:
        """The seal, adversarially, over the real link.

        pureboot 9 has no running-slot guard: what stops a mangled command from
        erasing the loader is the seal and nothing else. So this aims the worst
        command the protocol has — an SPM erase at the loader's own first page —
        and damages one header byte at a time. Every one must come back NAK with
        the slot untouched and the session still in step.

        On a board whose link drops or mangles bytes of its own accord this is
        also the stress test: `--seal-rounds` repeats it, and a link fault
        during a round is indistinguishable to the loader from the damage being
        injected, which is the point.
        """
        module = pbrig.load_pureboot(self.rig.d.pureboot)
        self.rig.reset()
        port = self.rig.open_port()
        try:
            loader = module.Loader(port)
            if self.rig.d.autobaud:
                loader.connect_autobaud(self.rig.d.wait)
            else:
                loader.connect(self.rig.d.wait)
            if loader.info.version < module.SEALED_LOADER:
                print(f"  skip  seal checks (loader is pureboot {loader.info.version})")
                return

            head_of = lambda dmg: self._sealed(module, module.OP_WRITE, module.SP_SPM,
                                               info.base, module.SPM_ERASE, dmg)
            refused = 0
            attempts = 0
            for _ in range(rounds):
                for index in range(6):
                    attempts += 1
                    port.write(head_of((index, 0x01)))
                    verdict = port.read_exact(1, 5.0)
                    if verdict != module.NAK:
                        self.check(f"damaged byte {index} refused", False,
                                   f"verdict {verdict.hex()}")
                        return
                    if port.read_exact(1, 5.0) != module.PROMPT:
                        self.check(f"re-prompt after byte {index}", False, "no prompt")
                        return
                    refused += 1
            self.check("damaged headers refused", refused == attempts,
                       f"{refused}/{attempts}, every header byte")
            self.check("session still in step", loader.identity().raw == info.raw)

            # And the slot itself, read back over the link: the loader is the
            # thing that would have been erased, so its own account of its
            # first bytes is a real witness — an erased page reads all 0xff.
            head = loader.read_flash(info.base, 16)
            self.check("loader slot intact", set(head) != {0xFF}, head[:8].hex())
        except Exception as error:  # noqa: BLE001 — a dead link is a result
            self.check("seal checks", False, str(error)[:70])
        finally:
            try:
                port.close()
            except Exception:  # noqa: BLE001
                pass

    @staticmethod
    def _sealed(module, op, space, address, count, damage=None):
        """A sealed header, damaged after sealing — the shape a link fault has."""
        head = bytearray((op, module.selector(space, address), address & 0xFF,
                          (address >> 8) & 0xFF, count & 0xFF))
        seal = module.SEAL
        for byte in head:
            seal ^= byte
        out = bytearray(head + bytes((seal,)))
        if damage:
            out[damage[0]] ^= damage[1]
        return bytes(out)

    def refusals(self, info) -> None:
        # One word too many: a patched-vector part spends the slot's last word
        # on the trampoline, so its application stops two bytes short.
        limit = info.base - 2 if info.patch_vector else info.base
        oversized = self.work / "oversized.bin"
        oversized.write_bytes(bytes(limit + 2))
        rc, out = self.rig.pureboot("--flash", str(oversized))
        self.check(f"image over {limit} B refused", rc != 0, self._brief(out))

    # ------------------------------------------------------------------- run

    def run(self, app: pathlib.Path | None, loader_image: pathlib.Path | None,
            marker: str, marker_wait: float = 2.5, seal_rounds: int = 1) -> int:
        print("identity")
        info = self.identity()
        if info is None:
            print("\nthe loader never answered; nothing below can be trusted")
            return 1

        if not self.rig.d.autobaud:
            print("\nscan")
            self.scan()

        print("\nEEPROM")
        self.eeprom(info)

        if app:
            print("\napplication")
            self.application(info, app, marker, marker_wait)
        else:
            print("\nskip  application checks (pass --app <image.hex>)")

        print("\nerase and the slot boundary")
        self.erase_and_slot(info, loader_image)

        print("\nthe seal")
        self.seal(info, seal_rounds)

        print("\nrefusals")
        self.refusals(info)

        passed = sum(1 for _, ok, _ in self.results if ok)
        print(f"\n{passed}/{len(self.results)} passed")
        return 0 if passed == len(self.results) else 1


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="hardware acceptance suite for one pureboot deployment",
        epilog="overwrites the board's application flash and EEPROM — back them up first")
    pbrig.Deployment.add_arguments(parser)
    parser.add_argument("--app", type=pathlib.Path,
                        help="application image to flash (test/pbapp.cpp built for this deployment)")
    parser.add_argument("--loader", type=pathlib.Path,
                        help="the resident loader's .bin, to prove the slot survives an erase")
    parser.add_argument("--seal-rounds", type=int, default=1,
                        help="repeat the adversarial seal sweep N times (a lossy board's stress test)")
    parser.add_argument("--marker", default="",
                        help="text the application emits when it runs, e.g. APP")
    parser.add_argument("--marker-wait", type=float, default=2.5,
                        help="seconds to listen for it. On a board whose DTR is wired to "
                             "reset, opening the port resets the part, so this must outlast "
                             "the activation window (default 2.5)")
    args = parser.parse_args(argv)

    rig = pbrig.Rig(pbrig.Deployment.from_args(args))
    print(f"rig: {args.part} on {args.programmer}, link {args.port} at {args.baud} Bd"
          f"{' (autobaud)' if args.autobaud else ''}")
    print("this overwrites the application flash and EEPROM\n")
    with tempfile.TemporaryDirectory(prefix="pbhw-") as temporary:
        return Suite(rig, pathlib.Path(temporary)).run(args.app, args.loader, args.marker,
                                                      args.marker_wait, args.seal_rounds)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except pbrig.Error as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)