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

Two positions. Address 0, a raw .bin handed to a programmer: the staging
install and the word-0 redirect run from copies outside page 0's slot, so the
running-slot guard never blocks them. And the staging slot itself, where a
loader already sitting there IS the staging copy — recognized by its embedded
block and left in place, then streaming the new resident like any staged copy.

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 - pb.SLOT
    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()