"""RPC command encoder/decoder for BCM4775 flash operations.

The RPC layer sits atop the reliable transport and encodes the 6 flash
commands that the bootloader patch exposes after loading. Each command
has a specific request format and response format.

RPC commands are registered via the OnLhReady callback after the
bootloader Scratch APP starts (FSM state 16). Command IDs from
FlashUpdater decompilation:

    6  = FLASH_READ     — Read flash contents to host
    9  = FLASH_WRITE    — Write data to flash
    10 = FLASH_ERASE    — Erase flash sectors (4KB granularity)
    35 = CONFIG_QUERY   — Query flash size and sector erase time
    36 = CRC_VERIFY     — Compute CRC over flash region
    39 = STATUS         — Query flash operation status

GHIDRA: The exact RPC encoding (header layout, field order) needs
verification via decompilation of rpc_engine.cpp functions.
"""

import logging
import struct

from .bcm_models import CrcResult, FlashConfig, FlashStatus, RpcCommandId
from .bcm_protocol import BcmReliableTransport

logger = logging.getLogger(__name__)

# Flash sector size — from FlashUpdater: 4KB (0x1000) sector-aligned operations
FLASH_SECTOR_SIZE = 0x1000  # 4096 bytes

# RPC status codes (inferred from FlashUpdater debug strings)
RPC_STATUS_OK = 0x00
RPC_STATUS_ERROR = 0x01
RPC_STATUS_BUSY = 0x02


def _encode_rpc_request(cmd_id: int, params: bytes = b"") -> bytes:
    """Encode an RPC request payload.

    GHIDRA: The exact format needs verification. Using a simple layout:
        [cmd_id: 2 bytes LE][params: variable]

    Args:
        cmd_id: RPC command identifier.
        params: Command-specific parameter bytes.

    Returns:
        Encoded RPC request payload (goes inside a transport packet).
    """
    return struct.pack("<H", cmd_id) + params


def _decode_rpc_response(data: bytes) -> tuple[int, int, bytes]:
    """Decode an RPC response payload.

    GHIDRA: The exact format needs verification. Using:
        [cmd_id: 2 bytes LE][status: 1 byte][response_data: variable]

    Returns:
        Tuple of (cmd_id, status_code, response_data).

    Raises:
        ValueError: If response is too short.
    """
    if len(data) < 3:
        raise ValueError(f"RPC response too short: {len(data)} bytes (minimum 3)")
    cmd_id = struct.unpack_from("<H", data, 0)[0]
    status = data[2]
    response_data = data[3:]
    return cmd_id, status, response_data


class RpcEngine:
    """RPC command executor for BCM4775 flash operations.

    Provides async methods for each of the 6 flash commands. Each method
    encodes the request, sends it via the reliable transport, and decodes
    the response into a typed result.

    All methods raise ValueError on protocol errors and TimeoutError on
    communication failures.
    """

    def __init__(self, transport: BcmReliableTransport):
        self._transport = transport

    async def _execute(self, cmd_id: int, params: bytes = b"", timeout: float = 5.0) -> bytes:
        """Send an RPC command and return the response data.

        Args:
            cmd_id: RPC command ID.
            params: Command-specific parameters.
            timeout: Time to wait for response.

        Returns:
            Response data bytes (after cmd_id and status).

        Raises:
            TimeoutError: If no response received.
            ValueError: If response indicates error.
        """
        request = _encode_rpc_request(cmd_id, params)
        ack = await self._transport.send_packet(request, timeout=timeout)

        if ack is None:
            raise TimeoutError(f"RPC command {cmd_id} got no ACK after retries")

        # For commands that return data, we need to receive the response packet
        response = await self._transport.receive_packet(timeout=timeout)
        if response is None:
            raise TimeoutError(f"RPC command {cmd_id} got no response")

        resp_cmd, status, data = _decode_rpc_response(response.payload)

        if resp_cmd != cmd_id:
            raise ValueError(
                f"RPC response command mismatch: sent {cmd_id}, got {resp_cmd}"
            )

        if status != RPC_STATUS_OK:
            raise ValueError(
                f"RPC command {cmd_id} failed with status 0x{status:02X}"
            )

        return data

    async def config_query(self) -> FlashConfig:
        """Query flash device configuration (RPC 35).

        Returns flash size and sector erase time. This is the first
        RPC command sent after bootloader starts — it tells us the
        flash geometry for backup/erase/write calculations.
        """
        data = await self._execute(RpcCommandId.CONFIG_QUERY)

        if len(data) < 8:
            raise ValueError(f"CONFIG_QUERY response too short: {len(data)} bytes")

        flash_size, erase_time = struct.unpack_from("<II", data, 0)
        logger.info("Flash config: size=%d bytes, sector_erase=%dms", flash_size, erase_time)

        return FlashConfig(flash_size=flash_size, sector_erase_time_ms=erase_time)

    async def flash_read(self, address: int, length: int) -> bytes:
        """Read flash contents (RPC 6).

        Args:
            address: Flash start address.
            length: Number of bytes to read.

        Returns:
            Flash data bytes.
        """
        params = struct.pack("<II", address, length)
        # Flash reads can be slow for large regions — scale timeout
        timeout = max(5.0, length / 1024)  # ~1s per KB minimum
        data = await self._execute(RpcCommandId.FLASH_READ, params, timeout=timeout)

        if len(data) < length:
            logger.warning(
                "Flash read returned %d bytes, expected %d", len(data), length
            )

        return data

    async def flash_write(self, address: int, data: bytes) -> None:
        """Write data to flash (RPC 9).

        The address must be 4KB sector-aligned. Data is written in a
        single RPC call — the caller is responsible for chunking large
        writes if needed.

        Args:
            address: Flash destination address (must be sector-aligned).
            data: Bytes to write.

        Raises:
            ValueError: If address is not sector-aligned.
        """
        if address % FLASH_SECTOR_SIZE != 0:
            raise ValueError(
                f"Flash write address 0x{address:08X} is not "
                f"{FLASH_SECTOR_SIZE}-byte sector aligned"
            )

        params = struct.pack("<II", address, len(data)) + data
        timeout = max(5.0, len(data) / 512)  # ~2s per KB for writes
        await self._execute(RpcCommandId.FLASH_WRITE, params, timeout=timeout)
        logger.info("Flash write: %d bytes at 0x%08X", len(data), address)

    async def flash_erase(self, address: int, length: int) -> None:
        """Erase flash sectors (RPC 10).

        Erases the flash region starting at `address` for `length` bytes.
        Both must be 4KB sector-aligned. Erase sets all bits to 0xFF.

        Args:
            address: Flash start address (must be sector-aligned).
            length: Number of bytes to erase (must be sector-aligned).

        Raises:
            ValueError: If address or length is not sector-aligned.
        """
        if address % FLASH_SECTOR_SIZE != 0:
            raise ValueError(
                f"Flash erase address 0x{address:08X} is not sector aligned"
            )
        if length % FLASH_SECTOR_SIZE != 0:
            raise ValueError(
                f"Flash erase length {length} is not sector aligned"
            )

        sectors = length // FLASH_SECTOR_SIZE
        params = struct.pack("<II", address, length)
        # Erase timeout: allow sector_erase_time per sector + overhead
        timeout = max(10.0, sectors * 0.5)
        await self._execute(RpcCommandId.FLASH_ERASE, params, timeout=timeout)
        logger.info("Flash erase: %d sectors at 0x%08X", sectors, address)

    async def crc_verify(self, address: int, length: int) -> CrcResult:
        """Compute CRC over a flash region (RPC 36).

        From FlashUpdater format string:
            RpcFlashResponse::Crc(addr=0x%08x, size=%lu, crc=0x%02x)

        The CRC is a single byte (0x00-0xFF). The caller compares
        this against a locally computed CRC of the expected data.

        Args:
            address: Flash start address.
            length: Number of bytes to verify.

        Returns:
            CrcResult with address, size, crc value, and match flag (caller sets match).
        """
        params = struct.pack("<II", address, length)
        timeout = max(5.0, length / 1024)
        data = await self._execute(RpcCommandId.CRC_VERIFY, params, timeout=timeout)

        if len(data) < 1:
            raise ValueError("CRC_VERIFY response has no CRC byte")

        crc = data[0]
        logger.info("Flash CRC: addr=0x%08X size=%d crc=0x%02X", address, length, crc)

        return CrcResult(address=address, size=length, crc=crc, match=False)

    async def flash_status(self) -> FlashStatus:
        """Query flash operation status (RPC 39).

        Returns the current status of the flash controller. Used to
        check if a previous erase or write operation is still in progress.
        """
        data = await self._execute(RpcCommandId.STATUS)

        code = data[0] if data else 0
        busy = code == RPC_STATUS_BUSY
        logger.debug("Flash status: code=0x%02X busy=%s", code, busy)

        return FlashStatus(code=code, busy=busy)


def compute_local_crc(data: bytes) -> int:
    """Compute the same CRC used by the BCM4775 for flash verification.

    GHIDRA: The exact algorithm needs verification. Using XOR of all
    bytes, matching the transport checksum pattern. This must produce
    the same result as the BCM4775's CRC_VERIFY RPC command.

    Args:
        data: Bytes to checksum.

    Returns:
        Single-byte CRC value (0x00-0xFF).
    """
    result = 0
    for b in data:
        result ^= b
    return result & 0xFF
