#!/usr/bin/env python3
"""End-to-end TSB protocol test: spawn the simavr device, speak the TinySafeBoot
wire protocol over its pty (as the real host tools do), and actually flash it.

Usage: tsbtest.py <device_binary> <tsb.elf> <boot_base_hex>
Exits 0 if every scenario passes.
"""
import subprocess
import sys
import time

import serial

CONFIRM = 0x21  # '!'
REQUEST = 0x3F  # '?'
PAGE = 128      # ATmega328P: 64 words


class Device:
    """The simavr runner, exposing UART0 as a pty."""

    def __init__(self, binary, elf, boot_base, dump="/tmp/tsb_dump.bin"):
        self.proc = subprocess.Popen(
            [binary, elf, boot_base, dump],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        self.dump = dump
        self.pty = None
        deadline = time.time() + 5
        while time.time() < deadline:
            line = self.proc.stdout.readline()
            if not line:
                break
            if line.startswith("TSB_PTY"):
                self.pty = line.split()[1]
                break
        if not self.pty:
            self.stop()
            raise RuntimeError("device did not report a pty")

    def stop(self):
        self.proc.terminate()
        try:
            self.proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.proc.kill()


class Host:
    """A faithful TSB host, per the wire protocol."""

    def __init__(self, pty):
        self.s = serial.Serial(pty, 115200, timeout=1.5)
        self.info = None

    def _read(self, n):
        data = self.s.read(n)
        if len(data) != n:
            raise AssertionError(f"expected {n} bytes, got {len(data)}: {data.hex()}")
        return data

    def activate(self):
        self.s.reset_input_buffer()
        self.s.write(b"@@@")
        reply = self._read(17)
        if reply[16] != CONFIRM:
            raise AssertionError(f"activation reply not '!'-terminated: {reply.hex()}")
        self.info = reply[:16]
        return self.info

    # Parsed info-block fields (host math from the spec).
    @property
    def pagesize(self):
        return self.info[9] * 2

    @property
    def appflash(self):
        return (self.info[10] | (self.info[11] << 8)) * 2

    @property
    def eeprom_size(self):
        return (self.info[12] | (self.info[13] << 8)) + 1

    def _expect(self, byte, what):
        r = self._read(1)
        if r[0] != byte:
            raise AssertionError(f"{what}: expected {byte:#x}, got {r.hex()}")

    # Host-paced page read ('f'/'e'): send '!', take a page, repeat; stop with
    # anything else, then the Mainloop '!'.
    def _read_pages(self, cmd, npages):
        self.s.write(cmd.encode())
        data = b""
        for _ in range(npages):
            self.s.write(bytes([CONFIRM]))
            data += self._read(PAGE)
        self.s.write(bytes([REQUEST]))  # stop
        self._expect(CONFIRM, f"{cmd} end")
        return data

    # Device-paced page write ('F'/'E'): device offers '?', host sends '!'+page,
    # or anything else to stop.
    def _write_pages(self, cmd, data):
        if len(data) % PAGE:
            data += b"\xff" * (PAGE - len(data) % PAGE)
        self.s.write(cmd.encode())
        for off in range(0, len(data), PAGE):
            self._expect(REQUEST, f"{cmd} '?'")
            self.s.write(bytes([CONFIRM]) + data[off:off + PAGE])
        self._expect(REQUEST, f"{cmd} trailing '?'")
        self.s.write(bytes([REQUEST]))  # stop
        self._expect(CONFIRM, f"{cmd} end")

    def write_flash(self, data):
        self._write_pages("F", data)

    def read_flash(self, npages):
        return self._read_pages("f", npages)

    def write_eeprom(self, data):
        self._write_pages("E", data)

    def read_eeprom(self, npages):
        return self._read_pages("e", npages)

    def read_config(self):
        self.s.write(b"c")
        page = self._read(PAGE)
        self._expect(CONFIRM, "c end")
        return page

    def write_config(self, data):
        assert len(data) == PAGE
        self.s.write(b"C")
        self._expect(REQUEST, "C '?'")
        self.s.write(bytes([CONFIRM]) + data)
        echo = self._read(PAGE)  # device echoes what it programmed
        self._expect(CONFIRM, "C end")
        return echo


def check(cond, msg):
    if not cond:
        raise AssertionError(msg)
    print(f"  ok: {msg}")


def main():
    binary, elf, boot_base = sys.argv[1], sys.argv[2], sys.argv[3]
    dev = Device(binary, elf, boot_base)
    failures = []
    try:
        host = Host(dev.pty)

        # --- activation ---
        info = host.activate()
        check(info[0:3] == b"TSB", f"magic 'TSB' (got {info[0:3]!r})")
        check(info[6:9] == bytes([0x1E, 0x95, 0x0F]), f"signature 1E 95 0F (got {info[6:9].hex()})")
        check(info[14] == info[15], f"device-type bytes 14==15 (got {info[14]:#x},{info[15]:#x})")
        check(host.pagesize == PAGE, f"page size {PAGE} (got {host.pagesize})")
        check(host.eeprom_size == 1024, f"eeprom size 1024 (got {host.eeprom_size})")
        print(f"  info: {info.hex()}  appflash={host.appflash} eeprom={host.eeprom_size}")

        # --- flash write / read round-trip (actual self-programming) ---
        app = bytes(range(256))  # two pages of known data
        host.write_flash(app)
        back = host.read_flash(2)
        check(back == app, f"flash round-trip 2 pages ({'match' if back == app else 'MISMATCH'})")

        # --- EEPROM write / read round-trip ---
        edata = bytes((i * 7) & 0xFF for i in range(PAGE))
        host.write_eeprom(edata)
        eback = host.read_eeprom(1)
        check(eback == edata, "eeprom round-trip 1 page")

        # --- config page write / read round-trip ---
        cfg = bytes([0x00, 0x00, 0x40]) + b"\xff" * (PAGE - 3)  # appjump 0, timeout 0x40, no password
        echo = host.write_config(cfg)
        check(echo == cfg, "config write echoes the programmed page")
        check(host.read_config() == cfg, "config read-back matches")

    except AssertionError as e:
        failures.append(str(e))
        print(f"  FAIL: {e}")
    finally:
        dev.stop()

    if failures:
        print(f"FAILED ({len(failures)})")
        return 1
    print("ALL PASS")
    return 0


if __name__ == "__main__":
    sys.exit(main())