#!/usr/bin/env python3
"""Re-homing acceptance test: a pureboot image programmed somewhere other
than its canonical top slot must still be a working loader —
position-independent, guarding its accidental slot — and the ordinary
--update-loader flow must put a build into the top slot from there.

Two positions are exercised. Address 0 (a raw .bin handed to a programmer,
which defaults to offset 0): the staging install and the word-0 redirect
both run from copies whose slots are not page 0's, so the running-slot
guard never blocks the flow. The staging slot itself: a loader already
sitting there IS the installed staging copy — the tool recognizes it by
its embedded info block and leaves it in place instead of tripping the
copy's own guard on the composed through-word — and that (older) copy
streams the new resident like any staged copy. In both cases flashing an
application through the healed resident overwrites the stale copy, vector
surgery included, and the banner proves the launch.

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

import os
import sys


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


def rehome_from(pbsim, pb, device_bin, elf, place_hex, guard_probe, update_bin, base, page, baud, app_bin, workdir,
                mcu, hz):
    """Place the loader at `place_hex`, heal through --update-loader, flash
    the application, expect the banner."""
    dump = os.path.join(workdir, f"dump-{place_hex}.bin")
    state = os.path.join(workdir, f"rehome-{place_hex}.pbstate")
    if os.path.exists(state):
        os.unlink(state)
    device = pbsim.Device(device_bin, elf, mcu, hz, place_hex, page, baud, dump, reset_hex="0")
    try:
        port = pb.Port(device.pty, baud)
        loader = pb.Loader(port)
        info = loader.connect(25)
        if info.base != base:
            fail(f"the misplaced copy reports base {info.base:#06x} — the info block must stay canonical")

        # The accidental slot still guards itself; re-homing rides on the
        # canonical slots being writable from it.
        probe = int(guard_probe, 0)
        before = loader.read_flash(probe, info.page)
        loader.write_page(probe, bytes(info.page))
        if loader.read_flash(probe, info.page) != before:
            fail("the misplaced copy's guard let its own slot change")

        # The ordinary update flow puts the build into the top slot.
        pb.op_update_loader(loader, 25, update_bin, state, None)
        update = open(update_bin, "rb").read()
        if loader.read_flash(base, len(update)) != update:
            fail("the canonical slot does not hold the update image")

        # An application flashed through the healed resident overwrites the
        # stale copy (surgery included) and launches.
        pages = pb.plan_flash(open(app_bin, "rb").read(), loader.info)
        for address in pb.covered(pages, loader.info, skip_blank=False):
            loader.write_page(address, pages[address])
        pb.verify_pages(loader, pages)
        loader.run_application()
        if port.read_exact(3, 5.0) != b"APP":
            fail(f"application does not banner after the re-home from {place_hex}")
        port.close()
    finally:
        device.stop()


def main():
    (device_bin, elf, update_bin, mcu, hz, base_hex, page, baud, app_bin, tool, workdir) = sys.argv[1:]
    base, page, baud = int(base_hex, 0), int(page), int(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)

    # Address 0: the raw-.bin-to-a-programmer accident. The guard probe is
    # the copy's own page 0.
    rehome_from(pbsim, pb, device_bin, elf, "0x0", "0x0", update_bin, base, page, baud, app_bin, workdir, mcu, hz)
    print("re-home from address 0: converged")

    # The staging slot: erased flash with the loader sitting exactly where
    # a staging copy would — the tool must leave it in place and let it
    # stream the (different) update build into the resident slot.
    stage = base - 512
    rehome_from(pbsim, pb, device_bin, elf, hex(stage), hex(stage), update_bin, base, page, baud, app_bin, workdir,
                mcu, hz)
    print("re-home from the staging slot: converged")

    print("pbrehome: a misplaced loader re-homes through the ordinary update flow")


if __name__ == "__main__":
    main()