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

"""Boot image loader module.

This module is intended to enable basic level of read-only boot image handling.

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

from typing import Optional, BinaryIO
from pathlib import Path
from abc import ABC, abstractmethod
import struct


class BootHeader(ABC):
    """Abstract base class for boot image headers."""
    
    # Magic and identification values
    WIDTH_DETECTION = 0xAA995566
    HEADER_ID_WORD = 0x584C4E58  # "XLNX" in little-endian
    
    def __init__(self) -> None:
        """Initialize the boot header."""
        self.width_detection_word: Optional[bytes] = None
        self.header_signature: Optional[bytes] = None


class ImageHeaderTable(ABC):
    """Abstract base class for image header tables."""
    
    def __init__(self) -> None:
        """Initialize the image header table."""
        pass


class ImageHeader(ABC):
    """Abstract base class for image headers."""
    
    def __init__(self) -> None:
        """Initialize the image header."""
        pass


class PartitionHeader(ABC):
    """Abstract base class for partition headers."""
    
    def __init__(self) -> None:
        """Initialize the partition header."""
        pass


class ImageHeaderWrapper:
    """Wrapper to add partition_headers attribute to image headers.

    This provides compatibility with code that expects image_headers to be a list
    where each item has a partition_headers attribute.

    Uses composition: attribute access is delegated to the wrapped image header
    via __getattr__, so the wrapper transparently exposes all of the original
    header's attributes without copying them.
    """
    def __init__(self, image_header, partition_headers):
        self._image_header = image_header
        self.partition_headers = partition_headers

    def __getattr__(self, name):
        # Called only when normal attribute lookup fails, so self.partition_headers
        # and self._image_header are found via __dict__ without entering here.
        return getattr(self._image_header, name)


class BootImageLoader(ABC):
    """Abstract base class for loading and processing AMD boot images."""

    def __init__(self, image_path: Optional[Path] = None, image_path_offset: int = 0) -> None:
        """Initialize the boot image loader.
        
        Args:
            image_path: Path to the boot image file.
            image_path_offset: Offset within the image file to start reading from.  
        """
        self.image_fd: Optional[BinaryIO] = None
        self.image_path: Optional[Path] = image_path
        self.image_path_offset: int = image_path_offset

    def open(self, image_path: Optional[Path] = None, image_path_offset: Optional[int] = None) -> None:
        """Open a boot image from file for processing.
        
        Args:
            image_path: Path to the boot image file. If not provided,
                       uses the path from initialization.
            image_path_offset: Offset within the image file to start reading from. If not provided,
                               uses the offset from initialization.
        
        Raises:
            FileNotFoundError: If the image file doesn't exist.
            ValueError: If no image path is provided.
        """
        path = image_path or self.image_path
        if path is None:
            raise ValueError("No image path provided")
        
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Boot image not found: {path}")
        
        self.image_fd = open(path, 'rb')
        self.image_path = path
        if image_path_offset is not None:
            self.image_path_offset = image_path_offset

        if not self.verify_magic():
            raise ValueError(f"Invalid boot image: magic bytes do not match {self.boot_header.header_signature.hex()}")
        
    @abstractmethod
    def parse(self) -> None:
        """Parse the boot image."""
        pass

    def close(self) -> None:
        """Close the opened boot image file."""
        if self.image_fd:
            self.image_fd.close()
            self.image_fd = None
            self.boot_header.header_signature = None
            self.boot_header.width_detection_word = None

    def __enter__(self):
        """Enter the context manager, opening the boot image.
        
        Returns:
            self: The loader instance with the image opened.
        """
        self.open()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit the context manager, closing the boot image.
        
        Args:
            exc_type: Exception type if an exception was raised.
            exc_val: Exception value if an exception was raised.
            exc_tb: Exception traceback if an exception was raised.
        
        Returns:
            False: Always propagate exceptions.
        """
        self.close()
        return False

    def _safe_seek(self, offset: int, bytes_to_read: int = 0) -> bool:
        """Safely seek to an offset, checking bounds.
        
        Args:
            offset: Absolute offset to seek to.
            bytes_to_read: Number of bytes that will be read after seeking.
        
        Returns:
            True if seek is safe, False otherwise.
        
        Raises:
            IOError: If seek would go past end of file.
        """
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        # Get file size
        current_pos = self.image_fd.tell()
        self.image_fd.seek(0, 2)  # Seek to end
        file_size = self.image_fd.tell()
        self.image_fd.seek(current_pos)  # Restore position
        
        # Check bounds
        if offset < 0:
            raise IOError(f"Invalid seek offset: {offset} (negative)")
        if offset > file_size:
            raise IOError(f"Seek offset {offset} exceeds file size {file_size}")
        if bytes_to_read > 0 and (offset + bytes_to_read) > file_size:
            raise IOError(f"Read at offset {offset} for {bytes_to_read} bytes would exceed file size {file_size}")
        
        self.image_fd.seek(offset)
        return True

    @abstractmethod
    def verify_magic(self) -> bool:
        """Verify the magic bytes at the specified offset.
        
        Checks for "XLNX" (0x584c4e58 in little endian) at the given offset.
        
        Returns:
            True if the magic bytes are correct, False otherwise.
        """
        pass

    def _verify_magic_at_offset(self, offset: int) -> bool:
        """Verify the magic bytes at the specified offset.
        
        Helper method for subclasses to use in their verify_magic() implementation.
        Checks for "XLNX" (0x584c4e58 in little endian) at the given offset.
        
        Args:
            offset: Byte offset to read the magic bytes from.
        
        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.
        """
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        
        self.image_fd.seek(self.image_path_offset + offset)
        raw_data = self.image_fd.read(4)
        
        if len(raw_data) < 4:
            raise ValueError("Image data too small to contain magic bytes.")
        
        self.boot_header.header_signature = raw_data
        
        # Verify raw file bytes match the little-endian representation of HEADER_ID_WORD
        expected_magic = struct.pack('<I', self.boot_header.HEADER_ID_WORD)
        
        return self.boot_header.header_signature == expected_magic
    
    @abstractmethod
    def detect_width(self) -> bool:
        """Detect the bus width of the boot image.
        
        Returns:
            True if the width detection is successful, False otherwise.
        """
        pass

    def _detect_width_at_offset(self, offset: int) -> bool:
        """Detect single or double stacked flash.
        
        Helper method for subclasses to use in their detect_width() implementation.
        
        Args:
            offset: Byte offset to read the width information from.
        
        Returns:
            True if the width is single (the bytes match), False otherwise.
        
        Raises:
            RuntimeError: If no image data is loaded.
            ValueError: If image data is too small.
        """
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        
        self.image_fd.seek(self.image_path_offset + offset)
        raw_data = self.image_fd.read(4)
        
        if len(raw_data) < 4:
            raise ValueError("Image data too small to contain width information.")
        
        self.boot_header.width_detection_word = raw_data
        
        # Compare raw file bytes against the little-endian representation of WIDTH_DETECTION
        width_value = struct.pack('<I', self.boot_header.WIDTH_DETECTION)
        
        return self.boot_header.width_detection_word == width_value
    
    @abstractmethod
    def debug(self) -> None:
        """Output debug information about the loaded boot image."""
        pass

    def extract(self, image_index: int, partition_index: int = 0, filename: Optional[str] = None) -> bool:
        """Extract a partition from the boot image and write it to a file.
        
        Args:
            image_index: The image index to extract from.
            partition_index: The partition index within the image to extract (default: 0).
            filename: Optional filename to write the extracted partition to.
                     If not provided, uses the image name from the corresponding image header.
        
        Returns:
            True if the extraction was successful, False otherwise.
        
        Raises:
            RuntimeError: If no image data is loaded.
            IndexError: If the image or partition index is out of range.
        """
        if not self.image_fd:
            raise RuntimeError("No image data loaded. Call open() first.")
        
        if image_index < 0 or image_index >= len(self.image):
            raise IndexError(f"Image index {image_index} out of range (0-{len(self.image) - 1})")
        
        image_header, partition_headers = self.image[image_index]
        
        if partition_index < 0 or partition_index >= len(partition_headers):
            raise IndexError(f"Partition index {partition_index} out of range (0-{len(partition_headers) - 1})")
        
        partition_header = partition_headers[partition_index]
        
        # Determine output filename
        output_filename = self._get_output_filename(filename, image_header, partition_index)
        
        # Get byte offset and length
        data_byte_offset = partition_header.get_data_offset_bytes()
        data_length = partition_header.get_unencrypted_partition_length_bytes()
        
        # Device-specific override for first-stage loader partition
        data_length = self._get_first_stage_loader_length(data_byte_offset, data_length)
        
        # Seek and extract
        self._safe_seek(self.image_path_offset + data_byte_offset, data_length)
        self._write_partition_to_file(output_filename, data_length)
        
        return True

    def _get_output_filename(self, filename: Optional[str], image_header, partition_index: int) -> str:
        """Generate the output filename for extraction.
        
        Args:
            filename: User-provided filename, or None to use image header name.
            image_header: The image header containing the default image name.
            partition_index: The partition index within the image.
        
        Returns:
            The output filename to use for extraction.
        
        Raises:
            ValueError: If no filename can be determined.
        """
        output_filename = filename if filename else image_header.image_name
        if not output_filename:
            raise ValueError("No filename provided and image has no image name")
        output_filename = output_filename.replace('\x00', '')
        
        # For sub-partitions (partition_index > 0), append the partition index to the filename
        if partition_index > 0 and not filename:
            if '.' in output_filename:
                name, ext = output_filename.rsplit('.', 1)
                output_filename = f"{name}.{partition_index}.{ext}"
            else:
                output_filename = f"{output_filename}.{partition_index}"
        
        return output_filename

    @abstractmethod
    def _get_first_stage_loader_length(self, data_byte_offset: int, default_length: int) -> int:
        """Return the length for first-stage loader partition, or default_length otherwise.
        
        Override in subclasses to check for FSBL (Zynq/ZynqMP) or PLM (Versal family).
        
        Args:
            data_byte_offset: The data offset of the current partition.
            default_length: The default length from the partition header.
        
        Returns:
            The length to use for extraction (may be overridden for first-stage loader).
        """
        pass

    def _write_partition_to_file(self, filename: str, length: int) -> None:
        """Write partition data to file in chunks.
        
        Args:
            filename: The output filename to write to.
            length: The number of bytes to write.
        """
        chunk_size = 4096
        bytes_remaining = length
        with open(filename, 'wb') as output_fd:
            while bytes_remaining > 0:
                write_size = min(chunk_size, bytes_remaining)
                chunk = self.image_fd.read(write_size)
                if not chunk:
                    raise IOError(f"Unexpected end of file while extracting partition to '{filename}' "
                                  f"({length - bytes_remaining} of {length} bytes read)")
                output_fd.write(chunk)
                bytes_remaining -= len(chunk)

