"""BCM4775 firmware updater — 18-state FSM + flash operations.

Implements the complete firmware update sequence for the Broadcom
BCM4775 GNSS SoC as reverse-engineered from FlashUpdater4775.exe.

The update process:
    1. Auto-baud → ROM sync → version/chip ID validation
    2. Download bootloader.pkg to RAM → RAM sync → verify
    3. RPC flash operations: backup → erase → write → CRC verify

Safety guardrails (non-negotiable):
    1. Chip ID must validate against BCM477x table before any write
    2. Flash backup before erase (default, skip_backup opt-out)
    3. CRC verify after every write — abort on mismatch
    4. Flash offset must be 4KB sector-aligned
    5. CLI requires --yes, MCP requires confirm=True
    6. abort() waits for current sector write to finish
    7. No automatic retry on write failure — report and let user decide
"""

import logging
import struct
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path

from .bcm_models import (
    BCM_CHIP_IDS,
    BackupResult,
    FirmwareDownloaderState,
    FirmwareInfo,
    UpdateResult,
)
from .bcm_protocol import (
    BOOTLOADER_HEADER_SIZE,
    BOOTLOADER_MAGIC,
    BcmReliableTransport,
)
from .bcm_rpc import FLASH_SECTOR_SIZE, RpcEngine, compute_local_crc

logger = logging.getLogger(__name__)

# Type alias for progress callbacks.
# Called with (current_bytes, total_bytes, message)
ProgressCallback = Callable[[int, int, str], None]


def _null_progress(current: int, total: int, message: str) -> None:
    """Default no-op progress callback."""


class BootloaderPackage:
    """Parser for bootloader.pkg firmware container.

    The bootloader.pkg is a Broadcom proprietary format:
        - 4-byte magic: f3 f2 f1 f0
        - 4-byte identifier/version
        - 4-byte build info
        - 4-byte payload length (little-endian)
        - 16 bytes reserved/metadata
        - N bytes executable payload (starts at offset 0x20)

    The payload is a "Scratch APP" (type 0x41) that provides RPC
    flash command handlers when running in RAM on the BCM4775.
    """

    def __init__(self, data: bytes):
        """Parse a bootloader package from raw bytes.

        Args:
            data: Complete bootloader.pkg file contents.

        Raises:
            ValueError: If magic bytes don't match or file is too small.
        """
        if len(data) < BOOTLOADER_HEADER_SIZE:
            raise ValueError(
                f"Bootloader package too small: {len(data)} bytes "
                f"(minimum {BOOTLOADER_HEADER_SIZE})"
            )

        magic = data[:4]
        if magic != BOOTLOADER_MAGIC:
            raise ValueError(
                f"Invalid bootloader magic: {magic.hex()} "
                f"(expected {BOOTLOADER_MAGIC.hex()})"
            )

        self._data = data
        self._identifier = struct.unpack_from("<I", data, 4)[0]
        self._build_info = struct.unpack_from("<I", data, 8)[0]
        self._payload_length = struct.unpack_from("<I", data, 12)[0]

        # Validate payload length against file size
        expected_size = BOOTLOADER_HEADER_SIZE + self._payload_length
        if len(data) < expected_size:
            raise ValueError(
                f"Bootloader package truncated: file is {len(data)} bytes, "
                f"header declares {self._payload_length} byte payload "
                f"(need {expected_size} total)"
            )

    @classmethod
    def from_file(cls, path: str | Path) -> "BootloaderPackage":
        """Load bootloader package from a file.

        Args:
            path: Path to bootloader.pkg file.

        Returns:
            Parsed BootloaderPackage.
        """
        path = Path(path)
        data = path.read_bytes()
        logger.info("Loaded bootloader from %s (%d bytes)", path, len(data))
        return cls(data)

    @property
    def identifier(self) -> int:
        """Package identifier/version from header offset 0x04."""
        return self._identifier

    @property
    def build_info(self) -> int:
        """Build timestamp or version from header offset 0x08."""
        return self._build_info

    @property
    def payload_length(self) -> int:
        """Payload size in bytes from header offset 0x0C."""
        return self._payload_length

    @property
    def payload(self) -> bytes:
        """Executable payload (everything after the 32-byte header)."""
        return self._data[BOOTLOADER_HEADER_SIZE : BOOTLOADER_HEADER_SIZE + self._payload_length]

    @property
    def total_size(self) -> int:
        """Total package size including header."""
        return len(self._data)

    @property
    def header(self) -> bytes:
        """Raw 32-byte header."""
        return self._data[:BOOTLOADER_HEADER_SIZE]


class FirmwareUpdater:
    """Main firmware update orchestrator.

    Drives the 18-state FSM to establish communication with the
    BCM4775 bootloader, then uses RPC commands for flash operations.

    Usage:
        updater = FirmwareUpdater(transport)
        info = await updater.get_info()  # Safe — NMEA only
        result = await updater.update_firmware("new_fw.bin")
    """

    def __init__(
        self,
        transport: BcmReliableTransport,
        bootloader: BootloaderPackage | None = None,
        progress: ProgressCallback | None = None,
    ):
        """Initialize firmware updater.

        Args:
            transport: Reliable binary transport for BCM4775 communication.
            bootloader: Pre-loaded bootloader package. If None, must be
                       provided to update_firmware() via bootloader_path.
            progress: Optional callback for progress reporting.
        """
        self._transport = transport
        self._rpc = RpcEngine(transport)
        self._bootloader = bootloader
        self._progress = progress or _null_progress
        self._state = FirmwareDownloaderState.FD_IDLE
        self._chip_id: int = 0
        self._abort_requested = False
        self._bootloader_running = False

    @property
    def state(self) -> FirmwareDownloaderState:
        """Current FSM state."""
        return self._state

    @property
    def chip_id(self) -> int:
        """Chip ID read from the BCM4775 (0 if not yet read)."""
        return self._chip_id

    @property
    def chip_name(self) -> str:
        """Human-readable chip name, or 'Unknown' if not in table."""
        return BCM_CHIP_IDS.get(self._chip_id, f"Unknown (0x{self._chip_id:08X})")

    @property
    def bootloader_running(self) -> bool:
        """Whether the bootloader Scratch APP is running on the BCM4775."""
        return self._bootloader_running

    def abort(self) -> None:
        """Request graceful abort of the current operation.

        Sets a flag that the FSM checks between sectors. The current
        sector write (if any) finishes before the abort takes effect.
        """
        logger.warning("Abort requested")
        self._abort_requested = True

    async def enter_bootloader(self) -> int:
        """Run the FSM through bootloader download and return chip ID.

        Executes FSM states 1-17: auto-baud → ROM sync → version →
        bootloader download → RAM sync → verify.

        Returns:
            Validated chip ID (from BCM_CHIP_IDS table).

        Raises:
            ValueError: If chip ID not in known table.
            TimeoutError: If communication fails at any stage.
            RuntimeError: If FSM reaches an unexpected state.
        """
        if self._bootloader is None:
            raise ValueError("No bootloader package loaded")

        self._state = FirmwareDownloaderState.FD_REQUEST_DOWNLOAD
        return await self._run_fsm()

    async def get_info(self) -> FirmwareInfo:
        """Query module info without entering bootloader.

        This uses NMEA commands (safe, non-destructive). Does NOT enter
        the bootloader or touch flash.

        Note: chip_id will be 0 unless enter_bootloader() has been called.
        """
        return FirmwareInfo(
            chip_id=self._chip_id,
            chip_name=self.chip_name if self._chip_id else "",
            bootloader_active=self._bootloader_running,
        )

    async def backup_flash(self, output_path: str | Path | None = None) -> BackupResult:
        """Read entire flash to a local file.

        Requires the bootloader to be running (call enter_bootloader first).

        Args:
            output_path: Where to save the backup. If None, generates a
                        timestamped filename in the current directory.

        Returns:
            BackupResult with path, size, and CRC.

        Raises:
            RuntimeError: If bootloader is not running.
        """
        if not self._bootloader_running:
            raise RuntimeError("Bootloader not running — call enter_bootloader() first")

        # Query flash config to determine size
        config = await self._rpc.config_query()
        flash_size = config.flash_size

        if output_path is None:
            ts = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
            output_path = Path(f"FlashBackup_{ts}.bin")
        else:
            output_path = Path(output_path)

        self._progress(0, flash_size, "Backing up flash...")
        logger.info("Backing up %d bytes of flash to %s", flash_size, output_path)

        # Read flash in chunks
        chunk_size = 4096
        data = bytearray()
        offset = 0

        while offset < flash_size:
            if self._abort_requested:
                raise RuntimeError("Backup aborted by user")

            remaining = flash_size - offset
            read_size = min(chunk_size, remaining)
            chunk = await self._rpc.flash_read(offset, read_size)
            data.extend(chunk)
            offset += len(chunk)
            self._progress(offset, flash_size, f"Reading flash: {offset}/{flash_size}")

        output_path.write_bytes(bytes(data))
        crc = compute_local_crc(bytes(data))
        self._progress(flash_size, flash_size, "Backup complete")
        logger.info("Flash backup: %d bytes, CRC=0x%02X, saved to %s", len(data), crc, output_path)

        return BackupResult(path=str(output_path), size=len(data), crc=crc)

    async def update_firmware(
        self,
        firmware_path: str | Path,
        flash_offset: int = 0,
        skip_backup: bool = False,
        bootloader_path: str | Path | None = None,
    ) -> UpdateResult:
        """Full firmware update sequence.

        1. Enter bootloader (if not already running)
        2. Validate chip ID
        3. Backup flash (unless skip_backup=True)
        4. Erase target region
        5. Write new firmware
        6. CRC verify
        7. Report result

        Args:
            firmware_path: Path to the new firmware .bin file.
            flash_offset: Flash offset to write at (must be 4KB aligned).
            skip_backup: Skip flash backup before erase.
            bootloader_path: Path to bootloader.pkg (uses pre-loaded if None).

        Returns:
            UpdateResult with success status and details.
        """
        firmware_path = Path(firmware_path)
        self._abort_requested = False

        # Validate flash offset alignment (Safety guardrail #4)
        if flash_offset % FLASH_SECTOR_SIZE != 0:
            return UpdateResult(
                success=False,
                error=f"Flash offset 0x{flash_offset:08X} is not "
                f"{FLASH_SECTOR_SIZE}-byte sector aligned",
            )

        # Load firmware image
        if not firmware_path.exists():
            return UpdateResult(success=False, error=f"Firmware file not found: {firmware_path}")
        firmware_data = firmware_path.read_bytes()
        logger.info("Loaded firmware: %s (%d bytes)", firmware_path, len(firmware_data))

        # Load bootloader if needed
        if bootloader_path and self._bootloader is None:
            self._bootloader = BootloaderPackage.from_file(bootloader_path)

        # Enter bootloader if not already running
        backup_path = ""
        try:
            if not self._bootloader_running:
                self._progress(0, 100, "Entering bootloader...")
                chip_id = await self.enter_bootloader()
            else:
                chip_id = self._chip_id

            # Safety guardrail #1: Chip ID validation
            if chip_id not in BCM_CHIP_IDS:
                return UpdateResult(
                    success=False,
                    chip_id=chip_id,
                    error=f"Unknown chip ID 0x{chip_id:08X} — refusing to write. "
                    f"Known chips: {', '.join(BCM_CHIP_IDS.values())}",
                )

            # Safety guardrail #2: Backup before erase
            if not skip_backup:
                self._progress(0, 100, "Backing up flash...")
                backup = await self.backup_flash()
                backup_path = backup.path

            # Calculate erase region (round up to sector boundary)
            erase_length = len(firmware_data)
            if erase_length % FLASH_SECTOR_SIZE != 0:
                erase_length = ((erase_length // FLASH_SECTOR_SIZE) + 1) * FLASH_SECTOR_SIZE

            # Erase
            self._progress(0, erase_length, "Erasing flash sectors...")
            if self._abort_requested:
                return UpdateResult(
                    success=False, chip_id=chip_id, chip_name=self.chip_name,
                    backup_path=backup_path, error="Aborted before erase",
                )
            await self._rpc.flash_erase(flash_offset, erase_length)
            self._progress(erase_length, erase_length, "Erase complete")

            # Write in sector-sized chunks
            total = len(firmware_data)
            written = 0
            while written < total:
                if self._abort_requested:
                    return UpdateResult(
                        success=False, chip_id=chip_id, chip_name=self.chip_name,
                        bytes_written=written, backup_path=backup_path,
                        error=f"Aborted after writing {written}/{total} bytes",
                    )

                chunk_end = min(written + FLASH_SECTOR_SIZE, total)
                chunk = firmware_data[written:chunk_end]
                # Pad last chunk to sector boundary if needed
                if len(chunk) < FLASH_SECTOR_SIZE and chunk_end >= total:
                    chunk = chunk + b"\xff" * (FLASH_SECTOR_SIZE - len(chunk))

                await self._rpc.flash_write(flash_offset + written, chunk)
                written = chunk_end
                self._progress(written, total, f"Writing: {written}/{total} bytes")

            # Safety guardrail #3: CRC verify after write
            self._progress(0, total, "Verifying flash...")
            crc_result = await self._rpc.crc_verify(flash_offset, len(firmware_data))
            local_crc = compute_local_crc(firmware_data)
            verified = crc_result.crc == local_crc

            if not verified:
                return UpdateResult(
                    success=False, chip_id=chip_id, chip_name=self.chip_name,
                    bytes_written=written, verified=False,
                    backup_path=backup_path,
                    error=f"CRC mismatch: flash=0x{crc_result.crc:02X} "
                    f"local=0x{local_crc:02X}",
                )

            self._progress(total, total, "Update complete — verified")
            logger.info(
                "Firmware update complete: %d bytes written, CRC verified",
                written,
            )

            return UpdateResult(
                success=True, chip_id=chip_id, chip_name=self.chip_name,
                bytes_written=written, verified=True, backup_path=backup_path,
            )

        except Exception as e:
            logger.error("Firmware update failed: %s", e)
            return UpdateResult(
                success=False, chip_id=self._chip_id, chip_name=self.chip_name,
                backup_path=backup_path, error=str(e),
            )

    async def verify_flash(self, firmware_path: str | Path, flash_offset: int = 0) -> bool:
        """Read-back CRC comparison without writing.

        Computes CRC of the firmware file and compares against the
        flash contents. Requires bootloader to be running.

        Args:
            firmware_path: Path to firmware .bin to verify against.
            flash_offset: Flash offset where firmware should be.

        Returns:
            True if CRC matches, False otherwise.
        """
        if not self._bootloader_running:
            raise RuntimeError("Bootloader not running — call enter_bootloader() first")

        firmware_data = Path(firmware_path).read_bytes()
        local_crc = compute_local_crc(firmware_data)

        self._progress(0, len(firmware_data), "Verifying flash CRC...")
        crc_result = await self._rpc.crc_verify(flash_offset, len(firmware_data))
        self._progress(len(firmware_data), len(firmware_data), "Verification complete")

        match = crc_result.crc == local_crc
        if match:
            logger.info("Flash verification passed: CRC=0x%02X", local_crc)
        else:
            logger.error(
                "Flash verification FAILED: flash=0x%02X local=0x%02X",
                crc_result.crc,
                local_crc,
            )
        return match

    # ------------------------------------------------------------------
    # FSM implementation
    # ------------------------------------------------------------------

    async def _run_fsm(self) -> int:
        """Execute the 18-state firmware downloader FSM.

        Runs from FD_REQUEST_DOWNLOAD through FD_DONE, returning
        the validated chip ID on success.

        This is the core state machine from firmware_downloader.cpp
        (FUN_00517b80), reimplemented in Python.
        """
        raw = self._transport.raw_transport

        while self._state != FirmwareDownloaderState.FD_DONE:
            if self._abort_requested:
                raise RuntimeError("FSM aborted by user")

            match self._state:
                case FirmwareDownloaderState.FD_REQUEST_DOWNLOAD:
                    # State 1: Initiate download — flush input, prepare for auto-baud
                    logger.info("FSM: REQUEST_DOWNLOAD — starting bootloader sequence")
                    await raw.flush_input()
                    self._state = FirmwareDownloaderState.FD_WAIT4AUTOBAUD

                case FirmwareDownloaderState.FD_WAIT4AUTOBAUD:
                    # State 2: Send auto-baud sequence
                    logger.info("FSM: WAIT4AUTOBAUD — sending auto-baud")
                    await self._transport.send_autobaud()
                    self._state = FirmwareDownloaderState.FD_SENT_AUTOBAUD

                case FirmwareDownloaderState.FD_SENT_AUTOBAUD:
                    # State 3: Wait briefly for auto-baud to take effect
                    logger.debug("FSM: SENT_AUTOBAUD — waiting for sync")
                    self._state = FirmwareDownloaderState.FD_WAIT4SYNC

                case FirmwareDownloaderState.FD_WAIT4SYNC:
                    # State 4: Attempt ROM sync
                    logger.info("FSM: WAIT4SYNC — sending ROM sync")
                    if await self._transport.send_sync(target="rom"):
                        self._state = FirmwareDownloaderState.FD_GOT_SYNC
                    else:
                        self._state = FirmwareDownloaderState.FD_TIMEOUT_SYNC

                case FirmwareDownloaderState.FD_GOT_SYNC:
                    # State 5: Route to ROM (7) or RAM (14) based on sync target
                    # First time through: ROM sync → state 7
                    # After bootloader download: RAM sync → state 14
                    if self._bootloader_running:
                        logger.info("FSM: GOT_SYNC — routing to RAM sync confirm")
                        self._state = FirmwareDownloaderState.FD_GOT_SYNC_CONFIRM
                    else:
                        logger.info("FSM: GOT_SYNC — routing to ROM init")
                        self._state = FirmwareDownloaderState.FD_GOT_SYNC_INIT

                case FirmwareDownloaderState.FD_TIMEOUT_SYNC:
                    # State 6: Sync timed out — could retry or fail
                    raise TimeoutError("Failed to sync with BCM4775 ROM bootloader")

                case FirmwareDownloaderState.FD_GOT_SYNC_INIT:
                    # State 7: ROM sync achieved — send version request
                    logger.info("FSM: GOT_SYNC_INIT — requesting ROM version")
                    await self._send_version_request()
                    self._state = FirmwareDownloaderState.FD_WAIT4VERSION_INIT

                case FirmwareDownloaderState.FD_WAIT4VERSION_INIT:
                    # State 8: Wait for version response
                    logger.info("FSM: WAIT4VERSION_INIT — reading version")
                    self._chip_id = await self._read_version_response()
                    self._state = FirmwareDownloaderState.FD_GOT_VERSION_INIT

                case FirmwareDownloaderState.FD_GOT_VERSION_INIT:
                    # State 9: Validate chip ID against known table
                    logger.info(
                        "FSM: GOT_VERSION_INIT — chip_id=0x%08X (%s)",
                        self._chip_id,
                        self.chip_name,
                    )
                    if self._chip_id not in BCM_CHIP_IDS:
                        self._state = FirmwareDownloaderState.FD_ROMRESET
                    else:
                        self._state = FirmwareDownloaderState.FD_WAIT4DOWNLOAD_START

                case FirmwareDownloaderState.FD_ROMRESET:
                    # State 10: Chip ID mismatch — fail
                    raise ValueError(
                        f"Unsupported chip ID 0x{self._chip_id:08X}. "
                        f"Known BCM477x IDs: "
                        + ", ".join(f"0x{k:08X}" for k in BCM_CHIP_IDS)
                    )

                case FirmwareDownloaderState.FD_WAIT4DOWNLOAD_START:
                    # State 11: Begin bootloader download
                    assert self._bootloader is not None  # guaranteed by enter_bootloader()
                    logger.info("FSM: WAIT4DOWNLOAD_START — starting patch transfer")
                    self._progress(0, self._bootloader.payload_length, "Downloading bootloader...")
                    self._state = FirmwareDownloaderState.FD_DOWNLOADING

                case FirmwareDownloaderState.FD_DOWNLOADING:
                    # State 12: Transfer bootloader payload to BCM4775 RAM
                    await self._download_bootloader()
                    self._state = FirmwareDownloaderState.FD_WAIT4APP

                case FirmwareDownloaderState.FD_WAIT4APP:
                    # State 13: Wait for Scratch APP to start, then sync RAM
                    logger.info("FSM: WAIT4APP — waiting for bootloader to start")
                    self._bootloader_running = True
                    if await self._transport.send_sync(target="ram"):
                        self._state = FirmwareDownloaderState.FD_GOT_SYNC_CONFIRM
                    else:
                        raise TimeoutError("Failed to sync with RAM bootloader")

                case FirmwareDownloaderState.FD_GOT_SYNC_CONFIRM:
                    # State 14: RAM sync achieved — send version request
                    logger.info("FSM: GOT_SYNC_CONFIRM — requesting RAM version")
                    await self._send_version_request()
                    self._state = FirmwareDownloaderState.FD_WAIT4VERSION_CONFIRM

                case FirmwareDownloaderState.FD_WAIT4VERSION_CONFIRM:
                    # State 15: Wait for RAM version response
                    logger.info("FSM: WAIT4VERSION_CONFIRM — reading version")
                    ram_chip_id = await self._read_version_response()
                    if ram_chip_id != self._chip_id:
                        logger.warning(
                            "RAM chip ID 0x%08X differs from ROM chip ID 0x%08X",
                            ram_chip_id,
                            self._chip_id,
                        )
                    self._state = FirmwareDownloaderState.FD_GOT_VERSION_CONFIRM

                case FirmwareDownloaderState.FD_GOT_VERSION_CONFIRM:
                    # State 16: Bootloader verified — ready for flash operations
                    logger.info("FSM: GOT_VERSION_CONFIRM — bootloader ready")
                    self._progress(100, 100, "Bootloader active")
                    self._state = FirmwareDownloaderState.FD_DONE

        logger.info("FSM: DONE — chip_id=0x%08X (%s)", self._chip_id, self.chip_name)
        return self._chip_id

    async def _send_version_request(self) -> None:
        """Send a version request to the BCM4775.

        GHIDRA: The exact binary format of the version request needs
        verification. Using a minimal request packet.
        """
        # GHIDRA: Version request bytes need decompilation
        version_request = struct.pack("<H", 0x0001)  # GHIDRA: placeholder command ID
        await self._transport.send_packet(version_request)

    async def _read_version_response(self) -> int:
        """Read version response and extract chip ID.

        GHIDRA: The exact response format needs verification.
        Chip ID is stored at this+0x8AC in the FlashUpdater (4 bytes LE).
        """
        response = await self._transport.receive_packet(timeout=5.0)
        if response is None:
            raise TimeoutError("No version response received")

        # GHIDRA: Extract chip ID from response — offset needs verification
        if len(response.payload) >= 4:
            chip_id = struct.unpack_from("<I", response.payload, 0)[0]
        else:
            raise ValueError(
                f"Version response too short: {len(response.payload)} bytes"
            )

        return chip_id

    async def _download_bootloader(self) -> None:
        """Transfer bootloader.pkg payload to BCM4775 RAM.

        Sends the bootloader payload in chunks via the reliable transport.
        The BCM4775 ROM bootloader receives each chunk and writes it to
        the patch RAM address.

        GHIDRA: The exact chunk framing (how payload is segmented,
        address/offset headers per chunk) needs verification.
        """
        assert self._bootloader is not None  # guaranteed by enter_bootloader()
        payload = self._bootloader.payload
        total = len(payload)
        chunk_size = 4096  # GHIDRA: actual chunk size from decompilation
        sent = 0

        while sent < total:
            if self._abort_requested:
                raise RuntimeError("Bootloader download aborted")

            chunk_end = min(sent + chunk_size, total)
            chunk = payload[sent:chunk_end]

            # GHIDRA: Chunk may need address/offset header prepended
            ack = await self._transport.send_packet(chunk)
            if ack is None:
                raise TimeoutError(f"Bootloader download stalled at offset {sent}")

            sent = chunk_end
            self._progress(sent, total, f"Bootloader: {sent}/{total} bytes")

        logger.info("Bootloader download complete: %d bytes transferred", total)
