#!/usr/bin/env python3
"""Program an Intel HEX file through the Z80 UART monitor."""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

try:
    import serial
except ImportError:
    print(
        "error: pySerial is required; install it with: "
        "python3 -m pip install pyserial",
        file=sys.stderr,
    )
    raise SystemExit(2)


MAX_BLOCK_SIZE = 24  # Largest W command that fits the monitor's 80-byte line.


class ProgramError(Exception):
    pass


def load_ihex(path: Path) -> dict[int, int]:
    """Return the 16-bit-addressed data contained in an Intel HEX file."""
    memory: dict[int, int] = {}
    address_base = 0
    saw_eof = False

    try:
        lines = path.read_text(encoding="ascii").splitlines()
    except (OSError, UnicodeError) as exc:
        raise ProgramError(f"cannot read {path}: {exc}") from exc

    for line_number, text in enumerate(lines, 1):
        text = text.strip()
        if not text:
            continue
        if saw_eof:
            raise ProgramError(f"{path}:{line_number}: data after EOF record")
        if not text.startswith(":"):
            raise ProgramError(f"{path}:{line_number}: record does not start with ':'")
        try:
            record = bytes.fromhex(text[1:])
        except ValueError as exc:
            raise ProgramError(f"{path}:{line_number}: invalid hexadecimal record") from exc

        if len(record) < 5 or len(record) != record[0] + 5:
            raise ProgramError(f"{path}:{line_number}: invalid record length")
        if sum(record) & 0xFF:
            raise ProgramError(f"{path}:{line_number}: Intel HEX checksum mismatch")

        count = record[0]
        offset = (record[1] << 8) | record[2]
        record_type = record[3]
        data = record[4 : 4 + count]

        if record_type == 0x00:
            start = address_base + offset
            if start + count > 0x10000:
                raise ProgramError(
                    f"{path}:{line_number}: address is outside the Z80's 64 KiB space"
                )
            for index, value in enumerate(data):
                address = start + index
                old_value = memory.get(address)
                if old_value is not None and old_value != value:
                    raise ProgramError(
                        f"{path}:{line_number}: conflicting data at {address:04X}"
                    )
                memory[address] = value
        elif record_type == 0x01:
            if count != 0:
                raise ProgramError(f"{path}:{line_number}: malformed EOF record")
            saw_eof = True
        elif record_type == 0x02:
            if count != 2:
                raise ProgramError(
                    f"{path}:{line_number}: malformed extended segment address"
                )
            address_base = int.from_bytes(data, "big") << 4
        elif record_type == 0x04:
            if count != 2:
                raise ProgramError(
                    f"{path}:{line_number}: malformed extended linear address"
                )
            address_base = int.from_bytes(data, "big") << 16
        elif record_type in (0x03, 0x05):
            # Start-address records do not contain bytes to program.
            if count != 4:
                raise ProgramError(
                    f"{path}:{line_number}: malformed start-address record"
                )
            continue
        else:
            raise ProgramError(
                f"{path}:{line_number}: unsupported record type {record_type:02X}"
            )

    if not saw_eof:
        raise ProgramError(f"{path}: missing Intel HEX EOF record")
    if not memory:
        raise ProgramError(f"{path}: file contains no data")
    return memory


def make_blocks(memory: dict[int, int], block_size: int) -> list[tuple[int, bytes]]:
    """Combine adjacent addresses without filling gaps from the HEX file."""
    blocks: list[tuple[int, bytes]] = []
    addresses = sorted(memory)
    index = 0

    while index < len(addresses):
        start = addresses[index]
        data = bytearray([memory[start]])
        index += 1
        while (
            index < len(addresses)
            and len(data) < block_size
            and addresses[index] == start + len(data)
        ):
            data.append(memory[addresses[index]])
            index += 1
        blocks.append((start, bytes(data)))

    return blocks


class Monitor:
    def __init__(self, port: str, baud: int, timeout: float):
        try:
            self.serial = serial.Serial(
                port=port,
                baudrate=baud,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=timeout,
                write_timeout=timeout,
            )
        except serial.SerialException as exc:
            raise ProgramError(f"cannot open serial port {port}: {exc}") from exc

    def close(self) -> None:
        self.serial.close()

    def command(self, command: str) -> str:
        try:
            self.serial.write(command.encode("ascii") + b"\r")
            self.serial.flush()
            response = self.serial.read_until(b"> ")
        except serial.SerialException as exc:
            raise ProgramError(f"serial communication failed: {exc}") from exc
        if not response.endswith(b"> "):
            raise ProgramError(f"timeout waiting for monitor after '{command}'")
        return response.decode("ascii", errors="replace")

    def synchronize(self) -> None:
        try:
            self.serial.reset_input_buffer()
        except serial.SerialException as exc:
            raise ProgramError(f"serial communication failed: {exc}") from exc
        response = self.command("")
        if not response.endswith("> "):
            raise ProgramError("monitor prompt not found")


def response_has_ok(response: str) -> bool:
    return any(line.strip() == "OK" for line in response.splitlines())


def response_checksum(response: str) -> int | None:
    matches = [
        int(line.strip(), 16)
        for line in response.splitlines()
        if re.fullmatch(r"[0-9A-Fa-f]{4}", line.strip())
    ]
    return matches[-1] if matches else None


def write_and_verify(monitor: Monitor, start: int, data: bytes) -> tuple[bool, str]:
    values = " ".join(f"{value:02X}" for value in data)
    response = monitor.command(f"W {start:04X} {values}")
    if not response_has_ok(response):
        return False, "monitor did not acknowledge the write"

    end = start + len(data) - 1
    expected = sum(data) & 0xFFFF
    response = monitor.command(f"C {start:04X} {end:04X}")
    actual = response_checksum(response)
    if actual is None:
        return False, "monitor did not return a checksum"
    if actual != expected:
        return False, f"checksum mismatch (expected {expected:04X}, received {actual:04X})"
    return True, ""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Program an Intel HEX file through the Z80 UART monitor."
    )
    parser.add_argument("port", help="serial port, for example /dev/ttyUSB0 or COM3")
    parser.add_argument("file", type=Path, help="Intel HEX input file")
    parser.add_argument("--baud", type=int, default=2400, help="baud rate (default: 2400)")
    parser.add_argument(
        "--block-size",
        type=int,
        default=16,
        help=f"bytes per verified block, 1-{MAX_BLOCK_SIZE} (default: 16)",
    )
    parser.add_argument(
        "--retries",
        type=int,
        default=3,
        help="retries after the first attempt (default: 3)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=5.0,
        help="serial response timeout in seconds (default: 5)",
    )
    args = parser.parse_args()
    if not 1 <= args.block_size <= MAX_BLOCK_SIZE:
        parser.error(f"--block-size must be between 1 and {MAX_BLOCK_SIZE}")
    if args.retries < 0:
        parser.error("--retries cannot be negative")
    if args.baud <= 0:
        parser.error("--baud must be positive")
    if args.timeout <= 0:
        parser.error("--timeout must be positive")
    return args


def main() -> int:
    args = parse_args()
    monitor: Monitor | None = None

    try:
        memory = load_ihex(args.file)
        blocks = make_blocks(memory, args.block_size)
        print(
            f"Loaded {len(memory)} bytes in {len(blocks)} blocks "
            f"({min(memory):04X}-{max(memory):04X})"
        )

        monitor = Monitor(args.port, args.baud, args.timeout)
        monitor.synchronize()

        for number, (start, data) in enumerate(blocks, 1):
            end = start + len(data) - 1
            for attempt in range(args.retries + 1):
                try:
                    if attempt:
                        monitor.synchronize()
                    success, reason = write_and_verify(monitor, start, data)
                except ProgramError as exc:
                    success, reason = False, str(exc)
                if success:
                    print(
                        f"\rBlock {number}/{len(blocks)} {start:04X}-{end:04X} verified",
                        end="",
                        flush=True,
                    )
                    break
                if attempt == args.retries:
                    raise ProgramError(
                        f"block {start:04X}-{end:04X} failed after "
                        f"{args.retries + 1} attempts: {reason}"
                    )
                print(
                    f"\nRetrying {start:04X}-{end:04X} "
                    f"({attempt + 1}/{args.retries}): {reason}",
                    file=sys.stderr,
                )

        print(f"\nProgrammed and verified {len(memory)} bytes")
        return 0
    except ProgramError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("\nCancelled", file=sys.stderr)
        return 130
    finally:
        if monitor is not None:
            monitor.close()


if __name__ == "__main__":
    raise SystemExit(main())