#!/usr/bin/env python3
"""pureboot host tool — the smart half of the pureboot protocol (README.md).

The device exposes primitives; this tool composes them: image loading (raw
binary or Intel HEX), flash programming with read-back verification, erase as
writing 0xff, EEPROM programming, fuse and info readout, the hand-over jump,
and — on chips without a hardware boot section — the reset-vector surgery
that re-homes the application's entry through the trampoline word below the
loader. Page 0 and the trampoline are written first, so every interruption
point of a flash leaves the chip reset-recoverable into the loader.

It also updates the loader itself (--update-loader): pureboot's image is
position-independent, so the tool installs the identical binary one 512-byte
slot below the resident loader, jumps into that staging copy, lets it rewrite
the resident slot, and restores what the staging slot held — resumable at
every phase from the flash state plus a host-side state file carrying the
saved bytes.

Python standard library only; the serial port is driven with termios on POSIX
and the Win32 serial API (through ctypes) on Windows, so any tty or COM port
works — a USB adapter as well as a simavr pty.
"""

import argparse
import json
import os
import sys
import time

if os.name == "nt":
    import ctypes
    from ctypes import wintypes
else:
    import select
    import termios

PROMPT = b"+"
PROTOCOL_VERSION = 1
SLOT = 512  # the loader slot size; also the self-update staging distance


class Error(Exception):
    pass


# ---------------------------------------------------------------- serial ---


class PosixPort:
    """A raw serial port with deadline-based reads, over termios."""

    def __init__(self, path, baud):
        self.fd = os.open(path, os.O_RDWR | os.O_NOCTTY)
        attrs = termios.tcgetattr(self.fd)
        attrs[0] = 0  # iflag
        attrs[1] = 0  # oflag
        attrs[2] = termios.CREAD | termios.CLOCAL | termios.CS8  # cflag
        attrs[3] = 0  # lflag
        try:
            speed = getattr(termios, f"B{baud}")
        except AttributeError:
            raise Error(f"unsupported baud rate {baud}") from None
        attrs[4] = attrs[5] = speed
        attrs[6][termios.VMIN] = 0
        attrs[6][termios.VTIME] = 0
        termios.tcsetattr(self.fd, termios.TCSANOW, attrs)

    def close(self):
        os.close(self.fd)

    def write(self, data):
        os.write(self.fd, data)

    def flush_input(self):
        termios.tcflush(self.fd, termios.TCIFLUSH)

    def read_available(self, wait):
        """Everything that arrives within `wait` seconds of quiet start."""
        ready, _, _ = select.select([self.fd], [], [], wait)
        return os.read(self.fd, 4096) if ready else b""

    def read_exact(self, count, timeout):
        data = b""
        deadline = time.monotonic() + timeout
        while len(data) < count:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise Error(f"timeout: got {len(data)} of {count} bytes")
            ready, _, _ = select.select([self.fd], [], [], remaining)
            if ready:
                data += os.read(self.fd, count - len(data))
        return data


if os.name == "nt":
    # The same port, over the Win32 serial API — kernel32 through ctypes, so
    # the tool stays standard-library only. Timeouts live in the driver
    # (COMMTIMEOUTS) rather than in a readiness call: Windows has no select()
    # for a COM handle, so each read asks the driver for its own deadline.

    _GENERIC_READ, _GENERIC_WRITE = 0x80000000, 0x40000000
    _OPEN_EXISTING, _PURGE_RXCLEAR = 3, 0x0008
    _INVALID_HANDLE = wintypes.HANDLE(-1).value
    # A gap this long ends a read_available(): longer than the coalescing a
    # USB-serial adapter's latency timer imposes (16 ms on FTDI parts), so a
    # burst is not split, short enough to stay responsive.
    _GAP_MS = 30

    class _DCB(ctypes.Structure):
        _fields_ = [
            ("DCBlength", wintypes.DWORD),
            ("BaudRate", wintypes.DWORD),
            ("fBits", wintypes.DWORD),  # the packed flag bitfield, set below
            ("wReserved", wintypes.WORD),
            ("XonLim", wintypes.WORD),
            ("XoffLim", wintypes.WORD),
            ("ByteSize", wintypes.BYTE),
            ("Parity", wintypes.BYTE),
            ("StopBits", wintypes.BYTE),
            ("XonChar", ctypes.c_char),
            ("XoffChar", ctypes.c_char),
            ("ErrorChar", ctypes.c_char),
            ("EofChar", ctypes.c_char),
            ("EvtChar", ctypes.c_char),
            ("wReserved1", wintypes.WORD),
        ]

    class _COMMTIMEOUTS(ctypes.Structure):
        _fields_ = [
            ("ReadIntervalTimeout", wintypes.DWORD),
            ("ReadTotalTimeoutMultiplier", wintypes.DWORD),
            ("ReadTotalTimeoutConstant", wintypes.DWORD),
            ("WriteTotalTimeoutMultiplier", wintypes.DWORD),
            ("WriteTotalTimeoutConstant", wintypes.DWORD),
        ]

    _k32 = ctypes.WinDLL("kernel32", use_last_error=True)
    _LPDWORD = ctypes.POINTER(wintypes.DWORD)
    # Declared, not inferred: a HANDLE is a pointer, and a defaulted int
    # return would truncate it on 64-bit.
    _k32.CreateFileW.restype = wintypes.HANDLE
    _k32.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD,
                                 wintypes.LPVOID, wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE]
    _k32.ReadFile.argtypes = [wintypes.HANDLE, wintypes.LPVOID, wintypes.DWORD, _LPDWORD, wintypes.LPVOID]
    _k32.WriteFile.argtypes = [wintypes.HANDLE, wintypes.LPCVOID, wintypes.DWORD, _LPDWORD, wintypes.LPVOID]
    _k32.GetCommState.argtypes = [wintypes.HANDLE, ctypes.POINTER(_DCB)]
    _k32.SetCommState.argtypes = [wintypes.HANDLE, ctypes.POINTER(_DCB)]
    _k32.SetCommTimeouts.argtypes = [wintypes.HANDLE, ctypes.POINTER(_COMMTIMEOUTS)]
    _k32.PurgeComm.argtypes = [wintypes.HANDLE, wintypes.DWORD]
    _k32.CloseHandle.argtypes = [wintypes.HANDLE]

    def _fail(what):
        code = ctypes.get_last_error()
        raise Error(f"{what}: {ctypes.FormatError(code).strip()} (Windows error {code})")

    class WindowsPort:
        """A raw serial port with deadline-based reads, over Win32."""

        def __init__(self, path, baud):
            # Win32 takes the rate as a plain integer, so unlike termios any
            # rate the hardware can divide down to is available — but a driver
            # may also accept one it cannot produce (an FT232R takes a baud of
            # 3, reports it back, and goes on using the previous divisor).
            # Only obvious nonsense is refusable; the rest is the driver's word.
            if baud < 50:
                raise Error(f"unsupported baud rate {baud}")
            # \\.\COM6: the device-namespace form. A bare COMn resolves only
            # for n < 10, and double-digit ports are routine on Windows.
            if path.lower().startswith("com") and path[3:].isdigit():
                path = rf"\\.\{path}"
            self.handle = _k32.CreateFileW(
                path, _GENERIC_READ | _GENERIC_WRITE, 0, None, _OPEN_EXISTING, 0, None
            )
            if self.handle == _INVALID_HANDLE:
                _fail(f"cannot open {path}")
            self.timeouts = None
            try:
                dcb = _DCB()
                dcb.DCBlength = ctypes.sizeof(_DCB)
                if not _k32.GetCommState(self.handle, ctypes.byref(dcb)):
                    _fail(f"cannot read the state of {path}")
                dcb.BaudRate, dcb.ByteSize, dcb.Parity, dcb.StopBits = baud, 8, 0, 0  # 8N1
                # fBinary, and DTR/RTS asserted (fDtrControl and fRtsControl,
                # two bits each, = _ENABLE); every other flag clear, so no
                # parity and no flow control. Raising both matches what opening
                # a POSIX tty does — including the reset pulse on the boards
                # that wire DTR to it.
                dcb.fBits = 0x1 | (1 << 4) | (1 << 12)
                if not _k32.SetCommState(self.handle, ctypes.byref(dcb)):
                    _fail(f"cannot configure {path} for {baud} baud 8N1")
                # Arm them once here too: reads re-arm per call, but the write
                # timeout would otherwise stay at the driver's default — which
                # may be "wait forever" — until the first read.
                self._deadline(_GAP_MS, 1000)
            except Error:
                # An open port outlives the exception otherwise, and a COM
                # handle is exclusive: the next attempt would meet its own
                # leftover as "Access is denied".
                self.close()
                raise

        def close(self):
            _k32.CloseHandle(self.handle)

        def _deadline(self, interval, total):
            """Arm the driver's read timeouts: `interval` ms of quiet ends a
            read once bytes have arrived, `total` ms ends it regardless."""
            if self.timeouts == (interval, total):
                return
            spec = _COMMTIMEOUTS()
            spec.ReadIntervalTimeout = interval
            spec.ReadTotalTimeoutConstant = total
            spec.WriteTotalTimeoutConstant = 5000
            if not _k32.SetCommTimeouts(self.handle, ctypes.byref(spec)):
                _fail("cannot set the port timeouts")
            self.timeouts = (interval, total)

        def _read(self, count):
            buffer = ctypes.create_string_buffer(count)
            got = wintypes.DWORD()
            if not _k32.ReadFile(self.handle, buffer, count, ctypes.byref(got), None):
                _fail("read failed")
            return buffer.raw[: got.value]

        def write(self, data):
            written = wintypes.DWORD()
            if not _k32.WriteFile(self.handle, data, len(data), ctypes.byref(written), None):
                _fail("write failed")
            if written.value != len(data):
                raise Error(f"short write: {written.value} of {len(data)} bytes")

        def flush_input(self):
            if not _k32.PurgeComm(self.handle, _PURGE_RXCLEAR):
                _fail("cannot flush the input buffer")

        def read_available(self, wait):
            """Everything that arrives within `wait` seconds of quiet start."""
            # A zero total means *no* timeout to the driver, so never round
            # down to it — the same trap on the deadline below.
            self._deadline(_GAP_MS, max(1, round(wait * 1000)))
            return self._read(4096)

        def read_exact(self, count, timeout):
            data = b""
            deadline = time.monotonic() + timeout
            while len(data) < count:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise Error(f"timeout: got {len(data)} of {count} bytes")
                # No interval timeout here: only the count or the deadline
                # ends the read, so a gap mid-reply is simply waited out.
                self._deadline(0, max(1, round(remaining * 1000)))
                data += self._read(count - len(data))
            return data


Port = WindowsPort if os.name == "nt" else PosixPort


# -------------------------------------------------------------- protocol ---


class Info:
    """The 12-byte info block."""

    def __init__(self, raw):
        if len(raw) != 12 or raw[0:2] != b"PB":
            raise Error(f"bad info block: {raw.hex()}")
        if raw[2] != PROTOCOL_VERSION:
            raise Error(f"protocol version {raw[2]}, tool speaks {PROTOCOL_VERSION}")
        self.raw = bytes(raw)
        self.signature = raw[3:6]
        self.page = raw[6]
        self.base = raw[7] | (raw[8] << 8)
        self.eeprom_size = raw[9] | (raw[10] << 8)
        self.patch_vector = bool(raw[11] & 1)
        self.flash_size = self.base + SLOT
        self.stage = self.base - SLOT  # where a staging copy of the loader goes
        # The hand-over target, as the word address 'J' takes: the trampoline
        # below the loader (tinies), or word 0 (mega — the application's own
        # reset vector; BOOTRST re-vectors a reset into the loader instead).
        self.app_entry_word = (self.base - 2) // 2 if self.patch_vector else 0

    def describe(self):
        sig = " ".join(f"{b:02x}" for b in self.signature)
        vector = "host-patched reset vector" if self.patch_vector else "hardware boot section"
        return (
            f"signature {sig}, page {self.page} B, "
            f"app flash {self.base} B (loader at {self.base:#06x}), "
            f"EEPROM {self.eeprom_size} B, {vector}"
        )


class Loader:
    """A pureboot session. Between commands the loader has prompted `+` and
    awaits a command byte; every method restores that invariant — except
    jump(), after which the target must be knocked afresh."""

    def __init__(self, port):
        self.port = port
        self.info = None

    def connect(self, wait):
        """Knock until the activation window answers, then read the info
        block. Also converges when the loader already sits in its command
        loop: the knock bytes are ignored-or-executed there, and the drain
        absorbs whatever they produced."""
        self.port.flush_input()
        deadline = time.monotonic() + wait
        while True:
            self.port.write(b"pb")
            if PROMPT in self.port.read_available(0.4):
                break
            if time.monotonic() > deadline:
                raise Error("no answer — reset the device within its activation window")
        while self.port.read_available(0.3):
            pass
        self.port.write(b"b")
        self.info = Info(self.port.read_exact(12, 2.0))
        self._expect_prompt()
        return self.info

    def _expect_prompt(self, timeout=2.0):
        byte = self.port.read_exact(1, timeout)
        if byte != PROMPT:
            raise Error(f"expected prompt, got {byte.hex()}")

    def _command(self, tx, reply_len=0, timeout=2.0):
        self.port.write(tx)
        reply = self.port.read_exact(reply_len, timeout) if reply_len else b""
        self._expect_prompt(timeout)
        return reply

    def _stream_read(self, command, address, count):
        data = b""
        while count:
            chunk = min(count, 256)
            head = bytes((ord(command), address & 0xFF, address >> 8, chunk & 0xFF))
            data += self._command(head, chunk, 5.0)
            address += chunk
            count -= chunk
        return data

    def read_flash(self, address, count):
        return self._stream_read("R", address, count)

    def read_eeprom(self, address, count):
        return self._stream_read("r", address, count)

    def write_page(self, address, data):
        assert len(data) == self.info.page and address % self.info.page == 0
        head = bytes((ord("W"), address & 0xFF, address >> 8))
        self._command(head + data, 0, 2.0)

    def write_eeprom(self, address, data):
        offset = 0
        while offset < len(data):
            chunk = data[offset : offset + 256]
            head = bytes((ord("w"), address & 0xFF, address >> 8, len(chunk) & 0xFF))
            self.port.write(head)
            for byte in chunk:
                self.port.write(bytes((byte,)))
                self._expect_prompt()  # per-byte ack: the write has begun
            self._expect_prompt()  # the next command prompt
            address += len(chunk)
            offset += len(chunk)

    def read_fuses(self):
        return self._command(b"F", 4, 2.0)

    def jump(self, word_address):
        """'J': the device acks, then execution continues at the word
        address — a loader slot's base (whose copy must then be knocked
        afresh) or the application entry."""
        self.port.write(bytes((ord("J"), word_address & 0xFF, word_address >> 8)))
        self._expect_prompt()

    def enter_copy(self, byte_address, wait):
        """Jump into the loader copy at `byte_address` and knock it. Ending
        up in the copy addressed is guaranteed by construction: a jump to a
        slot base lands in that slot's entry stub."""
        self.jump(byte_address // 2)
        return self.connect(wait)

    def run_application(self):
        self.jump(self.info.app_entry_word)


# ---------------------------------------------------------------- images ---


def load_image(path):
    """Raw binary, or Intel HEX by extension (.hex/.ihx/.ihex)."""
    data = open(path, "rb").read()
    if not path.lower().endswith((".hex", ".ihx", ".ihex")):
        if not data:
            raise Error(f"{path}: empty image")
        return data
    memory = {}
    for number, line in enumerate(data.decode("ascii", "replace").splitlines(), 1):
        line = line.strip()
        if not line:
            continue
        if not line.startswith(":"):
            raise Error(f"{path}:{number}: not an Intel HEX record")
        record = bytes.fromhex(line[1:])
        if sum(record) & 0xFF:
            raise Error(f"{path}:{number}: checksum mismatch")
        count, address, kind = record[0], (record[1] << 8) | record[2], record[3]
        payload = record[4 : 4 + count]
        if kind == 0:
            for i, byte in enumerate(payload):
                memory[address + i] = byte
        elif kind == 1:
            break
        elif kind in (2, 4) and not any(payload):
            continue  # a zero base extends nothing
        elif kind in (3, 5):
            continue  # start address: irrelevant, reset is the entry
        else:
            raise Error(f"{path}:{number}: record type {kind} reaches beyond the 16-bit space")
    if not memory:
        raise Error(f"{path}: empty image")
    return bytes(memory.get(i, 0xFF) for i in range(max(memory) + 1))


# --------------------------------------------------------------- surgery ---


def rjmp_target(word_address, opcode, flash_words):
    return (word_address + 1 + (opcode & 0x0FFF)) % flash_words


def rjmp_to(word_address, destination, flash_words):
    return 0xC000 | ((destination - word_address - 1) % flash_words % 0x1000)


def plan_flash(image, info):
    """The pages to program, as {page_address: bytes}, already carrying the
    reset-vector surgery where the chip needs it."""
    page = info.page
    limit = info.base - (2 if info.patch_vector else 0)
    if len(image) > limit:
        raise Error(f"image is {len(image)} B, application flash ends at {limit}")
    final = bytearray(image) + bytearray([0xFF] * (-len(image) % page))

    if info.patch_vector:
        flash_words = info.flash_size // 2
        word0 = final[0] | (final[1] << 8)
        if word0 & 0xF000 != 0xC000:
            raise Error(
                "the image's reset vector is not an rjmp — pureboot's vector "
                "surgery cannot re-home it (crt-less entry at address 0?)"
            )
        entry = rjmp_target(0, word0, flash_words)
        if entry >= info.base // 2:
            raise Error(
                "the image's reset vector already targets the loader — this "
                "is a read-back of a patched image; flash the original"
            )
        trampoline_word = (info.base - 2) // 2
        patch = rjmp_to(0, info.base // 2, flash_words)
        final[0], final[1] = patch & 0xFF, patch >> 8
        trampoline_page = info.base - page
        if len(final) < trampoline_page + page:
            final += bytearray([0xFF] * (trampoline_page + page - len(final)))
        jump = rjmp_to(trampoline_word, entry, flash_words)
        final[info.base - 2], final[info.base - 1] = jump & 0xFF, jump >> 8

    pages = {a: bytes(final[a : a + page]) for a in range(0, len(final), page)}
    return pages


def covered(pages, info, skip_blank):
    """Pages in programming order; optionally dropping all-0xff pages (sound
    only over erased flash) — never a load-bearing one.

    With a patched vector (tinies), the patched page 0 goes first and the
    trampoline page second: from the first write on, a reset lands in the
    loader and the loader's own fall-through lands on the application entry,
    so every interruption point of the flash is recoverable. With a hardware
    boot section a reset re-vectors to the loader regardless; ascending
    order, page 0 last, maximizes what an interrupted image retains."""
    trampoline_page = info.base - info.page if info.patch_vector else None
    first = [0, trampoline_page] if info.patch_vector else []
    rest = [a for a in sorted(pages) if a not in first]
    if skip_blank:
        rest = [a for a in rest if pages[a].count(0xFF) != len(pages[a])]
    order = [a for a in first if a in pages] + rest
    if not info.patch_vector:
        order = [a for a in order if a != 0] + ([0] if 0 in pages else [])
    return order


# ----------------------------------------------------------------- fuses ---


def mega_boot(high_fuse):
    """Decode the ATmega328P high fuse's boot configuration (DS40002061B
    §27.3, Table 27-13/27-16): BOOTSZ1:0 in bits 2:1 select the boot-section
    words, BOOTRST in bit 0 (programmed = 0) re-vectors reset to its start.
    Returns (bootrst_programmed, boot_section_start_byte)."""
    bootsz = (high_fuse >> 1) & 0x03
    words = {0b11: 256, 0b10: 512, 0b01: 1024, 0b00: 2048}[bootsz]
    return (high_fuse & 1) == 0, 0x8000 - words * 2


# ---------------------------------------------------------- loader update ---


def image_info(image):
    """The info block embedded in a pureboot binary, or None."""
    at = image.find(b"PB" + bytes((PROTOCOL_VERSION,)))
    return Info(image[at : at + 12]) if 0 <= at <= len(image) - 12 else None


def staging_content(image, info):
    """The 512-byte staging-slot content: the image, padding, and — on
    chips whose hand-over jumps through the word below the resident loader —
    that word, which for a staging copy is the slot's own last word: an rjmp
    to the resident base. The staging copy's fall-through and 'J'-free exit
    both land in a loader instead of garbage."""
    if len(image) > (SLOT - 2 if info.patch_vector else SLOT):
        raise Error(f"loader image is {len(image)} B, the slot holds {SLOT - 2 if info.patch_vector else SLOT}")
    content = bytearray(image) + bytearray([0xFF] * (SLOT - len(image)))
    if info.patch_vector:
        through = rjmp_to((info.base - 2) // 2, info.base // 2, info.flash_size // 2)
        content[SLOT - 2], content[SLOT - 1] = through & 0xFF, through >> 8
    return bytes(content)


def update_preflight(image, info, fuse_bytes):
    """Errors and warnings before any flash is touched. Returns warnings."""
    embedded = image_info(image)
    if embedded is None:
        raise Error("no pureboot info block in the update image — not a pureboot binary?")
    if embedded.raw[3:] != info.raw[3:]:
        raise Error(
            f"update image is for another target: it declares "
            f"[{embedded.describe()}], the device says [{info.describe()}]"
        )
    warnings = []
    if not info.patch_vector:
        if fuse_bytes is None:
            raise Error("a loader update on this chip needs its fuses — unreadable? pass --assume-fuses")
        high = fuse_bytes[3]
        bootrst, bls_start = mega_boot(high)
        if info.stage < bls_start:
            raise Error(
                f"cannot self-update: the staging slot {info.stage:#06x} lies below the "
                f"boot section ({bls_start:#06x}, high fuse {high:#04x}) where SPM is disabled "
                f"— a boot section of at least 1 KB (BOOTSZ) is required, and only an "
                f"external programmer can change fuses"
            )
        if not bootrst:
            warnings.append(
                "BOOTRST unprogrammed: reset boots the application throughout the update; "
                "an interruption is recovered by re-running this update"
            )
        elif bls_start == info.stage:
            warnings.append(
                "BOOTRST targets the staging slot: brief unrecoverable windows exist while "
                "the staging copy itself is being installed or retired (page-write scale)"
            )
        else:
            warnings.append(
                f"BOOTRST targets {bls_start:#06x}, inside application flash: reset reaches a "
                f"loader only across erased flash from there"
            )
    return warnings


class UpdateState:
    """The host-side memory of an update in flight: what the staging slot
    held (and page 0, where the update repoints it). Losing this file after
    the staging slot was overwritten loses those saved bytes — the update
    still completes, but the staging region can then only be restored by
    reflashing the application."""

    def __init__(self, path):
        self.path = path
        self.data = None

    def load_or_save(self, loader):
        info = loader.info
        if os.path.exists(self.path):
            self.data = json.load(open(self.path))
            if bytes.fromhex(self.data["signature"]) != info.signature or self.data["base"] != info.base:
                raise Error(f"{self.path} belongs to a different device — remove it to start over")
            return
        self.data = {
            "signature": info.signature.hex(),
            "base": info.base,
            "staging": loader.read_flash(info.stage, SLOT).hex(),
            "page0": loader.read_flash(0, info.page).hex() if info.patch_vector else "",
        }
        with open(self.path, "w") as f:
            json.dump(self.data, f)

    @property
    def staging(self):
        return bytes.fromhex(self.data["staging"])

    @property
    def page0(self):
        return bytes.fromhex(self.data["page0"])

    def discard(self):
        os.unlink(self.path)


def write_differing(loader, base, content, order=None):
    """Program the pages of `content` at `base` that differ from flash —
    idempotent, so a resumed phase redoes only what an interruption left."""
    page = loader.info.page
    offsets = order if order is not None else range(0, len(content), page)
    written = 0
    for offset in offsets:
        want = content[offset : offset + page]
        if loader.read_flash(base + offset, page) != want:
            loader.write_page(base + offset, want)
            written += 1
    for at in range(0, len(content), 256):
        if loader.read_flash(base + at, min(256, len(content) - at)) != content[at : at + 256]:
            raise Error(f"verify failed at {base + at:#06x} after programming")
    return written


def patch_word0(loader, page0, target_base):
    """Rewrite page 0 with its word 0 re-aimed at `target_base` — the
    resume insurance around rewriting a loader slot the reset path uses."""
    info = loader.info
    patched = bytearray(page0)
    word = rjmp_to(0, target_base // 2, info.flash_size // 2)
    patched[0], patched[1] = word & 0xFF, word >> 8
    write_differing(loader, 0, bytes(patched))
    return bytes(patched)


def op_update_loader(loader, wait, path, state_path, fuse_bytes):
    """Replace the resident loader with `path`, using the loader itself as
    its own staging loader. Every phase is idempotent and keyed off the
    actual flash state, so a re-run after any interruption resumes; the
    state file carries the bytes the staging slot held."""
    info = loader.info
    image = load_image(path)
    for warning in update_preflight(image, info, fuse_bytes):
        print(f"note: {warning}")
    staged = staging_content(image, info)
    resident = bytes(image) + bytes([0xFF] * (SLOT - len(image)))
    page = info.page

    state = UpdateState(state_path)
    state.load_or_save(loader)

    # Install the staging copy. On a chip whose staging slot starts at
    # address 0 (the 1 KB tiny13A), its first page carries the reset vector:
    # written last, so any earlier interruption still resets into the old
    # resident, and from then on resets enter the staging copy.
    order = list(range(0, SLOT, page))
    if info.stage == 0:
        order = order[1:] + [0]
    if write_differing(loader, info.stage, staged, order):
        print(f"staging copy installed at {info.stage:#06x}")

    # Enter it and let it rewrite the resident slot. Where a patched reset
    # vector routes through the resident (a tiny with the staging slot away
    # from page 0), word 0 is re-aimed at the staging copy around the
    # rewrite, so a power failure mid-rewrite still resets into a loader.
    loader.enter_copy(info.stage, wait)
    redirect = info.patch_vector and info.stage != 0
    if redirect:
        patch_word0(loader, state.page0, info.stage)
    if write_differing(loader, info.base, resident):
        print(f"resident loader rewritten at {info.base:#06x}")

    # Enter the new resident and put the staging region back: page 0 first
    # where it lives in that region (word 0 then points at the new resident
    # for the rest of the restore), the saved trampoline with the rest.
    loader.enter_copy(info.base, wait)
    if redirect:
        write_differing(loader, 0, state.page0)
    order = list(range(0, SLOT, page))
    if info.stage == 0:
        order = [0] + order[1:]
    write_differing(loader, info.stage, state.staging, order)

    state.discard()
    print(f"loader updated: {len(image)} B at {info.base:#06x}, staging region restored")


def check_walk_region(pages, info, fuse_bytes, force):
    """With BOOTRST programmed but targeting below the loader, reset reaches
    the loader only by walking across erased flash from the boot-section
    start; application data in that span would divert reset into itself.
    Only checkable when the fuses are known (--fuses or --assume-fuses)."""
    if info.patch_vector or fuse_bytes is None:
        return
    bootrst, bls_start = mega_boot(fuse_bytes[3])
    if not bootrst or bls_start >= info.base:
        return
    overlap = [a for a in sorted(pages) if a >= bls_start and pages[a].count(0xFF) != len(pages[a])]
    if overlap and not force:
        raise Error(
            f"the image writes {overlap[0]:#06x}.. inside the reset walk region "
            f"[{bls_start:#06x}, {info.base:#06x}) (BOOTRST programmed): reset could no "
            f"longer reach the loader — --force to flash it anyway"
        )


# ------------------------------------------------------------ operations ---


def op_erase_flash(loader):
    """0xff over the whole application area. Descending on a patched-vector
    chip: page 0 — the patched reset vector — goes last, so an interrupted
    erase still resets into the loader, and once it is gone the whole area
    is erased and the reset walk reaches the loader anyway."""
    blank = bytes([0xFF] * loader.info.page)
    addresses = range(0, loader.info.base, loader.info.page)
    for address in reversed(addresses) if loader.info.patch_vector else addresses:
        loader.write_page(address, blank)
    print(f"erase: {loader.info.base // loader.info.page} pages")


def op_erase_eeprom(loader):
    loader.write_eeprom(0, bytes([0xFF] * loader.info.eeprom_size))
    print(f"erase: {loader.info.eeprom_size} B of EEPROM")


def op_flash(loader, path, erase, verify, fuse_bytes=None, force=False):
    image = load_image(path)
    pages = plan_flash(image, loader.info)
    check_walk_region(pages, loader.info, fuse_bytes, force)
    if erase:
        op_erase_flash(loader)
    order = covered(pages, loader.info, skip_blank=erase)
    for address in order:
        loader.write_page(address, pages[address])
    print(f"flash: {path}: {len(order)} pages")
    if verify:
        verify_pages(loader, pages)


def verify_pages(loader, pages):
    for address in sorted(pages):
        got = loader.read_flash(address, loader.info.page)
        if got != pages[address]:
            first = next(i for i in range(len(got)) if got[i] != pages[address][i])
            raise Error(
                f"verify failed at {address + first:#06x}: "
                f"wrote {pages[address][first]:02x}, read {got[first]:02x}"
            )
    print(f"verify: {len(pages)} pages ok")


def op_verify_flash(loader, path):
    verify_pages(loader, plan_flash(load_image(path), loader.info))


def op_read_flash(loader, path):
    data = loader.read_flash(0, loader.info.base)
    open(path, "wb").write(data)
    print(f"read flash: {len(data)} B -> {path}")


def op_eeprom(loader, path, erase, verify):
    image = load_image(path)
    if len(image) > loader.info.eeprom_size:
        raise Error(f"EEPROM image is {len(image)} B, device has {loader.info.eeprom_size}")
    if erase:
        op_erase_eeprom(loader)
    loader.write_eeprom(0, image)
    print(f"eeprom: {path}: {len(image)} B")
    if verify:
        got = loader.read_eeprom(0, len(image))
        if got != image:
            first = next(i for i in range(len(got)) if got[i] != image[i])
            raise Error(f"verify failed at EEPROM {first:#06x}: wrote {image[first]:02x}, read {got[first]:02x}")
        print(f"verify: {len(image)} B ok")


def op_verify_eeprom(loader, path):
    image = load_image(path)
    got = loader.read_eeprom(0, len(image))
    if got != image:
        first = next(i for i in range(len(got)) if got[i] != image[i])
        raise Error(f"verify failed at EEPROM {first:#06x}: expected {image[first]:02x}, read {got[first]:02x}")
    print(f"verify: {len(image)} B of EEPROM ok")


def op_read_eeprom(loader, path):
    data = loader.read_eeprom(0, loader.info.eeprom_size)
    open(path, "wb").write(data)
    print(f"read EEPROM: {len(data)} B -> {path}")


def op_fuses(loader):
    low, lock, extended, high = loader.read_fuses()
    print(f"fuses: low {low:02x} high {high:02x} extended {extended:02x} lock {lock:02x}")
    return bytes((low, lock, extended, high))


# -------------------------------------------------------------------- cli ---


def main():
    parser = argparse.ArgumentParser(
        description="pureboot host tool", epilog="operations run in the order listed above"
    )
    parser.add_argument("--port", required=True, help="serial device: COM6, /dev/ttyUSB0, or a simavr pty")
    parser.add_argument("--baud", type=int, default=115200, help="115200 mega, 57600 tinies")
    parser.add_argument("--wait", type=float, default=30.0, help="seconds to keep knocking")
    parser.add_argument("--info", action="store_true", help="print the device info block")
    parser.add_argument("--fuses", action="store_true", help="read the fuse and lock bytes")
    parser.add_argument("--update-loader", metavar="FILE", help="replace the loader with this pureboot binary")
    parser.add_argument("--state", metavar="FILE", help="update state file (default: FILE.pbstate)")
    parser.add_argument("--assume-fuses", metavar="HEX8", help="fuse bytes low,lock,ext,high as 8 hex digits "
                        "(overrides reading them — e.g. under a simulator that cannot)")
    parser.add_argument("--erase-flash", action="store_true", help="0xff over the application flash")
    parser.add_argument("--flash", metavar="FILE", help="program an application (bin or ihex)")
    parser.add_argument("--no-verify", action="store_true", help="skip read-back after writes")
    parser.add_argument("--read-flash", metavar="FILE", help="dump the application flash")
    parser.add_argument("--verify-flash", metavar="FILE", help="compare flash against an image")
    parser.add_argument("--erase-eeprom", action="store_true", help="0xff over the EEPROM")
    parser.add_argument("--eeprom", metavar="FILE", help="program the EEPROM (bin or ihex)")
    parser.add_argument("--read-eeprom", metavar="FILE", help="dump the EEPROM")
    parser.add_argument("--verify-eeprom", metavar="FILE", help="compare EEPROM against an image")
    parser.add_argument("--force", action="store_true", help="override refusable safety checks")
    parser.add_argument("--stay", action="store_true", help="leave the loader in its session")
    args = parser.parse_args()

    if args.update_loader and (args.flash or args.erase_flash):
        parser.error("--update-loader does not combine with application flash operations")
    fuse_override = None
    if args.assume_fuses:
        try:
            fuse_override = bytes.fromhex(args.assume_fuses)
            assert len(fuse_override) == 4
        except (ValueError, AssertionError):
            parser.error("--assume-fuses takes 8 hex digits: low,lock,extended,high")

    port = Port(args.port, args.baud)
    try:
        loader = Loader(port)
        info = loader.connect(args.wait)
        if args.info:
            print(f"device: {info.describe()}")
        fuse_bytes = fuse_override
        if args.fuses or (args.update_loader and not info.patch_vector and fuse_bytes is None):
            read = op_fuses(loader)
            if fuse_bytes is None:
                fuse_bytes = read
        if args.update_loader:
            state = args.state or args.update_loader + ".pbstate"
            op_update_loader(loader, args.wait, args.update_loader, state, fuse_bytes)
        if args.flash:
            op_flash(loader, args.flash, args.erase_flash, not args.no_verify, fuse_bytes, args.force)
        elif args.erase_flash:
            op_erase_flash(loader)
        if args.read_flash:
            op_read_flash(loader, args.read_flash)
        if args.verify_flash:
            op_verify_flash(loader, args.verify_flash)
        if args.eeprom:
            op_eeprom(loader, args.eeprom, args.erase_eeprom, not args.no_verify)
        elif args.erase_eeprom:
            op_erase_eeprom(loader)
        if args.read_eeprom:
            op_read_eeprom(loader, args.read_eeprom)
        if args.verify_eeprom:
            op_verify_eeprom(loader, args.verify_eeprom)
        if args.stay:
            print("loader stays in its session (reset to leave)")
        else:
            loader.run_application()
            print("application running")
    finally:
        port.close()


if __name__ == "__main__":
    try:
        main()
    except Error as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)