---

import sys
import subprocess
import os
import re
import struct

def patch_boot_data_length(bin_path, final_length):
    """
    Patches the NXP Boot Data 'Length' field inside the final application binary.
    This ensures that the internal ROM bootloader copies the exact required 
    amount of bytes from Flash to SDRAM during boot initialization.
    """
    if not os.path.exists(bin_path):
        return
    try:
        with open(bin_path, "r+b") as f:
            # Move pointer to NXP Boot Data 'Length' field (fixed offset 0x24)
            f.seek(0x24)
            # Write exact calculated image size as a 32-bit little-endian integer
            f.write(struct.pack("<I", final_length))
        print(f"[LMA Filter] BootROM Patch: Successfully updated Boot Data Length to {final_length} bytes ({hex(final_length)})")
    except Exception as e:
        print(f"[LMA Filter] BootROM Patch Error: Could not modify binary header. Reason: {e}")

def main():
    if len(sys.argv) < 2:
        print("[LMA Filter] Error: No real objcopy path provided.")
        sys.exit(1)

    real_objcopy = sys.argv[1]
    args = sys.argv[2:]
    elf_path = "zephyr.elf"
    
    for arg in args:
        if arg.endswith("zephyr.elf"):
            elf_path = arg
            break

    # Skip processing if this is an internal ELF-to-ELF adjustment step by Zephyr
    is_elf_to_elf_adjust = any("--change-section-lma" in arg for arg in args) and ("zephyr.elf" in elf_path) and not any(fmt in " ".join(args) for fmt in ["ihex", "binary"])
    if is_elf_to_elf_adjust:
        sys.exit(0)

    # Detect if the current command is generating the final .hex or .bin file
    is_hex_or_bin_generation = any("--output-target=" in arg for arg in args) or any(fmt in " ".join(args) for fmt in ["ihex", "binary"])

    if is_hex_or_bin_generation and os.path.exists(elf_path):
        sdk_bin_dir = os.path.dirname(real_objcopy)
        readelf_bin = os.path.join(sdk_bin_dir, "arm-zephyr-eabi-readelf")
        if not os.path.exists(readelf_bin):
            for file in os.listdir(sdk_bin_dir):
                if file.endswith("readelf"):
                    readelf_bin = os.path.join(sdk_bin_dir, file)
                    break

        # Step 1: Find the absolute physical end of the CODE region (0x80000000 to 0x803FFFFF)
        code_end_vma = 0x80000000
        if os.path.exists(readelf_bin):
            res = subprocess.run([readelf_bin, "-SW", elf_path], capture_output=True, text=True)
            if res.returncode == 0:
                section_regex = re.compile(r'^\s*\[\s*\d+\]\s+([\.\w\-_]+)\s+(\w+)\s+([0-9a-fA-F]{8,16})\s+([0-9a-fA-F]{6,8})\s+([0-9a-fA-F]{6,8})')
                for line in res.stdout.splitlines():
                    match = section_regex.match(line)
                    if match:
                        sec_name, sec_type, vma_str, _, size_str = match.groups()
                        vma_val = int(vma_str, 16)
                        size_val = int(size_str, 16)
                        # Monitor code sections located below the 0x80400000 boundary (4MB Flash limits)
                        if 0x80000000 <= vma_val < 0x80400000 and sec_type != "NOBITS":
                            end_pos = vma_val + size_val
                            if end_pos > code_end_vma:
                                code_end_vma = end_pos

        code_end_lma = code_end_vma - 0x20000000
        
        # MPU-SAFE ALIGNMENT FIX:
        # Align the starting point of the RAM data payload to a strict 16 KB (0x4000) boundary.
        # This prevents Cortex-M7 Memory Protection Unit (MPU) faults during hardware initialization.
        flash_data_start_lma = (code_end_lma + 0x3FFF) & ~0x3FFF
        
        print(f"[LMA Filter] Dynamic MPU-Safe Packing Initiated:")
        print(f"  -> Code region ends at Flash LMA: {hex(code_end_lma)}")
        print(f"  -> Data payload will start at aligned Flash LMA: {hex(flash_data_start_lma)}")

        # Step 2: Build individual LMA adjustments based on Flags, Type and VMA boundaries
        dynamic_offsets = []
        final_payload_end_lma = flash_data_start_lma
        
        if os.path.exists(readelf_bin):
            res = subprocess.run([readelf_bin, "-SW", elf_path], capture_output=True, text=True)
            if res.returncode == 0:
                # Regex captures: Section Name, Type, VMA, Offset, Size, and Flags string
                section_regex = re.compile(
                    r'^\s*\[\s*\d+\]\s+([\.\w\-_]+)\s+(\w+)\s+([0-9a-fA-F]{8,16})\s+[0-9a-fA-F]+\s+([0-9a-fA-F]+)\s+[0-9a-fA-F]+\s+([AXMSILG01-9\s]+)'
                )
                current_flash_lma = flash_data_start_lma
                
                for line in res.stdout.splitlines():
                    match = section_regex.match(line)
                    if match:
                        sec_name, sec_type, vma_str, size_str, flags_str = match.groups()
                        
                        vma_val = int(vma_str, 16)
                        size_val = int(size_str, 16)
                        
                        # Check if the section requires device memory allocation ('A' flag present)
                        is_alloc = 'A' in flags_str
                        # Ensure the section contains actual binary data payload (is not NOBITS)
                        has_contents = (sec_type != "NOBITS") and (size_val > 0)
                        # Ensure the VMA address space belongs to the physical 32MB SDRAM window
                        is_in_sdram = (0x80000000 <= vma_val < 0x82000000)
                        
                        # Strictly skip metadata, debug, comments, and untargeted architectures segments
                        if not (is_alloc and has_contents and is_in_sdram):
                            continue
                        
                        # Split processing into CODE partition (<0x80400000) and DATA partition (>=0x80400000)
                        if vma_val >= 0x80400000:
                            # Pack RAM-destined variables sequentially right after the code block in Flash LMA space.
                            # KEEP VMA INTACT to allow application execution from correct high SDRAM addresses.
                            dynamic_offsets += ["--change-section-lma", f"{sec_name}={hex(current_flash_lma)}"]
                            current_flash_lma += (size_val + 3) & ~3
                            final_payload_end_lma = current_flash_lma
                        else:
                            # Standard code mapping: shift SDRAM execution address (0x80xxxxxx VMA) 
                            # back to physical FlexSPI Flash storage destination (0x60xxxxxx LMA).
                            standard_lma = vma_val - 0x20000000
                            dynamic_offsets += ["--change-section-lma", f"{sec_name}={hex(standard_lma)}"]

        # Run objcopy tool with custom tailored per-section LMA targets
        final_cmd = [real_objcopy] + dynamic_offsets + args
        subprocess.run(final_cmd, check=True)

        # Step 3: Patch the generated binary file for standalone boot validation
        calculated_image_size = (final_payload_end_lma - 0x60000000)
        elf_dir = os.path.dirname(elf_path)
        potential_bin = os.path.join(elf_dir, "zephyr.bin")
        if os.path.exists(potential_bin):
            patch_boot_data_length(potential_bin, calculated_image_size)
        else:
            fallback_bin = os.path.join(os.getcwd(), "zephyr", "zephyr.bin")
            if os.path.exists(fallback_bin):
                patch_boot_data_length(fallback_bin, calculated_image_size)

        # Step 4: Print the final beautifully packed memory map based on exact same ELF flags
        if os.path.exists(readelf_bin):
            res = subprocess.run([readelf_bin, "-SW", elf_path], capture_output=True, text=True)
            if res.returncode == 0:
                print(f"\n[LMA Filter] === VERIFIED DYNAMICALLY PACKED MEMORY MAP FOR FLASH ===")
                print(f"{'Section Name':<28} | {'Execution VMA':<13} | {'Storage LMA (Flash)':<22}")
                print("-" * 72)
                section_regex = re.compile(
                    r'^\s*\[\s*\d+\]\s+([\.\w\-_]+)\s+(\w+)\s+([0-9a-fA-F]{8,16})\s+[0-9a-fA-F]+\s+([0-9a-fA-F]+)\s+[0-9a-fA-F]+\s+([AXMSILG01-9\s]+)'
                )
                current_flash_lma = flash_data_start_lma
                
                for line in res.stdout.splitlines():
                    match = section_regex.match(line)
                    if not match:
                        continue
                    sec_name, sec_type, vma_str, size_str, flags_str = match.groups()
                    vma_val = int(vma_str, 16)
                    size_val = int(size_str, 16)
                    
                    is_alloc = 'A' in flags_str
                    has_contents = (sec_type != "NOBITS") and (size_val > 0)
                    is_in_sdram = (0x80000000 <= vma_val < 0x82000000)
                    
                    if not is_in_sdram:
                        print(f"{sec_name:<28} | 0x{vma_val:08X}    | {'--------':<22} [System / Untouched]")
                        continue
                    if sec_type == "NOBITS":
                        print(f"{sec_name:<28} | 0x{vma_val:08X}    | {'--------':<22} [RAM-only / Ignored]")
                        continue
                    
                    if vma_val >= 0x80400000:
                        print(f"{sec_name:<28} | 0x{vma_val:08X}    | {hex(current_flash_lma)}")
                        current_flash_lma += (size_val + 3) & ~3
                    else:
                        print(f"{sec_name:<28} | 0x{vma_val:08X}    | {hex(vma_val - 0x20000000)}")

if __name__ == "__main__":
    main()