#!/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, activation-timeout
configuration, 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, writing page 0 last so an interrupted
flash still falls through to the loader.

Python standard library only; the serial port is driven with termios, so any
tty works — a USB adapter as well as a simavr pty.
"""

import argparse
import os
import select
import sys
import termios
import time

PROMPT = b"+"
PROTOCOL_VERSION = 1


class Error(Exception):
    pass


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


class Port:
    """A raw serial port with deadline-based reads."""

    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


# -------------------------------------------------------------- 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.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 + 512

    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."""

    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 run_application(self):
        self.port.write(b"G")
        self._expect_prompt()


# ---------------------------------------------------------------- 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 0 must go last —
    callers get it separated."""
    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, skip_blank):
    """Pages in programming order: ascending, page 0 last; optionally
    dropping all-0xff pages (sound only over erased flash) — never the
    load-bearing page 0."""
    rest = [a for a in sorted(pages) if a != 0]
    if skip_blank:
        rest = [a for a in rest if pages[a].count(0xFF) != len(pages[a])]
    return rest + [0]


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


def op_erase_flash(loader):
    """0xff over the whole application area. Order does not matter here —
    every target byte is the same value — and a blank word 0 still falls
    through to the loader, so an interruption is harmless."""
    blank = bytes([0xFF] * loader.info.page)
    for address in range(0, loader.info.base, loader.info.page):
        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):
    image = load_image(path)
    pages = plan_flash(image, loader.info)
    if erase:
        op_erase_flash(loader)
    order = covered(pages, 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_timeout(loader, seconds):
    loader.write_eeprom(loader.info.eeprom_size - 1, bytes((seconds,)))
    label = f"{seconds} s" if seconds else "the device default"
    print(f"activation timeout: {label}")


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}")


# -------------------------------------------------------------------- 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 (or 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("--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("--timeout", type=int, metavar="S", help="activation window, 1-254 s (0: default)")
    parser.add_argument("--stay", action="store_true", help="leave the loader in its session")
    args = parser.parse_args()

    if args.timeout is not None and not 0 <= args.timeout <= 254:
        parser.error("--timeout must be 0..254 (255 is the erased cell)")

    port = Port(args.port, args.baud)
    try:
        loader = Loader(port)
        info = loader.connect(args.wait)
        if args.info:
            print(f"device: {info.describe()}")
        if args.fuses:
            op_fuses(loader)
        if args.flash:
            op_flash(loader, args.flash, args.erase_flash, not args.no_verify)
        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.timeout is not None:
            op_timeout(loader, args.timeout)
        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)