---

#!/usr/bin/env python3
import sys
import subprocess
import os
import re

def main():
    # We expect only ONE argument — the absolute path to zephyr.elf
    if len(sys.argv) < 2:
        print("[ELF LMA Shifter] Error: Please provide path to zephyr.elf")
        sys.exit(1)

    elf_path = sys.argv[1]
    if not os.path.exists(elf_path):
        print(f"[ELF LMA Shifter] Error: {elf_path} not found.")
        sys.exit(1)

    # Locate toolchain binaries relative to the execution environment
    # Standard Zephyr SDK pathing fallback
    sdk_bin_dir = "/home/vovan/ZEPHYR_RTOS/zephyrsdk/zephyr-sdk-0.16.3/arm-zephyr-eabi/bin"
    readelf_bin = os.path.join(sdk_bin_dir, "arm-zephyr-eabi-readelf")
    objcopy_bin = os.path.join(sdk_bin_dir, "arm-zephyr-eabi-objcopy")

    if not os.path.exists(readelf_bin) or not os.path.exists(objcopy_bin):
        print("[ELF LMA Shifter] Error: GNU Toolchain binaries not found in SDK path.")
        sys.exit(1)

    # Step 1: Find the absolute physical end of the CODE region in SDRAM (0x80000000 - 0x803FFFFF)
    code_end_vma = 0x80000000
    res = subprocess.run([readelf_bin, "-SW", elf_path], capture_output=True, text=True, env={"LC_ALL": "C"})
    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 only valid code sections located below the data boundary
                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
    # Align the starting point of the RAM data payload to an MPU-safe 16 KB boundary
    flash_data_start_lma = (code_end_lma + 0x3FFF) & ~0x3FFF
    
    print(f"[ELF LMA Shifter] Target Analysis Complete:")
    print(f"  -> Code region ends at Flash LMA: {hex(code_end_lma)}")
    print(f"  -> Data payload packing starts at aligned LMA: {hex(flash_data_start_lma)}")

    # Step 2: Parse and collect sections based strictly on ELF Flags and Types
    dynamic_offsets = []
    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]+\s+([0-9a-fA-F]+)\s+[0-9a-fA-F]+\s+([AXMSILG01-9\s]+)'
        )
        
        parsed_sections = []
        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)
                
                is_alloc = 'A' in flags_str
                has_contents = (sec_type != "NOBITS") and (size_val > 0)
                is_in_sdram = (0x80000000 <= vma_val < 0x82000000)
                
                # CRITICAL: Only touch real data/code allocated in SDRAM! 
                # Ignore debug symbols, comments, attributes, and unallocated sections.
                if is_alloc and has_contents and is_in_sdram:
                    parsed_sections.append((vma_val, size_val, sec_name))

        # Sort strictly by VMA to prevent internal segment cross-overlapping
        parsed_sections.sort(key=lambda x: x[0])

        # Build targeted LMA change arguments without affecting VMA
        current_flash_lma = flash_data_start_lma
        for vma_val, size_val, sec_name in parsed_sections:
            if vma_val >= 0x80400000:
                dynamic_offsets += ["--change-section-lma", f"{sec_name}={hex(current_flash_lma)}"]
                current_flash_lma += (size_val + 3) & ~3
            else:
                standard_lma = vma_val - 0x20000000
                dynamic_offsets += ["--change-section-lma", f"{sec_name}={hex(standard_lma)}"]

    # Step 3: Apply LMA patches directly to 'zephyr.elf' via a safe temporary file
    if dynamic_offsets:
        print(f"[ELF LMA Shifter] Patching headers for {len(parsed_sections)} SDRAM sections...")
        temp_elf = elf_path + ".tmp"
        
        # objcopy cannot perform in-place modification safely, so we output to a .tmp file
        elf_patch_cmd = [objcopy_bin] + dynamic_offsets + [elf_path, temp_elf]
        patch_res = subprocess.run(elf_patch_cmd, capture_output=True, text=True)
        
        if patch_res.returncode == 0 and os.path.exists(temp_elf):
            # Atomic replacement of the original unshifted ELF
            os.replace(temp_elf, elf_path)
            print("[ELF LMA Shifter] Success: zephyr.elf updated natively.")
        else:
            print(f"[ELF LMA Shifter] Error: objcopy failed to patch ELF. Reason: {patch_res.stderr}")
            if os.path.exists(temp_elf):
                os.remove(temp_elf)
            sys.exit(1)
    else:
        print("[ELF LMA Shifter] Warning: No active SDRAM sections found to shift.")

if __name__ == "__main__":
    main()