rt4d_uart.py

"""RT-4D UART Communication Protocol

Ported from the C# RT_4D_UART class
"""

import os
import serial
import time
from typing import Optional

# Firmware flash constants
FIRMWARE_SIZE = 251904  # 246 KB - exact firmware image size
FIRMWARE_CHUNK_SIZE = 1024  # 1 KB chunks
_FW_CHECKSUM_BASE = 0x48
_CMD_HANDSHAKE = 0x39
_CMD_WRITE_FW = 0x57
_RESP_ACK = 0x06
_RESP_ERROR = 0xFF
_PROBE_MAX_ATTEMPTS = 50
_HANDSHAKE_MAX_RETRIES = 10
_HANDSHAKE_STEPS = 3
_ERASE_TRIGGER_MAX_RETRIES = 10
_WRITE_MAX_RETRIES = 100
_RETRY_DELAY = 0.05  # 50 ms


def _fw_checksum(command: bytearray) -> None:
    """Calculate firmware checksum (base 0x48) and store in last byte."""
    total = _FW_CHECKSUM_BASE
    for b in command[:-1]:
        total = (total + b) & 0xFF
    command[-1] = total


def validate_firmware_file(file_path: str) -> tuple[bool, str, int]:
    """Validate a firmware file for flashing.

    Returns:
        (is_valid, error_message_or_None, file_size)
    """
    if not os.path.isfile(file_path):
        return False, "File does not exist", 0

    file_size = os.path.getsize(file_path)

    if file_size == 0:
        return False, "File is empty", 0

    if file_size > FIRMWARE_SIZE:
        return (False,
                f"File is too large ({file_size} bytes). "
                f"Maximum allowed is {FIRMWARE_SIZE} bytes ({FIRMWARE_SIZE // 1024} KB)",
                file_size)

    try:
        with open(file_path, 'rb') as f:
            f.read(1)
    except Exception as e:
        return False, f"Cannot read file: {e}", file_size

    return True, None, file_size


def prepare_firmware_data(file_path: str) -> bytearray:
    """Read firmware file and pad to FIRMWARE_SIZE bytes with zeros.

    Raises:
        Exception: If file is invalid.
    """
    is_valid, error_msg, _ = validate_firmware_file(file_path)
    if not is_valid:
        raise Exception(f"Invalid firmware file: {error_msg}")

    with open(file_path, 'rb') as f:
        data = bytearray(f.read())

    if len(data) < FIRMWARE_SIZE:
        data.extend(bytearray(FIRMWARE_SIZE - len(data)))

    return data


class RT4DUART:
    """Communication with RT-4D radio via UART/Serial"""

    def __init__(self):
        self.port: Optional[serial.Serial] = None
        self.read_timeout = 10.0  # seconds
        self.write_timeout = 2.0  # seconds
        self.actual_baudrate: int = 0  # Track which baud rate is actually being used

    def open(self, port_name: str, baudrate: int = 115200):
        """Open serial port connection"""
        self.port = serial.Serial(
            port=port_name,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=self.read_timeout,
            write_timeout=self.write_timeout
        )
        if not self.port.is_open:
            raise IOError(f"Failed to open port {port_name}")
        self.actual_baudrate = baudrate

    def open_with_fallback(self, port_name: str, baudrate: int = 115200) -> tuple[bool, int, str]:
        """Open serial port with automatic fallback to 115200 if higher speed fails

        Args:
            port_name: Serial port name
            baudrate: Requested baud rate

        Returns:
            tuple: (success, actual_baudrate, message)
                - success: True if connection established
                - actual_baudrate: The baud rate that worked
                - message: Description of what happened
        """
        # Try requested baud rate first
        try:
            self.open(port_name, baudrate)

            # Test if connection works with a notify command
            if self.command_notify():
                return (True, baudrate, f"Connected at {baudrate} baud")

            # Connection test failed - try fallback if using high speed
            self.close()

            if baudrate > 115200:
                print(f"Connection failed at {baudrate} baud, trying 115200...")
                print("Note: High-speed mode (256000) requires holding # during radio boot")

                try:
                    self.open(port_name, 115200)

                    if self.command_notify():
                        return (True, 115200, f"Fallback: Connected at 115200 baud (requested {baudrate} failed)")

                    self.close()
                    return (False, 0, f"Connection failed at both {baudrate} and 115200 baud")

                except Exception as e:
                    return (False, 0, f"Fallback connection error: {str(e)}")
            else:
                return (False, 0, f"Connection test failed at {baudrate} baud")

        except Exception as e:
            return (False, 0, f"Error opening port: {str(e)}")

    def close(self):
        """Close serial port"""
        if self.port and self.port.is_open:
            self.port.close()

    def _checksum(self, command: bytearray) -> int:
        """Calculate checksum for command"""
        total = sum(command[:-1])
        return total & 0xFF

    def _verify(self, command: bytes) -> bool:
        """Verify checksum of received data"""
        if len(command) == 0:
            return False
        total = sum(command[:-1])
        return (total & 0xFF) == command[-1]

    def command_notify(self) -> bool:
        """Send notify command to radio"""
        command = bytearray([0x34, 0x00, 0x00, 0x10, 0x00])
        command[-1] = self._checksum(command)

        self.port.write(command)
        self.port.flush()

        try:
            response = self.port.read(1)
            if len(response) == 0:
                print("Timeout waiting for notify response")
                return False
            if response[0] != 0x06:
                print(f"Unexpected notify response: 0x{response[0]:02X}")
                return False
            return True
        except Exception as e:
            print(f"Error in command_notify: {e}")
            return False

    def command_close(self) -> bool:
        """Send close command to radio"""
        command = bytes([0x34, 0x52, 0x05, 0xEE, 0x79])

        try:
            self.port.write(command)
            self.port.flush()
            return True
        except Exception as e:
            print(f"Error closing connection: {e}")
            return False

    def command_read_spi(self, offset: int) -> Optional[bytes]:
        """Read 1KB block from SPI flash at given offset (in KB)"""
        command = bytearray([0x52, 0x00, 0x00, 0x00])
        command[1] = (offset >> 8) & 0xFF
        command[2] = (offset >> 0) & 0xFF
        command[3] = self._checksum(command)

        self.port.write(command)
        self.port.flush()

        # Read 1024 + 4 bytes (data + header + checksum)
        block = bytearray()
        while len(block) < 1028:
            chunk = self.port.read(1028 - len(block))
            if len(chunk) == 0:
                print(f"Timeout reading SPI block at offset 0x{offset:04X}")
                return None
            block.extend(chunk)

        # Check for error response
        if block[0] == 0xFF:
            return None

        # Verify checksum
        if not self._verify(block):
            print(f"Checksum error reading SPI block at 0x{offset:04X}")
            return None

        # Return data only (skip 3-byte header and checksum)
        return bytes(block[3:1027])

    def command_write_spi(self, data: bytes, region: int, start: int, size: int) -> bool:
        """Write data to SPI flash region"""
        command = bytearray(1028)  # 1024 data + 4 header

        for i in range(0, size, 1024):
            command[0] = region
            command[1] = ((i // 1024) >> 8) & 0xFF
            command[2] = ((i // 1024) >> 0) & 0xFF

            # Copy 1KB of data
            chunk_size = min(1024, size - i)
            command[3:3 + chunk_size] = data[start + i:start + i + chunk_size]

            # Calculate checksum
            command[1027] = self._checksum(command)

            # Send command
            self.port.write(command)
            self.port.flush()

            # Wait for ACK
            try:
                response = self.port.read(1)
                if len(response) == 0:
                    print(f"Timeout waiting for write ACK at offset {i}")
                    return False
                if response[0] != 0x06:
                    print(f"Unexpected write response: 0x{response[0]:02X}")
                    return False
            except Exception as e:
                print(f"Error writing SPI: {e}")
                return False

        return True

    def is_bootloader_mode(self) -> bool:
        """Check if radio is in bootloader mode"""
        data = self.command_read_spi(0)

        if data is not None:
            return False  # Normal mode

        # Clear any pending data
        for _ in range(8):
            if self.port.in_waiting == 0:
                break
            self.port.read(1)

        return True  # Bootloader mode

    def read_spi_dump(self, output_file: str, progress_callback=None) -> bool:
        """Read complete SPI flash to file (4MB)

        Args:
            output_file: Path to save the dump
            progress_callback: Optional callback function(current, total) for progress updates
        """
        print("Reading SPI flash...")

        with open(output_file, 'wb') as f:
            for i in range(4096):  # 4096 KB = 4 MB
                percentage = (i / 4096) * 100
                print(f"\rReading SPI flash: {percentage:.1f}%", end='', flush=True)

                # Call progress callback if provided
                if progress_callback:
                    progress_callback(i, 4096)

                data = self.command_read_spi(i)
                if data is None:
                    print(f"\nFailed to read SPI flash at 0x{i:04X}")
                    return False

                f.write(data)

        print("\rReading complete" + " " * 30)

        # Final progress update
        if progress_callback:
            progress_callback(4096, 4096)

        return True

    def read_spi_region(self, address: int, size: int) -> Optional[bytes]:
        """Read a specific region from SPI flash"""
        data = bytearray()
        kb_offset = address // 1024  # Convert byte address to KB offset
        byte_offset = address % 1024  # Offset within first block
        # Account for byte offset when calculating blocks needed
        num_blocks = (byte_offset + size + 1023) // 1024

        for i in range(num_blocks):
            block_data = self.command_read_spi(kb_offset + i)
            if block_data is None:
                print(f"\nFailed to read SPI block at 0x{kb_offset + i:04X}")
                return None
            data.extend(block_data)

        # Return bytes starting from the offset within first block
        return bytes(data[byte_offset:byte_offset + size])

    def write_spi_region(self, data: bytes, region_name: str) -> bool:
        """Write data to a specific SPI region"""
        from rt4d_codeplug.constants import SPI_REGIONS

        if region_name not in SPI_REGIONS:
            print(f"Unknown region: {region_name}")
            return False

        region_info = SPI_REGIONS[region_name]
        region_id = region_info["region_id"]
        address = region_info["address"]
        size = region_info["size"]

        if len(data) > size:
            print(f"Data too large for region {region_name}: {len(data)} > {size}")
            return False

        print(f"Writing {len(data)} bytes to {region_name}...")
        return self.command_write_spi(data, region_id, 0, len(data))

    def command_write_addressbook(self, data: bytes, progress_callback=None) -> bool:
        """Write global contacts (address book) to radio

        Args:
            data: GBK-encoded CSV data (first 6 columns, no header)
            progress_callback: Optional callback(current_block, total_blocks) for progress

        Returns:
            True if successful, False otherwise
        """
        # Prepare data with header
        MAX_SIZE = 29360124  # 28MB max
        data_len = min(len(data), MAX_SIZE)

        # Build payload: 4-byte length header + data
        total_len = data_len + 4
        payload = bytearray(total_len)

        # Write length as 32-bit big-endian
        payload[0] = (total_len >> 24) & 0xFF
        payload[1] = (total_len >> 16) & 0xFF
        payload[2] = (total_len >> 8) & 0xFF
        payload[3] = (total_len) & 0xFF

        # Copy data
        payload[4:4 + data_len] = data[:data_len]

        # Calculate number of 1KB blocks
        num_blocks = (total_len + 1023) // 1024

        print(f"Writing {len(data):,} bytes ({num_blocks} blocks) to address book...")

        # Send blocks
        for block_num in range(num_blocks):
            # Prepare 1028-byte packet
            packet = bytearray(1028)

            # Command and block number
            packet[0] = 0xA4  # Address book write command
            packet[1] = (block_num >> 8) & 0xFF
            packet[2] = (block_num) & 0xFF

            # Copy 1KB of data
            start_offset = block_num * 1024
            end_offset = min(start_offset + 1024, total_len)
            chunk_size = end_offset - start_offset

            # Fill data bytes
            for i in range(1024):
                if start_offset + i < total_len:
                    packet[3 + i] = payload[start_offset + i]
                else:
                    packet[3 + i] = 0xFF  # Pad with 0xFF

            # Calculate checksum
            packet[1027] = self._checksum(packet)

            # Send packet
            try:
                self.port.write(packet)
                self.port.flush()
            except Exception as e:
                print(f"\nError writing block {block_num}: {e}")
                return False

            # Wait for response
            try:
                response = self.port.read(1)
                if len(response) == 0:
                    print(f"\nTimeout waiting for ACK at block {block_num}")
                    return False

                if response[0] == 0x06:
                    # Success - continue
                    if progress_callback:
                        progress_callback(block_num + 1, num_blocks)
                elif response[0] == 0xA4:
                    print(f"\nFlash IC capacity mismatch!")
                    return False
                elif response[0] == 0x4A:
                    print(f"\nFlash IC capacity limit reached!")
                    return False
                else:
                    print(f"\nUnexpected response: 0x{response[0]:02X}")
                    return False

            except Exception as e:
                print(f"\nError reading response at block {block_num}: {e}")
                return False

        print("\nAddress book write complete!")
        return True

    def read_messages(self, region_name: str, progress_callback=None) -> Optional[bytes]:
        """Read messages from a specific SPI region.

        Args:
            region_name: One of 'presets', 'drafts', 'inbox', 'outbox'
            progress_callback: Optional callback(current_kb, total_kb) for progress

        Returns:
            Raw bytes from the message region, or None on error
        """
        from rt4d_codeplug.constants import MESSAGE_REGIONS

        if region_name not in MESSAGE_REGIONS:
            print(f"Unknown message region: {region_name}")
            return None

        region_info = MESSAGE_REGIONS[region_name]
        address = region_info["address"]
        size = region_info["size"]

        print(f"Reading {region_name} messages from 0x{address:06X}...")

        # Read using existing read_spi_region
        kb_offset = address // 1024
        num_blocks = (size + 1023) // 1024

        data = bytearray()
        for i in range(num_blocks):
            if progress_callback:
                progress_callback(i, num_blocks)

            block_data = self.command_read_spi(kb_offset + i)
            if block_data is None:
                print(f"\nFailed to read message block at 0x{(kb_offset + i) * 1024:06X}")
                return None
            data.extend(block_data)

        # Final progress
        if progress_callback:
            progress_callback(num_blocks, num_blocks)

        return bytes(data[:size])

    def write_messages(self, region_name: str, data: bytes, progress_callback=None) -> bool:
        """Write messages to a specific SPI region.

        Note: Message regions may require different write commands than codeplug regions.
        Currently uses standard SPI write which may not work for all message regions.

        Args:
            region_name: One of 'presets', 'drafts', 'inbox', 'outbox'
            data: Raw bytes to write
            progress_callback: Optional callback(current_kb, total_kb) for progress

        Returns:
            True if successful, False otherwise
        """
        from rt4d_codeplug.constants import MESSAGE_REGIONS

        if region_name not in MESSAGE_REGIONS:
            print(f"Unknown message region: {region_name}")
            return False

        region_info = MESSAGE_REGIONS[region_name]
        address = region_info["address"]
        size = region_info["size"]

        if len(data) > size:
            print(f"Data too large for region {region_name}: {len(data)} > {size}")
            return False

        print(f"Writing {len(data)} bytes to {region_name} at 0x{address:06X}...")

        # Pad data to full region size
        padded_data = bytearray(data)
        if len(padded_data) < size:
            padded_data.extend([0xFF] * (size - len(padded_data)))

        # Write using direct address-based SPI write
        # Note: This uses the default_sms region ID (0x97) as a starting point
        # The actual region IDs for message areas may need verification
        num_blocks = (size + 1023) // 1024

        for i in range(num_blocks):
            if progress_callback:
                progress_callback(i, num_blocks)

            # Build write command packet
            packet = bytearray(1028)
            packet[0] = 0x97  # Use default_sms region ID (may need adjustment)
            packet[1] = (i >> 8) & 0xFF
            packet[2] = i & 0xFF

            # Copy 1KB of data
            start = i * 1024
            chunk = padded_data[start:start + 1024]
            packet[3:3 + len(chunk)] = chunk

            # Calculate checksum
            packet[1027] = self._checksum(packet)

            try:
                self.port.write(packet)
                self.port.flush()

                response = self.port.read(1)
                if len(response) == 0:
                    print(f"\nTimeout waiting for write ACK at block {i}")
                    return False
                if response[0] != 0x06:
                    print(f"\nUnexpected write response: 0x{response[0]:02X} at block {i}")
                    return False

            except Exception as e:
                print(f"\nError writing message block {i}: {e}")
                return False

        if progress_callback:
            progress_callback(num_blocks, num_blocks)

        print(f"\n{region_name} write complete!")
        return True

    # ── Firmware flash commands ──────────────────────────────────

    def probe_bootloader(self, max_attempts: int = _PROBE_MAX_ATTEMPTS) -> bool:
        """Probe bootloader by sending 0xFF until it responds.

        Raises:
            IOError: If bootloader does not respond within max_attempts.
        """
        probe_byte = bytearray([_RESP_ERROR])

        for attempt in range(max_attempts):
            try:
                self.port.reset_input_buffer()
                self.port.write(probe_byte)
                time.sleep(_RETRY_DELAY)
                response = self.port.read(1)
                if len(response) == 1 and response[0] == _RESP_ERROR:
                    return True
            except serial.SerialException:
                continue

        raise IOError(f"Bootloader did not respond after {max_attempts} attempts")

    def command_handshake(self, max_retries: int = _HANDSHAKE_MAX_RETRIES) -> bool:
        """Send 3-step handshake to prepare bootloader for firmware update.

        Raises:
            IOError: If any handshake step fails.
        """
        command = bytearray([_CMD_HANDSHAKE, 0x33, 0x05, 0x10, 0x00])
        _fw_checksum(command)

        for step in range(_HANDSHAKE_STEPS):
            success = False
            for _ in range(max_retries):
                try:
                    self.port.reset_input_buffer()
                    self.port.write(command)
                    time.sleep(_RETRY_DELAY)
                    response = self.port.read(1)
                    if len(response) == 1 and response[0] == _RESP_ACK:
                        success = True
                        break
                except serial.SerialException:
                    continue

            if not success:
                raise IOError(f"Handshake step {step + 1}/{_HANDSHAKE_STEPS} failed")

        return True

    def command_trigger_erase(self, max_retries: int = _ERASE_TRIGGER_MAX_RETRIES) -> bool:
        """Trigger firmware erase mode after successful handshake.

        Raises:
            IOError: If erase trigger fails.
        """
        command = bytearray([_CMD_HANDSHAKE, 0x33, 0x05, 0x55, 0x00])
        _fw_checksum(command)

        for _ in range(max_retries):
            try:
                self.port.reset_input_buffer()
                self.port.write(command)
                time.sleep(_RETRY_DELAY)
                response = self.port.read(1)
                if len(response) == 1 and response[0] == _RESP_ACK:
                    return True
            except serial.SerialException:
                continue

        raise IOError(f"Erase trigger failed after {max_retries} retries")

    def command_write_firmware(self, offset: int, data: bytes,
                               max_retries: int = _WRITE_MAX_RETRIES) -> bool:
        """Write a 1 KB chunk of firmware data.

        Args:
            offset: Byte offset in firmware (must be multiple of 1024).
            data: Exactly 1024 bytes of firmware data.

        Raises:
            ValueError: If data is not 1024 bytes.
            IOError: If write fails after retries.
        """
        if len(data) != FIRMWARE_CHUNK_SIZE:
            raise ValueError(f"Data must be exactly {FIRMWARE_CHUNK_SIZE} bytes, got {len(data)}")

        command = bytearray(1028)
        command[0] = _CMD_WRITE_FW
        command[1] = (offset >> 8) & 0xFF
        command[2] = offset & 0xFF
        command[3:3 + FIRMWARE_CHUNK_SIZE] = data
        _fw_checksum(command)

        for _ in range(max_retries):
            try:
                self.port.reset_input_buffer()
                self.port.write(command)
                time.sleep(_RETRY_DELAY)
                response = self.port.read(1)
                if len(response) == 1 and response[0] == _RESP_ACK:
                    return True
            except serial.SerialException:
                continue

        raise IOError(f"Firmware write failed at offset 0x{offset:06X} after {max_retries} retries")