"""Core firmware flashing primitives using SRIO-over-TFTP.

This module provides the high-level erase/write/verify operations
derived from `_OLD/linux_app/dwl_fw.cpp` but adapted to use TFTP transport
with SRIO command encoding (compatible with Vecchia_app protocol).

Reference: dwl_fw.cpp main() loop and fpgaflashengine.cpp operations.
"""
from __future__ import annotations

import time
from typing import Callable, Optional

from pydownloadfwviasrio.tftp_client import SRIOTFTPClient
from pydownloadfwviasrio.profiles import FlashProfile
from pydownloadfwviasrio.core.srio_flash import SRIOFlashController

# Delay constants from dwl_fw.cpp (line 18-19)
WRITE_DELAY_US = 100000  # 100ms delay after write_flash
WAITFLASH_DELAY_US = 100000  # 100ms delay after wait_flash
MAX_RETRY_READ = 1000  # Max retries for wait_flash


class FirmwareFlasher:
    """High-level firmware flasher using TFTP+SRIO transport.

    Methods mirror the sequence in dwl_fw.cpp: init -> erase -> write -> verify.
    Each method accepts an optional `log` callback for progress updates.

    Reference: `_OLD/linux_app/dwl_fw.cpp` (erase_section, write_flash, read_flash).
    """

    def __init__(
        self,
        profile: Optional[FlashProfile] = None,
        client: Optional[SRIOTFTPClient] = None,
        max_register_retries: int = 500,
    ) -> None:
        """Initialize firmware flasher.
        
        Args:
            profile: Flash profile with target configuration.
            client: SRIO TFTP client.
            max_register_retries: Max retries for register write verification.
        """
        self.profile = profile
        self.client = client
        self.chunk_size = 256  # SRIO write chunk (dwl_fw.cpp: DWL_CHUNK_SIZE)
        self.flash_controller: Optional[SRIOFlashController] = None
        self.qspi_initialized = False  # Track QSPI initialization
        self.max_register_retries = max_register_retries
        
        # Initialize flash controller if profile and client available
        if self.profile and self.client:
            self.flash_controller = SRIOFlashController(
                client=self.client,
                slot_address=self.profile.slot_address,
                flashport=0,  # TODO: make configurable
                max_verify_retries=max_register_retries,
            )

    def _log(self, message: str, log: Optional[Callable[[str], None]] = None) -> None:
        if log:
            log(message)
    
    def init_qspi(
        self,
        log: Optional[Callable[[str], None]] = None,
        tftp_log: Optional[Callable[[str], None]] = None,
    ) -> bool:
        """Initialize QSPI flash interface (from dwl_fw.cpp main line 513-514).
        
        MUST be called before any flash operations.
        Calls op_init_qspi() TWICE as in the original code.
        
        Args:
            log: Optional logging callback for general messages.
            tftp_log: Optional logging callback for TFTP commands.
        
        Returns:
            True on success.
        """
        if not self.flash_controller:
            self._log("ERROR: No flash controller configured", log)
            return False
        
        self._log("Initializing QSPI flash interface...", log)
        
        # Call op_init_qspi TWICE (exactly as dwl_fw.cpp main() does)
        if not self.flash_controller.op_init_qspi(log=tftp_log):
            self._log("QSPI init (1st) FAILED", log)
            return False
        
        if not self.flash_controller.op_init_qspi(log=tftp_log):
            self._log("QSPI init (2nd) FAILED", log)
            return False
        
        self.qspi_initialized = True
        self._log("QSPI initialized successfully", log)
        return True
    
    def init_qspi(
        self,
        log: Optional[Callable[[str], None]] = None,
        tftp_log: Optional[Callable[[str], None]] = None,
    ) -> bool:
        """Initialize QSPI flash interface (from dwl_fw.cpp main line 513-514).
        
        MUST be called before any flash operations.
        Calls op_init_qspi() TWICE as in the original code.
        
        Args:
            log: Optional logging callback for general messages.
            tftp_log: Optional logging callback for TFTP commands.
        
        Returns:
            True on success.
        """
        if not self.flash_controller:
            self._log("ERROR: No flash controller configured", log)
            return False
        
        self._log("Initializing QSPI flash interface...", log)
        
        # Call op_init_qspi TWICE (exactly as dwl_fw.cpp main() does)
        if not self.flash_controller.op_init_qspi(log=tftp_log):
            self._log("QSPI init (1st) FAILED", log)
            return False
        
        if not self.flash_controller.op_init_qspi(log=tftp_log):
            self._log("QSPI init (2nd) FAILED", log)
            return False
        
        self.qspi_initialized = True
        self._log("QSPI initialized successfully", log)
        return True

    def erase(
        self,
        address: int,
        length: int = 65536,
        log: Optional[Callable[[str], None]] = None,
        tftp_log: Optional[Callable[[str], None]] = None,
        progress_callback: Optional[Callable[[int, int], None]] = None,
        abort_check: Optional[Callable[[], bool]] = None,
    ) -> bool:
        """Erase a flash sector (64 KB default).

        Reference: dwl_fw.cpp erase_section() + wait_flash_with_retry().
        In the real target, this sends SRIO erase command (0xD8) via QSPI.
        Here we emulate by sending a dummy TFTP write (future: custom command).

        Args:
            address: Flash address (must be sector-aligned, typically 64 KB).
            length: Erase block size (default 65536 = 64 KB).
            log: Optional logging callback for general messages.
            tftp_log: Optional logging callback for TFTP commands.
            progress_callback: Optional callback(current, total) for progress bar.
            abort_check: Optional callback that returns True if operation should abort.

        Returns:
            True on success, False on failure.
        """
        self._log(f"Erasing sector at 0x{address:08X} ({length} bytes)", log)
        if not self.flash_controller:
            self._log("ERROR: No flash controller configured", log)
            return False
        
        # Set abort check on flash controller for TFTP-level interruption
        self.flash_controller.set_abort_check(abort_check)
        
        # Initialize QSPI if not done yet (dwl_fw.cpp main line 513-514)
        if not self.qspi_initialized:
            if not self.init_qspi(log=log, tftp_log=tftp_log):
                self.flash_controller.set_abort_check(None)
                return False

        # Check for abort request before starting
        if abort_check and abort_check():
            self._log("Erase operation aborted", log)
            self.flash_controller.set_abort_check(None)
            return False

        if progress_callback:
            progress_callback(0, 100)

        # Erase exactly as dwl_fw.cpp main line 528-532
        if not self.flash_controller.erase_section(address, log=tftp_log):
            self._log(f"Erase FAILED at 0x{address:08X}", log)
            return False
        
        # Check for abort after erase command
        if abort_check and abort_check():
            self._log("Erase operation aborted", log)
            return False
        
        # wait_flash_with_retry (dwl_fw.cpp line 529)
        if not self.flash_controller.wait_flash(max_retries=MAX_RETRY_READ, log=tftp_log):
            self._log(f"Erase wait timeout at 0x{address:08X}", log)
            return False

        if progress_callback:
            progress_callback(100, 100)

        self._log(f"Erase completed: 0x{address:08X}", log)
        self.flash_controller.set_abort_check(None)
        return True

    def write(
        self,
        address: int,
        data: bytes,
        log: Optional[Callable[[str], None]] = None,
        tftp_log: Optional[Callable[[str], None]] = None,
        progress_callback: Optional[Callable[[int, int], None]] = None,
        verify_after_write: bool = False,
        max_chunk_write_retries: int = 3,
        abort_check: Optional[Callable[[], bool]] = None,
    ) -> bool:
        """Write binary data to flash in 256-byte chunks.

        Reference: dwl_fw.cpp main() write loop with write_flash() + wait_flash().
        Each chunk is sent via TFTP with filename encoding the SRIO address.

        Args:
            address: Flash start address.
            data: Binary payload to write.
            log: Optional logging callback for general messages.
            tftp_log: Optional logging callback for TFTP commands.
            progress_callback: Optional callback(current, total) for progress bar.
            verify_after_write: If True, read-back and verify each chunk after writing.
            max_chunk_write_retries: Max retry attempts for chunk write+verify.
            abort_check: Optional callback that returns True if operation should abort.

        Returns:
            True on success, False on failure.
        """
        if not self.flash_controller:
            self._log("ERROR: No flash controller configured", log)
            return False
        
        # Set abort check on flash controller for TFTP-level interruption
        self.flash_controller.set_abort_check(abort_check)
        
        # Initialize QSPI if not done yet (dwl_fw.cpp main line 513-514)
        if not self.qspi_initialized:
            if not self.init_qspi(log=log, tftp_log=tftp_log):
                self.flash_controller.set_abort_check(None)
                return False

        total = len(data)
        self._log(f"Writing {total} bytes to 0x{address:08X}", log)

        written = 0
        current_address = address
        while written < total:
            # Check for abort request
            if abort_check and abort_check():
                self._log("Write operation aborted", log)
                return False
            
            # Chunk data (pad last chunk to 256 bytes with 0x00)
            # dwl_fw.cpp line 547-553
            chunk_data = data[written : written + self.chunk_size]
            if len(chunk_data) < self.chunk_size:
                chunk_data += b"\x00" * (self.chunk_size - len(chunk_data))
            
            # Retry loop at chunk level (write + optional verify)
            chunk_success = False
            for chunk_attempt in range(max_chunk_write_retries):
                # write_flash (dwl_fw.cpp line 561)
                if not self.flash_controller.write_flash(current_address, chunk_data, log=tftp_log):
                    if chunk_attempt < max_chunk_write_retries - 1:
                        self._log(f"↻ Chunk write retry {chunk_attempt+1}/{max_chunk_write_retries-1} at 0x{current_address:08X}", log)
                        time.sleep(0.5)  # Wait before retry
                        continue
                    self._log(f"⚠️ Write FAILED at 0x{current_address:08X} after {max_chunk_write_retries} attempts", log)
                    return False
                
                # usleep(WRITE_DELAY_US*value); (dwl_fw.cpp line 562)
                time.sleep(WRITE_DELAY_US / 1_000_000.0)
                
                # wait_flash_with_retry (dwl_fw.cpp line 563)
                if not self.flash_controller.wait_flash(max_retries=MAX_RETRY_READ, log=tftp_log):
                    if chunk_attempt < max_chunk_write_retries - 1:
                        self._log(f"↻ Wait timeout, retry {chunk_attempt+1}/{max_chunk_write_retries-1} at 0x{current_address:08X}", log)
                        time.sleep(0.5)
                        continue
                    self._log(f"⚠️ Wait timeout at 0x{current_address:08X} after {max_chunk_write_retries} attempts", log)
                    return False
                
                # usleep(WAITFLASH_DELAY_US*value); (dwl_fw.cpp line 564)
                time.sleep(WAITFLASH_DELAY_US / 1_000_000.0)
                
                # Optional: verify chunk immediately after write
                if verify_after_write:
                    self._log(f"Verifying chunk at 0x{current_address:08X}...", log)
                    read_data = self.flash_controller.read_flash(current_address, self.chunk_size, log=tftp_log)
                    if read_data is None:
                        if chunk_attempt < max_chunk_write_retries - 1:
                            self._log(f"↻ Verify read failed, retry {chunk_attempt+1}/{max_chunk_write_retries-1} at 0x{current_address:08X}", log)
                            time.sleep(0.5)
                            continue
                        self._log(f"⚠️ Verify read FAILED at 0x{current_address:08X} after {max_chunk_write_retries} attempts", log)
                        return False
                    
                    # Compare with expected data
                    if len(read_data) != len(chunk_data):
                        if chunk_attempt < max_chunk_write_retries - 1:
                            self._log(f"↻ Verify length mismatch, retry {chunk_attempt+1}/{max_chunk_write_retries-1} at 0x{current_address:08X}", log)
                            time.sleep(0.5)
                            continue
                        self._log(f"⚠️ Verify length mismatch at 0x{current_address:08X} after {max_chunk_write_retries} attempts", log)
                        return False
                    
                    if read_data != chunk_data:
                        # Find first mismatch for debugging
                        mismatch_offset = -1
                        for i in range(len(read_data)):
                            if read_data[i] != chunk_data[i]:
                                mismatch_offset = i
                                break
                        
                        if chunk_attempt < max_chunk_write_retries - 1:
                            self._log(f"↻ Verify data mismatch at offset {mismatch_offset}, retry {chunk_attempt+1}/{max_chunk_write_retries-1} at 0x{current_address:08X}", log)
                            time.sleep(0.5)
                            continue
                        
                        self._log(f"⚠️ Verify FAILED at 0x{current_address:08X}: mismatch at offset {mismatch_offset} (got 0x{read_data[mismatch_offset]:02X}, expected 0x{chunk_data[mismatch_offset]:02X}) after {max_chunk_write_retries} attempts", log)
                        return False
                    
                    if chunk_attempt > 0:  # Log solo se c'è stato retry
                        self._log(f"✓ Verified 0x{current_address:08X} after {chunk_attempt+1} attempt(s)", log)
                    else:
                        self._log(f"✓ Verified 0x{current_address:08X}", log)
                
                # Chunk written successfully (with or without verify)
                chunk_success = True
                break  # Exit retry loop
            
            if not chunk_success:
                return False  # Should not reach here, but safety check

            written += len(chunk_data)
            current_address += self.chunk_size
            self._log(f"Progress: {written}/{total} bytes", log)
            
            # Update progress bar if callback provided
            if progress_callback:
                progress_callback(written, total)

        self._log("Write completed", log)
        self.flash_controller.set_abort_check(None)
        return True

    def verify(
        self,
        address: int,
        data: bytes,
        log: Optional[Callable[[str], None]] = None,
        tftp_log: Optional[Callable[[str], None]] = None,
        progress_callback: Optional[Callable[[int, int], None]] = None,        abort_check: Optional[Callable[[], bool]] = None,    ) -> bool:
        """Verify flash contents by reading back and comparing.

        Reference: dwl_fw.cpp read_flash() for verification.
        Reads back data in chunks and compares with expected.

        Args:
            address: Flash start address.
            data: Expected binary data.
            log: Optional logging callback for general messages.
            tftp_log: Optional logging callback for TFTP commands.
            progress_callback: Optional callback(current, total) for progress bar.

        Returns:
            True if verification passes, False otherwise.
        """
        if not self.flash_controller:
            self._log("ERROR: No flash controller configured", log)
            return False
        
        # Set abort check on flash controller for TFTP-level interruption
        self.flash_controller.set_abort_check(abort_check)
        
        # Initialize QSPI if not done yet (dwl_fw.cpp main line 513-514)
        if not self.qspi_initialized:
            if not self.init_qspi(log=log, tftp_log=tftp_log):
                self.flash_controller.set_abort_check(None)
                return False

        total = len(data)
        self._log(f"Verifying {total} bytes at 0x{address:08X}", log)

        verified = 0
        current_address = address
        while verified < total:
            # Check for abort request
            if abort_check and abort_check():
                self._log("Verify operation aborted", log)
                return False
            
            chunk_size = min(self.chunk_size, total - verified)
            expected_chunk = data[verified : verified + chunk_size]

            # Read back via SRIO flash protocol (from dwl_fw.cpp)
            read_chunk = self.flash_controller.read_flash(
                current_address,
                chunk_size,
                log=tftp_log,
            )
            if read_chunk is None:
                self._log(f"Verify READ FAILED at 0x{current_address:08X}", log)
                return False

            # Compare (only compare actual data length)
            if read_chunk[:chunk_size] != expected_chunk:
                self._log(f"Verify MISMATCH at 0x{current_address:08X}", log)
                return False

            verified += chunk_size
            current_address += chunk_size
            self._log(f"Verify progress: {verified}/{total} bytes", log)
            
            # Update progress bar if callback provided
            if progress_callback:
                progress_callback(verified, total)

        self._log("Verify OK", log)
        self.flash_controller.set_abort_check(None)
        return True


def example_sequence(log: Optional[Callable[[str], None]] = None) -> None:
    """Example high-level sequence used by the GUI/tests.

    This demonstrates how the operations are called in order.
    """
    flasher = FirmwareFlasher()
    base = 0x100000
    data = b"\xFF" * 1024
    flasher.erase(base, len(data), log=log)
    flasher.write(base, data, log=log)
    flasher.verify(base, data, log=log)


if __name__ == "__main__":
    def _print(s: str) -> None:
        print(s)

    example_sequence(log=_print)
