#!/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 erase_and_guard(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))

        # The slot must be untouched by an application erase, which only an
        # independent read can show — so this one goes over ISP, not the link.
        whole = self.work / "whole.bin"
        if not self.rig.read_memory("flash", whole, "r"):
            self.check("loader slot survives the erase", False, "ISP read failed")
            return
        image = whole.read_bytes()
        image += b"\xff" * (info.flash_size - len(image))

        # 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(image[0:limit]) <= {0xFF}, f"0x0000..{limit:#06x}")

        if loader_image and loader_image.exists():
            want = loader_image.read_bytes()
            got = image[info.base:info.base + len(want)]
            self.check("loader slot survives the erase", got == want,
                       f"{len(want)} B at {info.base:#06x}")
        else:
            print("  skip  loader slot comparison (pass --loader <image.bin>)")

    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) -> 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 write guard")
        self.erase_and_guard(info, loader_image)

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


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