#!/usr/bin/env python3
"""Position-independence acceptance test: the identical binary, flashed one
slot below the resident, must serve the complete command set from there. The
info block must come back byte-identical, the write guard must refuse the
staged copy's own slot and permit the resident's, and the staged copy must be
able to rewrite the resident verbatim.

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

import os
import subprocess
import sys


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


def main():
    device_bin, elf, mcu, hz, base_hex, page, baud, tool, workdir = sys.argv[1:]
    base, page, baud = int(base_hex, 0), int(page), int(baud)
    stage = None  # derived from the device's own info (slot-sized) below
    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")
    image_path = os.path.join(workdir, "pureboot.bin")
    subprocess.run([objcopy, "-O", "binary", elf, image_path], check=True)
    image = open(image_path, "rb").read()

    device = pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, os.path.join(workdir, "dump.bin"))
    try:
        port = pb.Port(device.pty, baud)
        loader = pb.Loader(port)
        info = loader.connect(25)
        if info.base != base:
            fail(f"info reports base {info.base:#06x}")
        resident_info = info.raw

        # Install the staging copy exactly as the update flow would.
        stage = info.stage
        staged = pb.staging_content(image, info)
        pb.write_differing(loader, stage, staged)

        # Enter it; from here on, every command runs in the relocated copy.
        staged_info = loader.enter_copy(stage, 25)
        if staged_info.raw != resident_info:
            fail(f"staged info {staged_info.raw.hex()} != resident info {resident_info.hex()}")

        # 'R' from the staged copy already proved itself in the install
        # verify; 'F' must answer 4 bytes (values are unmodeled in simavr).
        if len(loader.read_fuses()) != 4:
            fail("fuse read from the staged copy")

        # EEPROM round-trip through the staged copy.
        pattern = bytes(range(0x50, 0x60))
        loader.write_eeprom(0, pattern)
        if loader.read_eeprom(0, len(pattern)) != pattern:
            fail("EEPROM round-trip through the staged copy")

        # The guard, both ways: its own slot refused (drained, unchanged), the
        # resident slot writable. The refusal leaves its drained words in the
        # SPM buffer, so the write that follows may take them — and clears
        # them by writing, so the retry must not.
        before = loader.read_flash(stage, page)
        loader.write_page(stage, bytes(page))
        if loader.read_flash(stage, page) != before:
            fail("the staged copy's guard let its own slot change")
        marker = bytes((i * 3) & 0xFF for i in range(page))
        loader.write_page(base, marker)
        if loader.read_flash(base, page) != marker:
            loader.write_page(base, marker)
            if loader.read_flash(base, page) != marker:
                fail("the staged copy could not write the resident slot, even on retry")

        # Restore the resident image through the staged copy, then 'J' back
        # into it and prove it lives.
        resident = image + b"\xff" * (pb.SLOT - len(image))
        pb.write_differing(loader, base, resident)
        back_info = loader.enter_copy(base, 25)
        if back_info.raw != resident_info:
            fail("the restored resident does not serve its info block")
        port.close()
    finally:
        device.stop()
    print("pbreloc: the relocated copy serves the full command set")


if __name__ == "__main__":
    main()