#!/usr/bin/env python3
"""Self-update across a link change: the host must follow the staging copy.

`--update-loader` installs the new image in the staging slot and then *enters
it* to have it rewrite the resident. That copy is the new image, so it speaks the
new image's baud and backend — but the host was talking to the *resident*. Where
the two differ, the host kept knocking at the old rate in the old mode, the
staging copy never answered, and the update stranded: staging installed, resident
untouched, and on a 1 KiB tiny the application region (which *is* the staging
slot there) already gone.

The wire cannot be probed for this — 512 bytes of position-independent code carry
no header saying what rate they were built for — so the operator declares it, and
a mismatch with nothing declared has to say so instead of reporting a bare
timeout.

Stdlib only, no device: host-tool logic, so it runs on every chip's preset beside
pureboot.planner.
"""
import importlib.util
import pathlib

PB = pathlib.Path(__file__).resolve().parents[1] / "pureboot" / "pureboot.py"
_spec = importlib.util.spec_from_file_location("pureboot", PB)
pb = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(pb)

IDENTITY = b"\x05\x1e\x95\x0f"  # pureboot 5 + m328p signature
P = F = 0


def check(name, ok, detail=""):
    global P, F
    P, F = P + (1 if ok else 0), F + (0 if ok else 1)
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))


class TwoLinkPort:
    """A board whose resident and staging copy answer on different links.

    Only the rate currently set decides who can be heard, which is the physical
    truth: a loader's bit timing is a cycle count, so a copy built for another
    rate is unreadable until the host retunes. The knock bytes carry the mode, so
    a backend mismatch is caught the same way.
    """

    def __init__(self, resident=(57600, False), staged=(38400, False)):
        self.resident, self.staged = resident, staged
        self.baud = resident[0]
        self.entered = False        # a 'J' has handed control to the staging copy
        self.switches = []          # every retune the host asked for
        self.pending = bytearray()  # what the device has queued to send

    # --- the part under test needs this to exist at all
    def set_baud(self, baud):
        self.baud = baud
        self.switches.append(baud)

    def flush_input(self):
        self.pending.clear()

    def _audible(self, knock=None):
        baud, autobaud = self.staged if self.entered else self.resident
        if self.baud != baud:
            return False
        if knock is None:
            return True
        return knock == (bytes((pb.CALIBRATE, ord("p"))) if autobaud else b"pb")

    def write(self, data):
        data = bytes(data)
        if data[:1] == b"J" and len(data) == 3:
            # The resident acks the jump, then control moves to the copy.
            if self._audible():
                self.pending += pb.PROMPT
            self.entered = True
        elif data in (b"pb", bytes((pb.CALIBRATE, ord("p")))):
            if self._audible(data):
                self.pending += pb.PROMPT
        elif data == b"b":
            if self._audible():
                self.pending += IDENTITY + pb.PROMPT

    def read_available(self, wait):
        out, self.pending = bytes(self.pending), bytearray()
        return out

    def read_exact(self, count, timeout):
        if len(self.pending) < count:
            raise pb.Error(f"timeout: got {len(self.pending)} of {count} bytes")
        out, self.pending = bytes(self.pending[:count]), self.pending[count:]
        return out


def connected(port):
    """A Loader already in session with the resident."""
    loader = pb.Loader(port)
    loader.connect(2.0)
    return loader


def main():
    # The control first: where the staged image keeps the resident's link, the
    # flow works and needs no retune. This is the case that always passed, and
    # it is what made the bug look like "self-update is broken" rather than
    # "self-update cannot change the link".
    port = TwoLinkPort(resident=(57600, False), staged=(57600, False))
    loader = connected(port)
    try:
        loader.enter_copy(0x7C00, 2.0)
        check("same link: staging copy entered", True)
    except pb.Error as error:
        check("same link: staging copy entered", False, str(error))

    # A baud change, declared. The host must retune before knocking.
    port = TwoLinkPort(resident=(57600, False), staged=(38400, False))
    loader = connected(port)
    try:
        loader.enter_copy(0x7C00, 2.0, link=(38400, False))
        check("baud change declared: entered after retuning", 38400 in port.switches,
              f"switches={port.switches}")
    except (pb.Error, TypeError) as error:
        check("baud change declared: entered after retuning", False, repr(error))

    # A backend change, declared: the knock itself has to become the calibration
    # pulse, or an autobaud staging copy never hears a thing.
    port = TwoLinkPort(resident=(57600, False), staged=(57600, True))
    loader = connected(port)
    try:
        loader.enter_copy(0x7C00, 2.0, link=(57600, True))
        check("backend change declared: entered as autobaud", True)
    except (pb.Error, TypeError) as error:
        check("backend change declared: entered as autobaud", False, repr(error))

    # Nothing declared against a changed link: it still cannot work, but the
    # error has to name the cause. A bare "no answer" sent the operator looking
    # at the wiring while the application region sat erased.
    port = TwoLinkPort(resident=(57600, False), staged=(38400, False))
    loader = connected(port)
    try:
        loader.enter_copy(0x7C00, 0.3)
        check("undeclared mismatch: reported", False, "unexpectedly succeeded")
    except pb.Error as error:
        text = str(error).lower()
        check("undeclared mismatch: error names the link, not just a timeout",
              "link" in text or "baud" in text or "backend" in text, str(error))
    except TypeError as error:
        check("undeclared mismatch: error names the link, not just a timeout",
              False, repr(error))

    # The resident's own link must be restored for the caller: a declared
    # staging link is for the copy, and the tool talks to the new resident after.
    port = TwoLinkPort(resident=(57600, False), staged=(38400, False))
    loader = connected(port)
    try:
        loader.enter_copy(0x7C00, 2.0, link=(38400, False))
        check("session records the link it is now speaking", loader.baud == 38400,
              f"loader.baud={getattr(loader, 'baud', None)}")
    except (pb.Error, TypeError, AttributeError) as error:
        check("session records the link it is now speaking", False, repr(error))

    print(f"\n  {P} passed, {F} failed")
    return 1 if F else 0


if __name__ == "__main__":
    raise SystemExit(main())