#!/usr/bin/env python3
"""End-to-end autobaud test: drive an autobaud loader in simavr through the
calibration handshake and a flash + EEPROM + fuse round-trip, cross-checked
against the simulator's ground-truth memory — then repeat at a second F_CPU with
the *same* loader binary, which is the property autobaud exists for: one
clock-agnostic image that locks onto whatever rate the host sends.

Usage: pbautobaud.py <device_bin> <loader_elf> <mcu> <base_hex> <page>
                     <app_bin> <app_hz> <app_baud> <tool_py> <workdir>

The loader is a software-serial build on PB0/PB1 (pureboot_add_autobaud's
default), so the runner drives it over the GPIO⇄pty bridge (-l sw:B0,B1). The
app fixture is built for (app_hz, app_baud); the hand-over is checked at that
point, and a second point at half the clock proves the lock is measured, not
baked in.
"""

import os
import sys


def fail(message):
    print(f"FAIL: {message}")
    sys.exit(1)


def main():
    (device_bin, elf, mcu, base_hex, page, app_bin, app_hz, app_baud, tool, workdir) = sys.argv[1:]
    base, page, app_hz, app_baud = int(base_hex, 0), int(page), int(app_hz), int(app_baud)
    sys.path.insert(0, os.path.dirname(os.path.abspath(tool)))
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    import pbsim
    import pureboot as pb

    os.makedirs(workdir, exist_ok=True)
    ee_image = bytes(range(0xA0, 0xB0))
    ee_path = os.path.join(workdir, "ee.bin")
    open(ee_path, "wb").write(ee_image)

    # The geometry the surgery planner needs, from the chip class the runner is
    # told — the same derivation pbtest.py makes: the boot-sectioned megas need
    # no vector surgery, the tinies and the boot-section-less m48s do, and the
    # large chips speak word addresses.
    mega = mcu.startswith("atmega")
    patch = not mega or mcu.startswith("atmega48")
    word_flash = base + pb.SLOT > 0x10000
    wire_base = base // 2 if word_flash else base
    flags = (1 if patch else 0) | (2 if word_flash else 0)
    ground_truth = pb.Info(bytes([ord("P"), ord("B"), pb.NEWEST_LOADER, 0, 0, 0, page & 0xFF,
                                  wire_base & 0xFF, wire_base >> 8, 0, 0, flags]))

    def round_trip(hz, baud, label, hand_over):
        """One clock point: reset, calibrate + knock, program, verify against the
        simulator's own flash, and (at the app's point) hand over to the fixture."""
        dump = os.path.join(workdir, f"flash_{label}.bin")
        device = pbsim.Device(device_bin, elf, mcu, str(hz), base_hex, page, baud, dump, link="sw:B0,B1")
        try:
            # The host tool, in autobaud mode, sends the 0xC0 calibration pulse
            # and a single knock at `baud`; the loader locks to it.
            out = pbsim.run_tool(tool, device.pty, baud, "--autobaud", "--info", "--fuses",
                                 "--flash", app_bin, "--eeprom", ee_path, "--stay")
            for needed in ("version", "signature", "fuses", "verify:", "stays"):
                if needed not in out:
                    fail(f"{label}: session output lacks {needed!r}\n{out}")
            # Read both memories back over the locked link and check them.
            read_flash = os.path.join(workdir, f"rf_{label}.bin")
            read_eeprom = os.path.join(workdir, f"re_{label}.bin")
            out = pbsim.run_tool(tool, device.pty, baud, "--autobaud", "--verify-flash", app_bin,
                                 "--verify-eeprom", ee_path, "--read-flash", read_flash,
                                 "--read-eeprom", read_eeprom, "--stay")
            if out.count("verify:") != 2:
                fail(f"{label}: did not verify both memories\n{out}")
            if open(read_eeprom, "rb").read()[: len(ee_image)] != ee_image:
                fail(f"{label}: EEPROM read-back mismatch")

            if hand_over:
                device.reset()
                port = pb.Port(device.pty, baud)
                try:
                    loader = pb.Loader(port)
                    live = loader.connect_autobaud(15)
                    if live.version != pb.NEWEST_LOADER:
                        fail(f"{label}: loader reports pureboot {live.version}")
                    loader.run_application()
                    banner = port.read_exact(3, 5.0)
                    if banner != b"APP":
                        fail(f"{label}: application banner was {banner!r}")
                finally:
                    port.close()
        finally:
            device.stop()

        # Ground truth (read after the runner exits and writes its dump): what
        # the tool programmed must be what the simulator actually holds.
        pages = pb.plan_flash(open(app_bin, "rb").read(), ground_truth)
        flash_true = open(dump, "rb").read()
        for address, data in pages.items():
            if flash_true[address : address + page] != data:
                fail(f"{label}: simulator flash differs from the programmed image at {address:#06x}")
        print(f"  {label}: locked at {hz} Hz / {baud} Bd, flash+EEPROM verified"
              + (", hand-over ok" if hand_over else ""))

    # The app fixture is built for one clock; the hand-over banners there. A
    # second point at double that clock, same loader binary, proves the lock is
    # measured, not baked in — the whole point of autobaud. (Doubling keeps the
    # bit period healthy; halving would drop it below the software UART's floor.)
    round_trip(app_hz, app_baud, "clock-a", hand_over=True)
    round_trip(app_hz * 2, app_baud, "clock-b", hand_over=False)
    print("pbautobaud: calibration lock and flash/EEPROM/fuse round-trip pass at both clocks")


if __name__ == "__main__":
    main()