#!/usr/bin/env python3
"""Self-update end-to-end: an application is flashed, the loader replaces
itself with a re-timed build, and every power-fail phase is rehearsed by
killing the device mid-write, restarting it from its flash dump, and letting
a re-run complete the update.

The boot-sectioned megas run the BOOTRST-unprogrammed profile — reset boots
the application, whose 'L' is the application-owned loader entry — with
--assume-fuses standing in for the fuse read simavr cannot model.

Usage: pbupdate.py <device_bin> <pureboot_elf> <update_elf> <mcu> <hz>
                   <base_hex> <page> <baud> <app_bin> <tool_py> <workdir>
"""

import os
import subprocess
import sys


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


def rjmp_decode(word, at, flash_words):
    """Written against the instruction-set definition, not with the tool's
    encoder, so an encoding bug cannot verify itself."""
    if word & 0xF000 != 0xC000:
        fail(f"word at {at * 2:#06x} is {word:#06x}, not an rjmp")
    offset = word & 0x0FFF
    if offset >= 0x800:
        offset -= 0x1000
    return (at + 1 + offset) % flash_words


class PowerFail(Exception):
    pass


def assumed_fuses(pb, image):
    """Synthetic 'F' bytes for --assume-fuses: the smallest boot section
    covering both the resident and the staging slot (two slots — what a
    self-update needs), BOOTRST unprogrammed — the per-chip BOOTSZ ladder
    and fuse byte come from the tool's own table, keyed by the update
    image's embedded signature."""
    info = pb.image_info(image)
    which, ladder = pb.BOOT_FUSE[bytes(info.signature[1:3])]
    bits = min((b for b in ladder if ladder[b] * 2 >= 2 * pb.SLOT), key=lambda b: ladder[b])
    fuses = bytearray((0xFF, 0xFF, 0xFF, 0xFF))
    fuses[which] = 0xF8 | (bits << 1) | 1
    return bytes(fuses)


def make_fault_loader(pb, base, slot, kill_region, kill_hits, device):
    """A Loader whose write_page kills the device (or, with device=None,
    just the host) at the Nth write into a region; the sequence
    stage->resident->stage distinguishes the install from the restore."""

    class FaultLoader(pb.Loader):
        def __init__(self, port):
            super().__init__(port)
            self.seen_resident = False
            self.hits = 0

        def write_page(self, address, data):
            if address >= base:
                phase = "resident"
                self.seen_resident = True
            elif address >= base - slot:
                phase = "stage_restore" if self.seen_resident else "stage"
            else:
                phase = "app"
            if phase == kill_region:
                self.hits += 1
                if self.hits == kill_hits:
                    if device is not None:
                        device.power_fail()
                    raise PowerFail(f"{kill_region} write {kill_hits}")
            super().write_page(address, data)

    return FaultLoader


def main():
    (device_bin, elf, update_elf, mcu, hz, base_hex, page, baud, app_bin, tool, workdir) = sys.argv[1:]
    base, page, baud = int(base_hex, 0), int(page), int(baud)
    mega = mcu.startswith("atmega")
    # The m48s are megas without a boot section: patched vector, no fuse
    # preflight, and the same reset-to-0 the tinies get.
    patch = not mega or mcu.startswith("atmega48")
    reset_hex = "0" if mega else None  # the boot-sectioned mega runs BOOTRST-unprogrammed here
    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

    slot = pb.SLOT
    os.makedirs(workdir, exist_ok=True)
    objcopy = os.environ.get("PB_OBJCOPY", "avr-objcopy")
    images = {}
    for name, source in (("v0", elf), ("v9", update_elf)):
        path = os.path.join(workdir, name + ".bin")
        subprocess.run([objcopy, "-O", "binary", source, path], check=True)
        images[name] = open(path, "rb").read()
    if images["v0"] == images["v9"]:
        fail("the update image is byte-identical to the resident build")
    dump = os.path.join(workdir, "dump.bin")
    state = os.path.join(workdir, "update.pbstate")
    fuses = assumed_fuses(pb, images["v0"]) if mega and not patch else None

    def connect(device):
        port = pb.Port(device.pty, baud)
        if mega:
            # Reset boots the application here; its 'L' is the loader entry.
            # To a live loader the same byte is an ignored command.
            port.read_available(0.5)
            port.write(b"L")
        loader = pb.Loader(port)
        loader.connect(25)
        return port, loader

    def padded(image):
        return image + b"\xff" * (slot - len(image))

    def resident_bytes(loader):
        return loader.read_flash(base, slot)

    def assert_state(loader, image, app_pages):
        if resident_bytes(loader) != padded(image):
            fail("resident loader does not match the update image")
        stage = base - slot
        got = loader.read_flash(stage, slot)
        for address, data in app_pages.items():
            if stage <= address < base:
                if got[address - stage : address - stage + page] != data:
                    fail(f"staging region page {address:#06x} not restored")

    device = pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, dump, reset_hex=reset_hex)
    final = "v0"
    try:
        # The application first — its planner output is the restore truth.
        pbsim.run_tool(tool, device.pty, baud, "--flash", app_bin, "--stay")
        port, loader = connect(device)
        app_pages = pb.plan_flash(open(app_bin, "rb").read(), loader.info)
        port.close()

        # A clean CLI update, resident -> v9.
        args = ["--update-loader", os.path.join(workdir, "v9.bin"), "--state", state, "--stay"]
        if fuses:
            args += ["--assume-fuses", fuses.hex()]
        out = pbsim.run_tool(tool, device.pty, baud, *args)
        if "loader updated" not in out:
            fail("update did not report success")
        if os.path.exists(state):
            fail("state file survived a completed update")
        port, loader = connect(device)
        assert_state(loader, images["v9"], app_pages)
        loader.run_application()
        if port.read_exact(3, 5.0) != b"APP":
            fail("application does not banner after the update")
        port.close()
        final = "v9"
        print("clean update: resident replaced, staging restored, application intact")

        # Power-fail rehearsal: kill mid-phase, restart from the dump,
        # re-run, and the update must still complete. Each round flips the
        # direction so the flash is never already at its target. The mega's
        # mid-resident-rewrite loss is exercised as a host crash instead:
        # with BOOTRST unprogrammed and the resident mid-erase, a power loss
        # there has no reset path into the staging copy — the documented
        # cost of that profile (README).
        for kill_region, kill_hits, kill_device in (
            ("stage", 2, True),
            ("resident", 1, patch),
            ("stage_restore", 2, True),
        ):
            device.reset()  # the previous round left the application running
            port, loader = connect(device)
            target = "v9" if resident_bytes(loader) == padded(images["v0"]) else "v0"
            image_path = os.path.join(workdir, target + ".bin")
            injected = make_fault_loader(pb, base, slot, kill_region, kill_hits, device if kill_device else None)(port)
            injected.info = loader.info
            try:
                pb.op_update_loader(injected, 25, image_path, state, fuses)
                fail(f"{kill_region}: fault never triggered")
            except PowerFail as event:
                print(f"power fail injected: {event}")
            port.close()
            if kill_device:
                device = pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, dump,
                                      reset_hex=reset_hex, resume=dump)
            port, loader = connect(device)
            pb.op_update_loader(loader, 25, image_path, state, fuses)
            assert_state(loader, images[target], app_pages)
            loader.run_application()
            if port.read_exact(3, 5.0) != b"APP":
                fail(f"{kill_region}: application lost after the resumed update")
            port.close()
            final = target
            print(f"resumed after {kill_region} loss: update completed, application intact")
    finally:
        device.stop()

    # Ground truth: the simulator's own flash against the final state, and
    # on the patched-vector chips an independent decode of the reset routing.
    flash = open(dump, "rb").read()
    if flash[base : base + slot] != padded(images[final]):
        fail("ground-truth resident region does not match the final image")
    if patch:
        flash_words = (base + slot) // 2
        word0 = flash[0] | (flash[1] << 8)
        if rjmp_decode(word0, 0, flash_words) != base // 2:
            fail("ground-truth reset vector does not land on the loader")
        app = open(app_bin, "rb").read()
        trampoline = flash[base - 2] | (flash[base - 1] << 8)
        if rjmp_decode(trampoline, (base - 2) // 2, flash_words) != rjmp_decode(app[0] | (app[1] << 8), 0, flash_words):
            fail("ground-truth trampoline does not land on the application entry")
    print("pbupdate: clean update + all power-fail phases recovered")


if __name__ == "__main__":
    main()