# Copyright (C) 2026, Advanced Micro Devices, Inc.  All rights reserved.
# Portions of this file consist of AI-generated content.
# SPDX-License-Identifier: MIT

"""Zynq boot image loader module.

Reference: https://docs.amd.com/r/en-US/ug1283-bootgen-user-guide
"""

from typing import Optional
from enum import IntEnum
import struct
from .loader import (
    BootImageLoader,
    BootHeader,
    ImageHeaderTable,
    ImageHeader,
    PartitionHeader,
)


class ZynqBootHeader(BootHeader):
    """Zynq-specific boot header with additional fields."""
    
    # Header versions
    VERSION_ZYNQ = 0x01010000
    
    # Boot header sizes and counts
    MAX_BOOT_VECTORS = 8
    UDF_BH_SIZE = 19  # User Defined Fields size in words
    UDF_BH_SIZE_BYTES = 76  # User Defined Fields size in bytes
    
    # Boot header offsets
    BH_ARM_VECTOR_TABLE_OFFSET = 0x00
    BH_WIDTH_DETECTION_OFFSET = 0x20
    BH_IDENTIFICATION_OFFSET = 0x24
    BH_KEY_SOURCE_OFFSET = 0x28
    BH_HEADER_VERSION_OFFSET = 0x2C
    BH_SOURCE_OFFSET = 0x30
    BH_FSBL_LENGTH_OFFSET = 0x34
    BH_FSBL_LOAD_ADDR_OFFSET = 0x38
    BH_FSBL_EXEC_ADDR_OFFSET = 0x3C
    BH_TOTAL_FSBL_LENGTH_OFFSET = 0x40
    BH_QSPI_CONFIG_OFFSET = 0x44
    BH_CHECKSUM_OFFSET = 0x48
    BH_UDF_OFFSET = 0x4C
    BH_IHT_OFFSET = 0x98
    BH_PHT_OFFSET = 0x9C
    
    def __init__(self):
        """Initialize the Zynq boot header."""
        super().__init__()
        self.arm_vector_table: Optional[bytes] = None # 0x00 - 0x1F
        # Width detection word at 0x20
        # Header signature at 0x24
        self.key_source: Optional[int] = None # 0x28
        self.header_version: Optional[bytes] = None # 0x2C - 0x01010000
        self.source_offset: Optional[int] = None # 0x30
        self.fsbl_image_length: Optional[int] = None # 0x34
        self.fsbl_load_address: Optional[int] = None # 0x38
        self.fsbl_execution_address: Optional[int] = None # 0x3C
        self.total_fsbl_length: Optional[int] = None # 0x40
        self.qspi_configuration_word: Optional[int] = None # 0x44
        self.boot_header_checksum: Optional[int] = None # 0x48
        self.user_defined_fields: Optional[bytes] = None  # 0x4c- 0x97 - 76 bytes
        self.image_header_table_offset: Optional[int] = None # 0x98
        self.partition_header_table_offset: Optional[int] = None # 0x9C

    def compute_checksum(self) -> int:
        """Compute the checksum over the boot header fields.
        
        The checksum covers words from 0x20 to 0x44 inclusive:
        width_detection_word, header_signature, key_source, header_version,
        source_offset, fsbl_image_length, fsbl_load_address, fsbl_execution_address,
        total_fsbl_length, and qspi_configuration_word.
        
        Returns:
            The computed checksum as a 32-bit unsigned integer.
        """
        # width_detection_word, header_signature, and header_version are stored
        # as raw file bytes; unpack as little-endian to get the integer values.
        width_word = struct.unpack('<I', self.width_detection_word or b'\x00\x00\x00\x00')[0]
        signature_word = struct.unpack('<I', self.header_signature or b'\x00\x00\x00\x00')[0]
        header_version_word = struct.unpack('<I', self.header_version or b'\x00\x00\x00\x00')[0]
        
        words = [
            width_word,                          # 0x20
            signature_word,                      # 0x24
            self.key_source or 0,                # 0x28
            header_version_word,                 # 0x2C
            self.source_offset or 0,             # 0x30
            self.fsbl_image_length or 0,         # 0x34
            self.fsbl_load_address or 0,         # 0x38
            self.fsbl_execution_address or 0,    # 0x3C
            self.total_fsbl_length or 0,         # 0x40
            self.qspi_configuration_word or 0,   # 0x44
        ]
        
        # Sum all words and mask to 32 bits
        checksum = sum(words) & 0xFFFFFFFF
        
        # Invert (one's complement)
        checksum = (~checksum) & 0xFFFFFFFF
        
        return checksum

    def verify_checksum(self) -> bool:
        """Verify the boot header checksum.
        
        Returns:
            True if the checksum is valid, False otherwise.
        """
        if self.boot_header_checksum is None:
            return False
        return self.compute_checksum() == self.boot_header_checksum


class ZynqImageHeaderTable(ImageHeaderTable):
    """Zynq-specific image header table."""
    
    # Valid versions
    VERSION_1_1 = 0x01010000
    VERSION_1_2 = 0x01020000
    
    def __init__(self):
        """Initialize the Zynq image header table."""
        super().__init__()
        self.version: Optional[int] = None # 0x00 - Versions 0x01010000 and 0x01020000
        self.count_of_image_headers: Optional[int] = None # 0x04
        self.first_partition_header_offset: Optional[int] = None # 0x08
        self.first_image_header_offset: Optional[int] = None # 0x0C
        self.header_authentication_certificate_offset: Optional[int] = None # 0x10 (0x0102000 only)
        # reserved 0x14

    def is_version_valid(self) -> bool:
        """Check if the version is valid.
        
        Returns:
            True if version is 0x01010000 or 0x01020000, False otherwise.
        """
        return self.version in (self.VERSION_1_1, self.VERSION_1_2)


class ZynqImageHeader(ImageHeader):
    """Zynq-specific image header."""
    
    def __init__(self):
        """Initialize the Zynq image header."""
        super().__init__()
        self.next_image_header: Optional[int] = None # 0x00
        self.corresponding_partition_header: Optional[int] = None # 0x04
        # Reserved 0x08 - always 0x00000000
        self.partition_count_length: Optional[int] = None # 0x0C
        self.image_name: Optional[str] = None # 0x10 - N: Null-terminated ASCII string
        # Reserved to 64 byte boundary
        self.partition_headers = []  # List of partition headers for this image


class ZynqPartitionHeader(PartitionHeader):
    """Zynq-specific partition header."""
    
    # Maximum partitions
    MAX_NUM_PARTITIONS = 14
    
    # Partition header reserved fields
    PHT_RESERVED = 4  # Number of reserved words in partition header
    
    # Partition attribute bit shifts
    PH_DEST_DEVICE_SHIFT = 4
    PH_CHECKSUM_SHIFT = 12
    PH_AC_FLAG_SHIFT = 15
    PH_OWNER_SHIFT = 16
    
    # Partition attribute bit masks
    PH_DEST_DEVICE_MASK = 0xF
    PH_CHECKSUM_MASK = 0x7
    PH_AC_FLAG_MASK = 0x1
    PH_OWNER_MASK = 0x3
    
    # Partition header offsets
    PH_ENCRYPTED_LENGTH_OFFSET = 0x00
    PH_UNENCRYPTED_LENGTH_OFFSET = 0x04
    PH_TOTAL_LENGTH_OFFSET = 0x08
    PH_LOAD_ADDR_OFFSET = 0x0C
    PH_EXEC_ADDR_OFFSET = 0x10
    PH_DATA_OFFSET = 0x14
    PH_ATTRIBUTES_OFFSET = 0x18
    PH_SECTION_COUNT_OFFSET = 0x1C
    PH_CHECKSUM_OFFSET_FIELD = 0x20
    PH_IH_OFFSET = 0x24
    PH_AC_OFFSET = 0x28
    PH_RESERVED_OFFSET = 0x2C
    PH_HEADER_CHECKSUM_OFFSET = 0x3C
    
    # Partition header size
    PARTITION_HEADER_SIZE = 64  # bytes
    
    # Encryption header signature (for detecting encrypted partitions)
    ENCRYPTION_HEADER = bytes([
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0xBB, 0x00, 0x00, 0x00, 0x44, 0x00, 0x22, 0x11,
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0x66, 0x55, 0x99, 0xAA
    ])
    ENCRYPTION_HEADER_SIZE = 28

    class PartitionOwner(IntEnum):
        """Partition owner types (bits 17:16)."""
        FSBL = 0
        UBOOT = 1
        RESERVED_2 = 2
        RESERVED_3 = 3

    class ChecksumType(IntEnum):
        """Checksum types (bits 14:12)."""
        NONE = 0
        MD5 = 1
        RESERVED_2 = 2
        RESERVED_3 = 3
        RESERVED_4 = 4
        RESERVED_5 = 5
        RESERVED_6 = 6
        RESERVED_7 = 7

    class DestinationDevice(IntEnum):
        """Destination device types (bits 7:4)."""
        NONE = 0
        PS = 1
        PL = 2
        INT = 3
        RESERVED_4 = 4
        RESERVED_5 = 5
        RESERVED_6 = 6
        RESERVED_7 = 7
        RESERVED_8 = 8
        RESERVED_9 = 9
        RESERVED_10 = 10
        RESERVED_11 = 11
        RESERVED_12 = 12
        RESERVED_13 = 13
        RESERVED_14 = 14
        RESERVED_15 = 15

    def __init__(self):
        """Initialize the Zynq partition header."""
        super().__init__()
        self.encrypted_partition_length: Optional[int] = None  # 0x00
        self.unencrypted_partition_length: Optional[int] = None  # 0x04
        self.total_partition_word_length: Optional[int] = None  # 0x08
        self.destination_load_address: Optional[int] = None  # 0x0C
        self.destination_execution_address: Optional[int] = None  # 0x10
        self.data_word_offset_in_image: Optional[int] = None  # 0x14
        self.attribute_bits: Optional[int] = None  # 0x18
        self.section_count: Optional[int] = None  # 0x1C
        self.checksum_word_offset: Optional[int] = None  # 0x20
        self.image_header_word_offset: Optional[int] = None  # 0x24
        self.authentication_certificate_word_offset: Optional[int] = None  # 0x28
        # Reserved 0x2C-0x38 (16 bytes)
        self.header_checksum: Optional[int] = None  # 0x3C

    def get_partition_owner(self) -> 'ZynqPartitionHeader.PartitionOwner':
        """Extract partition owner from attribute bits (bits 17:16).
        
        Returns:
            PartitionOwner enum value.
        """
        if self.attribute_bits is None:
            return ZynqPartitionHeader.PartitionOwner.FSBL
        return ZynqPartitionHeader.PartitionOwner((self.attribute_bits >> 16) & 0x3)

    def get_rsa_signature_present(self) -> bool:
        """Extract RSA signature present flag from attribute bits (bit 15).
        
        Returns:
            True if RSA signature is present, False otherwise.
        """
        if self.attribute_bits is None:
            return False
        return bool((self.attribute_bits >> 15) & 0x1)

    def get_checksum_type(self) -> 'ZynqPartitionHeader.ChecksumType':
        """Extract checksum type from attribute bits (bits 14:12).
        
        Returns:
            ChecksumType enum value.
        """
        if self.attribute_bits is None:
            return ZynqPartitionHeader.ChecksumType.NONE
        return ZynqPartitionHeader.ChecksumType((self.attribute_bits >> 12) & 0x7)

    def get_destination_device(self) -> 'ZynqPartitionHeader.DestinationDevice':
        """Extract destination device from attribute bits (bits 7:4).
        
        Returns:
            DestinationDevice enum value.
        """
        if self.attribute_bits is None:
            return ZynqPartitionHeader.DestinationDevice.NONE
        return ZynqPartitionHeader.DestinationDevice((self.attribute_bits >> 4) & 0xF)

    def get_attribute_summary(self) -> str:
        """Get a human-readable summary of the attribute bits.
        
        Returns:
            String describing the partition attributes.
        """
        return (f"Owner={self.get_partition_owner().name}, "
                f"RSA={self.get_rsa_signature_present()}, "
                f"Checksum={self.get_checksum_type().name}, "
                f"Dest={self.get_destination_device().name}")

    def get_encrypted_partition_length_bytes(self) -> int:
        """Get the encrypted partition length in bytes.
        
        The length is stored in words (4 bytes each) in the header.
        
        Returns:
            The encrypted partition length in bytes.
        """
        if self.encrypted_partition_length is None:
            return 0
        return self.encrypted_partition_length * 4

    def get_unencrypted_partition_length_bytes(self) -> int:
        """Get the unencrypted partition length in bytes.
        
        The length is stored in words (4 bytes each) in the header.
        
        Returns:
            The unencrypted partition length in bytes.
        """
        if self.unencrypted_partition_length is None:
            return 0
        return self.unencrypted_partition_length * 4

    def get_total_partition_length_bytes(self) -> int:
        """Get the total partition length in bytes.
        
        The length is stored in words (4 bytes each) in the header.
        
        Returns:
            The total partition length in bytes.
        """
        if self.total_partition_word_length is None:
            return 0
        return self.total_partition_word_length * 4

    def get_data_offset_bytes(self) -> int:
        """Get the data offset in bytes.
        
        The offset is stored in words (4 bytes each) in the header.
        
        Returns:
            The data offset in bytes.
        """
        if self.data_word_offset_in_image is None:
            return 0
        return self.data_word_offset_in_image * 4

    def compute_checksum(self) -> int:
        """Compute the checksum over the partition header fields.
        
        The checksum is the sum of all 32-bit words in the header,
        excluding the checksum field itself (last 4 bytes).
        
        Returns:
            The computed checksum as a 32-bit unsigned integer.
        """
        # Pack the header fields into words (excluding checksum and reserved)
        words = [
            self.encrypted_partition_length or 0,
            self.unencrypted_partition_length or 0,
            self.total_partition_word_length or 0,
            self.destination_load_address or 0,
            self.destination_execution_address or 0,
            self.data_word_offset_in_image or 0,
            self.attribute_bits or 0,
            self.section_count or 0,
            self.checksum_word_offset or 0,
            self.image_header_word_offset or 0,
            self.authentication_certificate_word_offset or 0,
            0, 0, 0, 0,  # Reserved fields (4 words)
        ]
        
        # Sum all words and mask to 32 bits
        checksum = sum(words) & 0xFFFFFFFF
        
        # Invert (one's complement)
        checksum = (~checksum) & 0xFFFFFFFF
        
        return checksum

    def verify_checksum(self) -> bool:
        """Verify the partition header checksum.
        
        Returns:
            True if the checksum is valid, False otherwise.
        """
        if self.header_checksum is None:
            return False
        return self.compute_checksum() == self.header_checksum


class ZynqBootImageLoader(BootImageLoader):
    """Zynq-specific boot image loader implementation."""

    def __init__(self, image_path=None, image_path_offset=0):
        """Initialize the Zynq boot image loader.
        
        Args:
            image_path: Path to the boot image file.
            image_path_offset: Offset within the image file to start reading from.
        """
        super().__init__(image_path, image_path_offset)
        self.boot_header = ZynqBootHeader()
        self.image_header_table = ZynqImageHeaderTable()
        self.image = []  # List of tuples (image_header, [partition_headers])

    def verify_magic(self) -> bool:
        """Verify the magic bytes at offset 0x24.
        
        Checks for "XLNX" (0x584c4e58 in little endian) at offset 0x24.
        
        Returns:
            True if the magic bytes are correct, False otherwise.
        
        Raises:
            RuntimeError: If no image data is loaded.
            ValueError: If image data is too small.
        """
        return self._verify_magic_at_offset(0x24)

    def detect_width(self) -> bool:
        """Detect the bus width at offset 0x20.
        
        Detects single or double stacked flash at offset 0x20.
        
        Returns:
            True if the width detection is successful, False otherwise.
        
        Raises:
            RuntimeError: If no image data is loaded.
            ValueError: If image data is too small.
        """
        return self._detect_width_at_offset(0x20)

    @property
    def image_headers(self):
        """Property to access image headers with their partition headers.
        
        Returns:
            List of image headers, each with partition_headers attribute set.
        """
        result = []
        for image_header, partition_headers in self.image:
            # Attach partition_headers to the image_header
            image_header.partition_headers = partition_headers
            result.append(image_header)
        return result

    def parse(self):
        """Parse the Zynq boot image.
        
        Raises:
            RuntimeError: If no image data is loaded.
        """
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        
        # Validate the boot header before reading too much

        # Header Signature at 0x24 (already read by verify_magic)
        if not self.verify_magic():
            raise ValueError(f"Invalid boot image: magic bytes do not match {self.boot_header.header_signature.hex()}")
        
        if not self.detect_width():
            raise ValueError(f"Width detection indicates an unsupported configuration: {self.boot_header.width_detection_word.hex()}") 
        
        # Key Source at 0x28 (4 bytes)
        self._safe_seek(self.image_path_offset + 0x28, 4)
        self.boot_header.key_source = struct.unpack('<I', self.image_fd.read(4))[0]
        if self.boot_header.key_source != 0x00000000:
            raise ValueError("Encryption is not supported")
        
        # Read ARM Vector Table at 0x00-0x1F (32 bytes)
        self._safe_seek(self.image_path_offset + 0x00, 32)
        word1 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word2 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word3 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word4 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word5 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word6 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word7 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        word8 = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
        self.boot_header.arm_vector_table = word1 + word2 + word3 + word4 + word5 + word6 + word7 + word8

        # Header Version at 0x2C (4 bytes)
        self._safe_seek(self.image_path_offset + 0x2C, 4)
        self.boot_header.header_version = self.image_fd.read(4)
        
        # Source Offset at 0x30 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x30, 4)
        self.boot_header.source_offset = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # FSBL Image Length at 0x34 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x34, 4)
        self.boot_header.fsbl_image_length = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # FSBL Load Address at 0x38 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x38, 4)
        self.boot_header.fsbl_load_address = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # FSBL Execution Address at 0x3C (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x3C, 4)
        self.boot_header.fsbl_execution_address = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # Total FSBL Length at 0x40 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x40, 4)
        self.boot_header.total_fsbl_length = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # QSPI Configuration Word at 0x44 (4 bytes)
        self._safe_seek(self.image_path_offset + 0x44, 4)
        self.boot_header.qspi_configuration_word = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # Boot Header Checksum at 0x48 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x48, 4)
        self.boot_header.boot_header_checksum = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # Verify boot header checksum
        if not self.boot_header.verify_checksum():
            computed = self.boot_header.compute_checksum()
            expected = self.boot_header.boot_header_checksum
            raise ValueError(f"Boot header checksum verification failed: computed=0x{computed:08X}, expected=0x{expected:08X}")
        
        # User Defined Fields at 0x4C-0x97 (76 bytes)
        self._safe_seek(self.image_path_offset + 0x4C, 76)
        self.boot_header.user_defined_fields = struct.unpack('>76s', self.image_fd.read(76))[0]
        
        # Image Header Table Offset at 0x98 (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x98, 4)
        self.boot_header.image_header_table_offset = struct.unpack('<I', self.image_fd.read(4))[0]
        
        # Partition Header Table Offset at 0x9C (4 bytes, little-endian)
        self._safe_seek(self.image_path_offset + 0x9C, 4)
        self.boot_header.partition_header_table_offset = struct.unpack('<I', self.image_fd.read(4))[0]

        # Load Image Header Table
        if self.boot_header.image_header_table_offset is not None:
            # Seek to Image Header Table offset
            self._safe_seek(self.image_path_offset + self.boot_header.image_header_table_offset, 20)
            
            # Version (4 bytes)
            self.image_header_table.version = struct.unpack('<I', self.image_fd.read(4))[0]
            
            # Validate version
            if not self.image_header_table.is_version_valid():
                raise ValueError(f"Invalid Image Header Table version: 0x{self.image_header_table.version:08X}")
            
            # Count of Image Headers (4 bytes, little-endian)
            self.image_header_table.count_of_image_headers = struct.unpack('<I', self.image_fd.read(4))[0]
            
            # First Partition Header Offset (4 bytes, little-endian)
            self.image_header_table.first_partition_header_offset = struct.unpack('<I', self.image_fd.read(4))[0]
            
            # First Image Header Offset (4 bytes, little-endian)
            self.image_header_table.first_image_header_offset = struct.unpack('<I', self.image_fd.read(4))[0]
            
            if self.image_header_table.version == ZynqImageHeaderTable.VERSION_1_2:
                # Header Authentication Certificate Offset (4 bytes, little-endian)
                self.image_header_table.header_authentication_certificate_offset = struct.unpack('<I', self.image_fd.read(4))[0]
            else:
                self.image_header_table.header_authentication_certificate_offset = None

        # Load Image Headers and corresponding Partition Headers
        if self.image_header_table.first_image_header_offset is not None:
            current_offset = self.image_header_table.first_image_header_offset * 4 # Word to byte offset
            
            while current_offset != 0 and current_offset != 0xffffffff:
                # Create a new image header instance
                image_header = ZynqImageHeader()
                
                # Seek to current image header offset
                self._safe_seek(self.image_path_offset + current_offset, 20)  # At least 20 bytes needed for header fields
                
                # Next Image Header Offset (4 bytes, little-endian)
                image_header.next_image_header = struct.unpack('<I', self.image_fd.read(4))[0]
                
                # Corresponding Partition Header (4 bytes, little-endian)
                image_header.corresponding_partition_header = struct.unpack('<I', self.image_fd.read(4))[0]
                
                # Reserved (4 bytes)
                _ = self.image_fd.read(4)
                
                # Partition Count (4 bytes, little-endian)
                image_header.partition_count_length = struct.unpack('<I', self.image_fd.read(4))[0]
                
                # Image Name (null-terminated ASCII string)
                image_name_bytes = b''
                while True:
                    word = struct.unpack('>I', self.image_fd.read(4))[0].to_bytes(4, 'little')
                    if word[0:1] == b'\x00' and word[1:2] == b'\x00' and word[2:3] == b'\x00' and word[3:4] == b'\x00':
                        break
                    image_name_bytes += word
                
                image_header.image_name = image_name_bytes.decode('ascii', errors='ignore')
                
                # Load all corresponding partition headers for this image
                partition_headers = []
                partition_header_offset = image_header.corresponding_partition_header * 4  # Convert word to byte offset
                partitions_to_read = image_header.partition_count_length or 1
                
                for _ in range(partitions_to_read):
                    if partition_header_offset == 0:
                        break
                    
                    partition_header = ZynqPartitionHeader()
                    
                    # Seek to partition header location
                    self._safe_seek(self.image_path_offset + partition_header_offset, ZynqPartitionHeader.PARTITION_HEADER_SIZE)
                    
                    # Encrypted Partition Length at 0x00 (4 bytes, little-endian)
                    partition_header.encrypted_partition_length = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Unencrypted Partition Length at 0x04 (4 bytes, little-endian)
                    partition_header.unencrypted_partition_length = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Total Partition Word Length at 0x08 (4 bytes, little-endian)
                    partition_header.total_partition_word_length = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Destination Load Address at 0x0C (4 bytes, little-endian)
                    partition_header.destination_load_address = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Destination Execution Address at 0x10 (4 bytes, little-endian)
                    partition_header.destination_execution_address = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Data Word Offset in Image at 0x14 (4 bytes, little-endian)
                    partition_header.data_word_offset_in_image = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Attribute Bits at 0x18 (4 bytes, little-endian)
                    partition_header.attribute_bits = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Section Count at 0x1C (4 bytes, little-endian)
                    partition_header.section_count = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Checksum Word Offset at 0x20 (4 bytes, little-endian)
                    partition_header.checksum_word_offset = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Image Header Word Offset at 0x24 (4 bytes, little-endian)
                    partition_header.image_header_word_offset = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Authentication Certificate Word Offset at 0x28 (4 bytes, little-endian)
                    partition_header.authentication_certificate_word_offset = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Reserved at 0x2C-0x38 (16 bytes)
                    _ = self.image_fd.read(16)
                    
                    # Header Checksum at 0x3C (4 bytes, little-endian)
                    partition_header.header_checksum = struct.unpack('<I', self.image_fd.read(4))[0]
                    
                    # Verify partition header checksum
                    if not partition_header.verify_checksum():
                        raise ValueError(f"Partition header checksum verification failed for '{image_header.image_name}'")
                    
                    partition_headers.append(partition_header)
                    
                    # Move to next partition header (sequential, 64 bytes apart)
                    partition_header_offset += ZynqPartitionHeader.PARTITION_HEADER_SIZE
                
                # Add tuple to list (image_header, list of partition_headers)
                self.image.append((image_header, partition_headers))
                
                # Move to next image header (0 means this was the last one)
                current_offset = image_header.next_image_header * 4 # Word to byte offset

        return

    def debug(self):
        """Output debug information about the loaded Zynq boot image."""
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        
        print("Zynq Boot Image Header:")
        print(f"  ARM Vector Table: {self.boot_header.arm_vector_table.hex() if self.boot_header.arm_vector_table else 'Not read'}")
        print(f"  Width Detection Word: {self.boot_header.width_detection_word.hex() if self.boot_header.width_detection_word else 'Not read'}")
        print(f"  Header Signature: {self.boot_header.header_signature.hex() if self.boot_header.header_signature else 'Not read'}")
        print(f"  Key Source: {hex(self.boot_header.key_source) if self.boot_header.key_source is not None else 'Not read'}")
        print(f"  Header Version: {hex(struct.unpack('<I', self.boot_header.header_version)[0]) if self.boot_header.header_version else 'Not read'}")
        print(f"  Source Offset: {self.boot_header.source_offset if self.boot_header.source_offset is not None else 'Not read'}")
        print(f"  FSBL Image Length: {self.boot_header.fsbl_image_length if self.boot_header.fsbl_image_length is not None else 'Not read'}")
        print(f"  FSBL Load Address: {self.boot_header.fsbl_load_address if self.boot_header.fsbl_load_address is not None else 'Not read'}")
        print(f"  FSBL Execution Address: {self.boot_header.fsbl_execution_address if self.boot_header.fsbl_execution_address is not None else 'Not read'}")
        print(f"  Total FSBL Length: {self.boot_header.total_fsbl_length if self.boot_header.total_fsbl_length is not None else 'Not read'}")
        print(f"  QSPI Configuration Word: {hex(self.boot_header.qspi_configuration_word) if self.boot_header.qspi_configuration_word else 'Not read'}")
        print(f"  Boot Header Checksum: {hex(self.boot_header.boot_header_checksum) if self.boot_header.boot_header_checksum is not None else 'Not read'}")
        print(f"  User Defined Fields: {self.boot_header.user_defined_fields if self.boot_header.user_defined_fields else 'Not read'}")
        print(f"  Image Header Table Offset: {hex(self.boot_header.image_header_table_offset) if self.boot_header.image_header_table_offset is not None else 'Not read'}")
        print(f"  Partition Header Table Offset: {hex(self.boot_header.partition_header_table_offset) if self.boot_header.partition_header_table_offset is not None else 'Not read'}")
        
        if self.boot_header.image_header_table_offset:
            print("\nZynq Image Header Table:")
            print(f"  Version: {hex(self.image_header_table.version) if self.image_header_table.version else 'Not read'}")
            print(f"  Count of Image Headers: {self.image_header_table.count_of_image_headers if self.image_header_table.count_of_image_headers is not None else 'Not read'}")
            print(f"  First Partition Header Offset: {hex(self.image_header_table.first_partition_header_offset) if self.image_header_table.first_partition_header_offset is not None else 'Not read'}")
            print(f"  First Image Header Offset: {hex(self.image_header_table.first_image_header_offset) if self.image_header_table.first_image_header_offset is not None else 'Not read'}")
            print(f"  Header Authentication Certificate Offset: {self.image_header_table.header_authentication_certificate_offset if self.image_header_table.header_authentication_certificate_offset is not None else 'Not read'}")
        
        # Display Image Headers
        if self.image:
            print(f"\nZynq Image Headers (Total: {len(self.image)}):")
            for idx, (image_header, partition_headers) in enumerate(self.image):
                print(f"\n  Image Header #{idx}:")
                print(f"    Next Image Header: {hex(image_header.next_image_header) if image_header.next_image_header is not None else 'Not read'}")
                print(f"    Corresponding Partition Header: {hex(image_header.corresponding_partition_header) if image_header.corresponding_partition_header is not None else 'Not read'}")
                print(f"    Partition Count Length: {image_header.partition_count_length if image_header.partition_count_length is not None else 'Not read'}")
                print(f"    Image Name: '{image_header.image_name if image_header.image_name else 'Not read'}'")
                
                # Display all Partition Headers for this image
                for part_idx, partition_header in enumerate(partition_headers):
                    print(f"\n  Partition Header #{idx}.{part_idx}:")
                    print(f"    Encrypted Partition Length: {partition_header.encrypted_partition_length if partition_header.encrypted_partition_length is not None else 'Not read'}")
                    print(f"    Unencrypted Partition Length: {partition_header.unencrypted_partition_length if partition_header.unencrypted_partition_length is not None else 'Not read'}")
                    print(f"    Total Partition Word Length: {partition_header.total_partition_word_length if partition_header.total_partition_word_length is not None else 'Not read'}")
                    print(f"    Destination Load Address: {hex(partition_header.destination_load_address) if partition_header.destination_load_address is not None else 'Not read'}")
                    print(f"    Destination Execution Address: {hex(partition_header.destination_execution_address) if partition_header.destination_execution_address is not None else 'Not read'}")
                    print(f"    Data Word Offset in Image: {hex(partition_header.data_word_offset_in_image) if partition_header.data_word_offset_in_image is not None else 'Not read'}")
                    print(f"    Attribute Bits: {hex(partition_header.attribute_bits) if partition_header.attribute_bits is not None else 'Not read'}")
                    print(f"      Owner: {partition_header.get_partition_owner().name}")
                    print(f"      RSA Signature Present: {partition_header.get_rsa_signature_present()}")
                    print(f"      Checksum Type: {partition_header.get_checksum_type().name}")
                    print(f"      Destination Device: {partition_header.get_destination_device().name}")
                    print(f"    Section Count: {partition_header.section_count if partition_header.section_count is not None else 'Not read'}")
                    print(f"    Checksum Word Offset: {hex(partition_header.checksum_word_offset) if partition_header.checksum_word_offset is not None else 'Not read'}")
                    print(f"    Image Header Word Offset: {hex(partition_header.image_header_word_offset) if partition_header.image_header_word_offset is not None else 'Not read'}")
                    print(f"    Authentication Certificate Word Offset: {hex(partition_header.authentication_certificate_word_offset) if partition_header.authentication_certificate_word_offset is not None else 'Not read'}")
                    print(f"    Header Checksum: {hex(partition_header.header_checksum) if partition_header.header_checksum is not None else 'Not read'}")

    def _get_first_stage_loader_length(self, data_byte_offset: int, default_length: int) -> int:
        """Return the length for FSBL partition, or default_length otherwise.
        
        For Zynq, the FSBL (First Stage Boot Loader) partition length is stored
        in the boot header's total_fsbl_length field.
        
        Args:
            data_byte_offset: The data offset of the current partition.
            default_length: The default length from the partition header.
        
        Returns:
            The FSBL length if this is the FSBL partition, otherwise default_length.
        """
        if (self.boot_header.source_offset is not None and 
            data_byte_offset == self.boot_header.source_offset):
            return self.boot_header.total_fsbl_length
        return default_length
