#!/usr/bin/env python3
import argparse
import sys
import time
from pathlib import Path

import serial


FRAME_START = 0x55
CMD_INFO = ord("I")
CMD_READ = ord("R")
CMD_WRITE_PAGE = ord("W")

REPLY_OK = ord("K")
REPLY_DATA = ord("D")
REPLY_ERROR = ord("E")

MAX_SIZE = 32768
DEFAULT_PAGE_SIZE = 64


class ProtocolError(RuntimeError):
    pass


def read_exact(port: serial.Serial, length: int) -> bytes:
    data = bytearray()
    while len(data) < length:
      chunk = port.read(length - len(data))
      if not chunk:
        raise ProtocolError(f"timeout while reading {length} bytes")
      data.extend(chunk)
    return bytes(data)


def expect_frame_start(port: serial.Serial) -> int:
    start = read_exact(port, 1)[0]
    if start != FRAME_START:
        raise ProtocolError(f"bad frame start: 0x{start:02x}")
    return read_exact(port, 1)[0]


def send_command(port: serial.Serial, command: int, payload: bytes = b"") -> int:
    port.write(bytes([FRAME_START, command]) + payload)
    port.flush()
    return expect_frame_start(port)


def query_info(port: serial.Serial) -> tuple[int, int]:
    reply = send_command(port, CMD_INFO)
    if reply != REPLY_DATA:
        raise ProtocolError(f"unexpected INFO reply: {reply!r}")
    length = read_exact(port, 1)[0]
    payload = read_exact(port, length)
    if length != 4:
        raise ProtocolError(f"unexpected INFO payload length: {length}")
    page_size = payload[0]
    total_size = payload[1] | (payload[2] << 8)
    return page_size, total_size


def read_chunk(port: serial.Serial, address: int, length: int) -> bytes:
    payload = bytes([address & 0xFF, (address >> 8) & 0xFF, length])
    reply = send_command(port, CMD_READ, payload)
    if reply == REPLY_ERROR:
        code = read_exact(port, 1)[0]
        raise ProtocolError(f"READ failed with error code {code}")
    if reply != REPLY_DATA:
        raise ProtocolError(f"unexpected READ reply: {reply!r}")
    returned_length = read_exact(port, 1)[0]
    if returned_length != length:
        raise ProtocolError(f"short READ reply: {returned_length} != {length}")
    return read_exact(port, returned_length)


def write_page(port: serial.Serial, address: int, data: bytes) -> None:
    if not data or len(data) > 255:
        raise ValueError("invalid page length")
    payload = bytes([address & 0xFF, (address >> 8) & 0xFF, len(data)]) + data
    reply = send_command(port, CMD_WRITE_PAGE, payload)
    if reply == REPLY_ERROR:
        code = read_exact(port, 1)[0]
        raise ProtocolError(f"WRITE failed with error code {code} at 0x{address:04x}")
    if reply != REPLY_OK:
        raise ProtocolError(f"unexpected WRITE reply: {reply!r}")


def read_image(port: serial.Serial, size: int, page_size: int) -> bytes:
    image = bytearray()
    for address in range(0, size, page_size):
        chunk_len = min(page_size, size - address)
        image.extend(read_chunk(port, address, chunk_len))
    return bytes(image)


def write_image(port: serial.Serial, data: bytes, page_size: int) -> None:
    for address in range(0, len(data), page_size):
        page = data[address:address + page_size]
        write_page(port, address, page)
        verify = read_chunk(port, address, len(page))
        if verify != page:
            raise ProtocolError(f"verify failed at 0x{address:04x}")
        print(f"wrote and verified page 0x{address:04x}")


def open_port(device: str, baud: int) -> serial.Serial:
    port = serial.Serial(device, baudrate=baud, timeout=2)
    time.sleep(2.0)
    port.reset_input_buffer()
    port.reset_output_buffer()
    return port


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="CAT28C256 UART programmer")
    parser.add_argument("mode", choices=["read", "write"])
    parser.add_argument("serial_port")
    parser.add_argument("file", type=Path)
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--size", type=int, default=MAX_SIZE, help="bytes to read in read mode")
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if args.mode == "write":
        data = args.file.read_bytes()
        if len(data) > MAX_SIZE:
            raise SystemExit(f"input file too large: {len(data)} > {MAX_SIZE}")
    else:
        if args.size < 1 or args.size > MAX_SIZE:
            raise SystemExit(f"--size must be between 1 and {MAX_SIZE}")

    with open_port(args.serial_port, args.baud) as port:
        page_size, total_size = query_info(port)
        if total_size != MAX_SIZE:
            raise ProtocolError(f"unexpected EEPROM size reported: {total_size}")
        if page_size != DEFAULT_PAGE_SIZE:
            print(f"using device page size {page_size}", file=sys.stderr)

        if args.mode == "read":
            image = read_image(port, args.size, page_size)
            args.file.write_bytes(image)
            print(f"read {len(image)} bytes to {args.file}")
            return 0

        write_image(port, data, page_size)
        print(f"programmed {len(data)} bytes from {args.file}")
        return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except ProtocolError as exc:
        print(f"error: {exc}", file=sys.stderr)
        raise SystemExit(1)