class Memory:

    def __init__(self, scheduler: Scheduler, gamepak: GamePak, bios: bytes | bytearray | Iterable[int]):
        self.scheduler = scheduler
        self.cpu: CPU | None = None
        self.io: IO | None = None
        self.bios = array('B', bios.ljust(MemoryRegion.BIOS_SIZE))
        self.ewram = array('B', [0] * MemoryRegion.EWRAM_SIZE)
        self.iwram = array('B', [0] * MemoryRegion.IWRAM_SIZE)
        self.gamepak = gamepak
        self.bios_last_opcode = 0
        self.wait_control = WaitstateControlRegister()
        self._init_access_times()

    def read_32(self, address: int, access_type: MemoryAccess) -> int:
        return self._read_32_internal(address, access_type)

    def read_32_ror(self, address: int, access_type: MemoryAccess) -> int:
        value = self.read_32(address, access_type)
        rotate = (address & 3) * 8
        return ror_32(value, rotate)

    def read_16(self, address: int, access_type: MemoryAccess) -> int:
        return self._read_16_internal(address, access_type)

    def read_16_signed(self, address: int, access_type: MemoryAccess) -> int:
        if get_bit(address, 0):
            return self.read_8_signed(address, access_type)
        return extend_sign_16(self.read_16(address, access_type))

    def read_16_ror(self, address: int, access_type: MemoryAccess) -> int:
        value = self.read_16(address, access_type)
        rotate = (address & 1) * 8
        return ror_32(value, rotate)

    def read_8(self, address: int, access_type: MemoryAccess) -> int:
        return self._read_8_internal(address, access_type)

    def read_8_signed(self, address: int, access_type: MemoryAccess) -> int:
        value = self.read_8(address, access_type)
        return extend_sign_8(value)

    def write_32(self, address: int, value: int, access_type: MemoryAccess):
        self._write_32_internal(address, value, access_type)

    def write_16(self, address: int, value: int, access_type: MemoryAccess):
        self._write_16_internal(address, value, access_type)

    def write_8(self, address: int, value: int, access_type: MemoryAccess):
        self._write_8_internal(address, value, access_type)

    def _read_32_internal(self, address: int, access_type: MemoryAccess) -> int:
        self._idle_for_access(address, 4, access_type)
        address = address & ~3
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            if address <= MemoryRegion.BIOS_END:
                if self.cpu.regs.get_pc() <= MemoryRegion.BIOS_END:
                    self.bios_last_opcode = array_read_32(self.bios, address & MemoryRegion.BIOS_MASK)
                return self.bios_last_opcode
            return self._read_unused_memory()
        elif region == MemoryRegion.EWRAM_REGION:
            return array_read_32(self.ewram, address & MemoryRegion.EWRAM_MASK)
        elif region == MemoryRegion.IWRAM_REGION:
            return array_read_32(self.iwram, address & MemoryRegion.IWRAM_MASK)
        elif MemoryRegion.IO_START <= address <= MemoryRegion.IO_END:
            return self.io.read_32(address)
        elif region == MemoryRegion.PALRAM_REGION:
            return self.io.ppu.memory.read_32_palram(address)
        elif region == MemoryRegion.VRAM_REGION:
            return self.io.ppu.memory.read_32_vram(address)
        elif region == MemoryRegion.OAM_REGION:
            return self.io.ppu.memory.read_32_oam(address)
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            return self.gamepak.read_32(address)
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to read SRAM: {address:#010x}')
            return 0
        else:
            return self._read_unused_memory()

    def _read_16_internal(self, address: int, access_type: MemoryAccess) -> int:
        self._idle_for_access(address, 2, access_type)
        address = address & ~1
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            if address <= MemoryRegion.BIOS_END:
                if self.cpu.regs.get_pc() <= MemoryRegion.BIOS_END:
                    self.bios_last_opcode = array_read_32(self.bios, address & ~3 & MemoryRegion.BIOS_MASK)
                return self.bios_last_opcode >> (address & 3) * 8 & 65535
            return self._read_unused_memory() >> (address & 3) * 8 & 65535
        elif region == MemoryRegion.EWRAM_REGION:
            return array_read_16(self.ewram, address & MemoryRegion.EWRAM_MASK)
        elif region == MemoryRegion.IWRAM_REGION:
            return array_read_16(self.iwram, address & MemoryRegion.IWRAM_MASK)
        elif MemoryRegion.IO_START <= address <= MemoryRegion.IO_END:
            return self.io.read_16(address)
        elif region == MemoryRegion.PALRAM_REGION:
            return self.io.ppu.memory.read_16_palram(address)
        elif region == MemoryRegion.VRAM_REGION:
            return self.io.ppu.memory.read_16_vram(address)
        elif region == MemoryRegion.OAM_REGION:
            return self.io.ppu.memory.read_16_oam(address)
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            return self.gamepak.read_16(address)
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to read SRAM: {address:#010x}')
            return 0
        else:
            return self._read_unused_memory() >> (address & 3) * 8 & 65535

    def _read_8_internal(self, address: int, access_type: MemoryAccess) -> int:
        self._idle_for_access(address, 1, access_type)
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            if address <= MemoryRegion.BIOS_END:
                if self.cpu.regs.get_pc() <= MemoryRegion.BIOS_END:
                    self.bios_last_opcode = array_read_32(self.bios, address & ~3 & MemoryRegion.BIOS_MASK)
                return self.bios_last_opcode >> (address & 3) * 8 & 255
            return self._read_unused_memory() >> (address & 3) * 8 & 255
        elif region == MemoryRegion.EWRAM_REGION:
            return self.ewram[address & MemoryRegion.EWRAM_MASK]
        elif region == MemoryRegion.IWRAM_REGION:
            return self.iwram[address & MemoryRegion.IWRAM_MASK]
        elif MemoryRegion.IO_START <= address <= MemoryRegion.IO_END:
            return self.io.read_8(address)
        elif region == MemoryRegion.PALRAM_REGION:
            return self.io.ppu.memory.read_8_palram(address)
        elif region == MemoryRegion.VRAM_REGION:
            return self.io.ppu.memory.read_8_vram(address)
        elif region == MemoryRegion.OAM_REGION:
            return self.io.ppu.memory.read_8_oam(address)
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            return self.gamepak.read_8(address)
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to read SRAM: {address:#010x}')
            return 0
        else:
            return self._read_unused_memory() >> (address & 3) * 8 & 255

    def _write_32_internal(self, address: int, value: int, access_type: MemoryAccess):
        self._idle_for_access(address, 4, access_type)
        address = address & ~3
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            pass
        elif region == MemoryRegion.EWRAM_REGION:
            array_write_32(self.ewram, address & MemoryRegion.EWRAM_MASK, value)
        elif region == MemoryRegion.IWRAM_REGION:
            array_write_32(self.iwram, address & MemoryRegion.IWRAM_MASK, value)
        elif region == MemoryRegion.IO_REGION:
            self.io.write_32(address, value)
        elif region == MemoryRegion.PALRAM_REGION:
            self.io.ppu.memory.write_32_palram(address, value)
        elif region == MemoryRegion.VRAM_REGION:
            self.io.ppu.memory.write_32_vram(address, value)
        elif region == MemoryRegion.OAM_REGION:
            self.io.ppu.memory.write_32_oam(address, value)
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            print(f'Attempt to write to SRAM: {address:#010x}')
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to write to SRAM: {address:#010x}')
        else:
            print(f'Attempt to write to unused memory: {address:#010x}')

    def _write_16_internal(self, address: int, value: int, access_type: MemoryAccess):
        self._idle_for_access(address, 2, access_type)
        address = address & ~1
        value = value & 65535
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            pass
        elif region == MemoryRegion.EWRAM_REGION:
            array_write_16(self.ewram, address & MemoryRegion.EWRAM_MASK, value)
        elif region == MemoryRegion.IWRAM_REGION:
            array_write_16(self.iwram, address & MemoryRegion.IWRAM_MASK, value)
        elif region == MemoryRegion.IO_REGION:
            self.io.write_16(address, value)
        elif region == MemoryRegion.PALRAM_REGION:
            self.io.ppu.memory.write_16_palram(address, value)
        elif region == MemoryRegion.VRAM_REGION:
            self.io.ppu.memory.write_16_vram(address, value)
        elif region == MemoryRegion.OAM_REGION:
            self.io.ppu.memory.write_16_oam(address, value)
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            print(f'Attempt to write to SRAM: {address:#010x}')
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to write to SRAM: {address:#010x}')
        else:
            print(f'Attempt to write to unused memory: {address:#010x}')

    def _write_8_internal(self, address: int, value: int, access_type: MemoryAccess):
        self._idle_for_access(address, 1, access_type)
        value = value & 255
        region = address >> 24
        if region == MemoryRegion.BIOS_REGION:
            pass
        elif region == MemoryRegion.EWRAM_REGION:
            self.ewram[address & MemoryRegion.EWRAM_MASK] = value
        elif region == MemoryRegion.IWRAM_REGION:
            self.iwram[address & MemoryRegion.IWRAM_MASK] = value
        elif region == MemoryRegion.IO_REGION:
            self.io.write_8(address, value)
        elif region == MemoryRegion.PALRAM_REGION:
            self.io.ppu.memory.write_8_palram(address, value)
        elif region == MemoryRegion.VRAM_REGION:
            self.io.ppu.memory.write_8_vram(address, value)
        elif region == MemoryRegion.OAM_REGION:
            pass
        elif MemoryRegion.GAMEPAK_REGION_START <= region <= MemoryRegion.GAMEPAK_REGION_END:
            print(f'Attempt to write to SRAM: {address:#010x}')
        elif region == MemoryRegion.SRAM_REGION:
            print(f'Attempt to write to SRAM: {address:#010x}')
        else:
            print(f'Attempt to write to unused memory: {address:#010x}')

    def _read_unused_memory(self) -> int:
        """
        Reading from unused memory returns the last prefetched instruction.
        https://problemkaputt.de/gbatek.htm#gbaunpredictablethings

        For THUMB code, the result consists of two 16-bit fragments and depends on
        the address area and alignment where the opcode was stored.
        """
        if self.cpu.regs.cpsr.get_state() == CPUState.ARM:
            return self.cpu.pipeline[1]
        region = self.cpu.regs.get_pc() >> 24
        if region == MemoryRegion.BIOS_REGION or region == MemoryRegion.OAM_REGION:
            if self.cpu.regs.get_pc() & 3 == 0:
                return self.cpu.pipeline[1] << 16 | self.cpu.pipeline[1]
            else:
                return self.cpu.pipeline[1] << 16 | self.cpu.pipeline[0]
        elif region == MemoryRegion.IWRAM_REGION:
            if self.cpu.regs.get_pc() & 3 == 0:
                return self.cpu.pipeline[0] << 16 | self.cpu.pipeline[1]
            else:
                return self.cpu.pipeline[1] << 16 | self.cpu.pipeline[0]
        else:
            return self.cpu.pipeline[1] << 16 | self.cpu.pipeline[1]

    def _init_access_times(self):
        """
        Source: GBATEK

        Region        Bus   Read      Write     Cycles
        =====================================================
        BIOS ROM      32    8/16/32   -         1/1/1
        Work RAM 32K  32    8/16/32   8/16/32   1/1/1
        I/O           32    8/16/32   8/16/32   1/1/1
        OAM           32    8/16/32   16/32     1/1/1 *
        Work RAM 256K 16    8/16/32   8/16/32   3/3/6 **
        Palette RAM   16    8/16/32   16/32     1/1/2 *
        VRAM          16    8/16/32   16/32     1/1/2 *
        GamePak ROM   16    8/16/32   -         5/5/8 **/***
        GamePak Flash 16    8/16/32   16/32     5/5/8 **/***
        GamePak SRAM  8     8         8         5     **

        Timing Notes:

          *   Plus 1 cycle if GBA accesses video memory at the same time.
          **  Default waitstate settings, see System Control chapter.
          *** Separate timings for sequential, and non-sequential accesses.

        """
        self.access_time_32 = [[1] * 16 for _ in range(2)]
        self.access_time_16 = [[1] * 16 for _ in range(2)]
        self.access_time_32[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.EWRAM_REGION)] = 6
        self.access_time_32[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.EWRAM_REGION)] = 6
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.EWRAM_REGION)] = 3
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.EWRAM_REGION)] = 3
        self.access_time_32[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.PALRAM_REGION)] = 2
        self.access_time_32[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.PALRAM_REGION)] = 2
        self.access_time_32[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.VRAM_REGION)] = 2
        self.access_time_32[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.VRAM_REGION)] = 2
        self.access_time_32[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.OAM_REGION)] = 2
        self.access_time_32[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.OAM_REGION)] = 2
        self.update_waitstates()

    def update_waitstates(self):
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_0_REGION_1)] = 1 + self.wait_control.get_ws0_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_0_REGION_2)] = 1 + self.wait_control.get_ws0_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_1_REGION_1)] = 1 + self.wait_control.get_ws1_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_1_REGION_2)] = 1 + self.wait_control.get_ws1_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_2_REGION_1)] = 1 + self.wait_control.get_ws2_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.GAMEPAK_2_REGION_2)] = 1 + self.wait_control.get_ws2_non_seq()
        self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(MemoryRegion.SRAM_REGION)] = 1 + self.wait_control.get_sram()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_0_REGION_1)] = 1 + self.wait_control.get_ws0_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_0_REGION_2)] = 1 + self.wait_control.get_ws0_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_1_REGION_1)] = 1 + self.wait_control.get_ws1_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_1_REGION_2)] = 1 + self.wait_control.get_ws1_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_2_REGION_1)] = 1 + self.wait_control.get_ws2_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.GAMEPAK_2_REGION_2)] = 1 + self.wait_control.get_ws2_seq()
        self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(MemoryRegion.SRAM_REGION)] = 1 + self.wait_control.get_sram()
        for region in range(MemoryRegion.GAMEPAK_0_REGION_1, MemoryRegion.SRAM_REGION + 1):
            self.access_time_32[int(MemoryAccess.NON_SEQUENTIAL)][int(region)] = self.access_time_16[int(MemoryAccess.NON_SEQUENTIAL)][int(region)] + self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(region)]
            self.access_time_32[int(MemoryAccess.SEQUENTIAL)][int(region)] = self.access_time_16[int(MemoryAccess.SEQUENTIAL)][int(region)] * 2

    def _idle_for_access(self, address: int, size: int, access_type: MemoryAccess):
        region = address >> 24 & 15
        cycles = self.access_time_32[int(access_type)][int(region)] if size == 4 else self.access_time_16[int(access_type)][int(region)]
        self.scheduler.idle(cycles)

class WaitstateControlRegister:

    def __init__(self):
        self.NON_SEQUENTIAL_CYCLES = [4, 3, 2, 8]
        self.reg = 0

    def get_sram(self):
        return self.NON_SEQUENTIAL_CYCLES[get_bits(self.reg, 0, 1)]

    def get_ws0_non_seq(self):
        return self.NON_SEQUENTIAL_CYCLES[get_bits(self.reg, 2, 3)]

    def get_ws0_seq(self):
        return 1 if get_bit(self.reg, 4) else 2

    def get_ws1_non_seq(self):
        return self.NON_SEQUENTIAL_CYCLES[get_bits(self.reg, 5, 6)]

    def get_ws1_seq(self):
        return 1 if get_bit(self.reg, 7) else 4

    def get_ws2_non_seq(self):
        return self.NON_SEQUENTIAL_CYCLES[get_bits(self.reg, 8, 9)]

    def get_ws2_seq(self):
        return 1 if get_bit(self.reg, 10) else 8