import sys
from makeelf.elf import ELF, EM, ELFDATA, ELFCLASS


def create_elf(output_path, segments, entry_point):
    elf = ELF(
        e_machine=EM.EM_XTENSA, e_data=ELFDATA.ELFDATA2LSB, e_class=ELFCLASS.ELFCLASS32
    )
    elf.Elf.Ehdr.e_entry = entry_point

    for name, addr, data_path in segments:
        with open(data_path, "rb") as f:
            data = f.read()
        elf.append_section(name, data, addr)

    print(f"Writing ELF to {output_path}...")
    with open(output_path, "wb") as f:
        f.write(bytes(elf))


if __name__ == "__main__":
    ultra_segments = [
        (
            ".flash.rodata",
            0x3C0E0020,
            "/home/sapient/esp32-analysis/extracted/seg0_DROM.bin",
        ),
        (
            ".dram0.data",
            0x3FC96440,
            "/home/sapient/esp32-analysis/extracted/seg1_DRAM.bin",
        ),
        (
            ".iram0.text",
            0x40374000,
            "/home/sapient/esp32-analysis/extracted/seg2_IRAM.bin",
        ),
        (
            ".flash.text",
            0x42000020,
            "/home/sapient/esp32-analysis/extracted/seg3_IROM.bin",
        ),
    ]
    create_elf(
        "/home/sapient/esp32-analysis/output_dumps/Ultra-TDeck-v9.2.elf",
        ultra_segments,
        0x4037736C,
    )

    meshos_segments = [
        (
            ".flash.rodata",
            0x3C190020,
            "/home/sapient/esp32-analysis/MeshOS_1.1.8/extracted/seg0_DROM.bin",
        ),
        (
            ".dram0.data",
            0x3FC98F50,
            "/home/sapient/esp32-analysis/MeshOS_1.1.8/extracted/seg1_DRAM.bin",
        ),
        (
            ".iram0.text",
            0x40374000,
            "/home/sapient/esp32-analysis/MeshOS_1.1.8/extracted/seg2_IRAM.bin",
        ),
        (
            ".flash.text",
            0x42000020,
            "/home/sapient/esp32-analysis/MeshOS_1.1.8/extracted/seg3_IROM.bin",
        ),
    ]
    create_elf(
        "/home/sapient/esp32-analysis/output_dumps/MeshOS-TDeck-1.1.8.elf",
        meshos_segments,
        0x40377A10,
    )