from ida_segment import *
from ida_bytes import *


def create_peripheral_segment(start_addr, name, size):
    seg = segment_t()
    seg.start_ea = start_addr
    seg.end_ea = start_addr + size
    seg.bitness = 1  # 32-bit
    seg.align = saRelByte
    seg.comb = scPub
    seg.perm = SEGPERM_READ | SEGPERM_WRITE

    if add_segm_ex(seg, name, "PERIPHERAL", ADDSEG_OR_DIE):
        getseg(start_addr).type = SEG_DATA
        set_segment_cmt(getseg(start_addr), "Volatile peripheral registers", True)
        return True
    return False


def verify_no_overlaps(peripherals):
    sorted_pairs = sorted(peripherals.items())
    for i in range(len(sorted_pairs) - 1):
        curr_addr, (curr_name, curr_size) = sorted_pairs[i]
        next_addr, (next_name, _) = sorted_pairs[i + 1]

        if curr_addr + curr_size > next_addr:
            print(f"ERROR: Overlap between {curr_name} and {next_name}")
            return False
    return True


def main():
    peripherals = {
        0x40000000: ("SYSTEM_CTRL", 0x200),
        0x40000200: ("PERI_ON", 0x80),
        0x40000280: ("PINMUX_REG", 0xD80),
        0x40001000: ("GPIO_REG", 0x1000),
        0x40002000: ("TIMER_REG", 0x800),
        0x40002800: ("VENDOR_REG", 0x800),
        0x40003000: ("LOG_UART_REG", 0x400),
        0x40003400: ("RTC", 0x800),
        0x40003C00: ("SPIC_CACHE", 0x400),
        0x40010000: ("ADC_REG", 0x1000),
        0x40020000: ("SPI_FLASH_CTRL", 0x1000),
        0x40040000: ("UART0_REG", 0x400),
        0x40040400: ("UART1_REG", 0x400),
        0x40042000: ("SPI0_REG", 0x400),
        0x40042400: ("SPI1_REG", 0x400),
        0x40044000: ("I2C0_REG", 0x400),
        0x40044400: ("I2C1_REG", 0x400),
        0x40050000: ("SDIO_DEVICE_REG", 0x1000),
        0x40060000: ("GDMA0_REG", 0x1000),
        0x40061000: ("GDMA1_REG", 0x1000),
        0x40062000: ("I2S0_REG", 0x1000),
        0x40070000: ("CRYPTO_REG", 0x1000),
        0x40080000: ("WIFI_REG", 0x40000),
        0x400C0000: ("SIE_REG", 0x2000),
        0x400C2000: ("USOC_REG", 0x2000),
    }

    if not verify_no_overlaps(peripherals):
        return

    for addr, (name, size) in sorted(peripherals.items()):
        if not create_peripheral_segment(addr, name, size):
            print(f"Failed to create segment {name}")


if __name__ == "__main__":
    main()