#!/usr/bin/env python3
"""Self-update end-to-end: an application is flashed, then the loader
replaces itself with a re-timed build through the host tool's
--update-loader — and the power-fail phases of that update are rehearsed by
killing the simulated device mid-write, restarting it from its flash dump,
and letting a re-run complete the update.

The mega runs the BOOTRST-unprogrammed profile (reset boots the application;
the fixture application's 'L' jump is the application-owned loader entry),
with --assume-fuses standing in for the fuse read simavr cannot model. The
tinies reset into a loader at every phase by construction — the t13a because
its staging slot carries the reset vector itself, the t85 through the word-0
redirect the tool plants around the resident rewrite.

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


MEGA_FUSES = "ffffffdd"  # high 0xdd: BOOTSZ = 1 KB, BOOTRST unprogrammed


def make_fault_loader(pb, base, 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 - 512:
                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 == "atmega328p"
    reset_hex = "0" if mega else None  # the 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

    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 = bytes.fromhex(MEGA_FUSES) if mega 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" * (512 - len(image))

    def resident_bytes(loader):
        return loader.read_flash(base, 256) + loader.read_flash(base + 256, 256)

    def assert_state(loader, image, app_pages):
        if resident_bytes(loader) != padded(image):
            fail("resident loader does not match the update image")
        stage = base - 512
        got = loader.read_flash(stage, 256) + loader.read_flash(stage + 256, 256)
        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 mega:
            args += ["--assume-fuses", MEGA_FUSES]
        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, not mega),
            ("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, 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 tinies an independent decode of the reset routing.
    flash = open(dump, "rb").read()
    if flash[base : base + 512] != padded(images[final]):
        fail("ground-truth resident region does not match the final image")
    if not mega:
        flash_words = (base + 512) // 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()