import cocotb
from cocotb.handle import Immediate

from cocotb.clock import Clock
from cocotb.triggers import Timer, RisingEdge, FallingEdge

from collections import defaultdict

import struct
import random

CLK_PERIOD = 5

memory = defaultdict(int)

def write_dword(addr: int, data: int):
    memory[addr + 0] = (data >>  0) & 0xff
    memory[addr + 1] = (data >>  8) & 0xff
    memory[addr + 2] = (data >> 16) & 0xff
    memory[addr + 3] = (data >> 24) & 0xff

def write_byte(addr: int, data: int):
    memory[addr] = data & 0xff

def write_bytes(addr: int, data: bytes| list[int]):
    for i, val in enumerate(data):
        memory[addr + i] = int(val)

async def handle_memory(dut):
    while True:
        await RisingEdge(dut.clk)
        addr = int(dut.AB.value)
        we = bool(dut.WE.value)

        dut.DI.value = memory[addr]

        if we:
            memory[addr] = int(dut.DO.value)

async def check_instruction_sequence(dut, instruction_sequence):
    for expected_output in instruction_sequence:
        await RisingEdge(dut.clk)

        if expected_output:
            expected_addr, expected_we, expected_do = expected_output
            dut_addr = int(dut.AB.value)
            dut_we = bool(dut.WE.value)
            dut_do = int(dut.DO.value)

            assert dut_addr == expected_addr
            assert dut_we == expected_we

            if dut_we:
                assert dut_do == expected_do

@cocotb.test
async def test_reset(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x12345678)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        (0x00000100, True,  (int(dut.PC.value) >> 24) & 0xff),  # High addr
        (0x000001ff, True,  (int(dut.PC.value) >> 16) & 0xff),  # Mid high addr
        (0x000001fe, True,  (int(dut.PC.value) >>  8) & 0xff),  # Mid low addr
        (0x000001fd, True,  (int(dut.PC.value) >>  0) & 0xff),  # Low addr
        (0x000001fc, True,  int(dut.P.value)),                  # Status
        (0xfffffff8, False, int(dut.regfile.value)),            # read vector byte 0
        (0xfffffff9, False, int(dut.regfile.value)),            # read vector byte 1
        (0xfffffffa, False, int(dut.regfile.value)),            # read vector byte 2
        (0xfffffffb, False, int(dut.regfile.value)),            # read vector byte 3
        (0x12345678, False, int(dut.regfile.value)),            # Read first instruction
        (0x12345679, False, int(dut.regfile.value)),            # Read second byte
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_absolute(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # lda $abcd1234
    # sta $50515253
    # wai
    write_bytes(0x200, [0xad, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x205, [0x8d, 0x53, 0x52, 0x51, 0x50])
    write_byte(0x20a, 0xcb)
    write_byte(0xabcd1234, 0x55)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # Read first instruction
        (0x00000201, False, None),            # Read address byte 0
        (0x00000202, False, None),            # Read address byte 1
        (0x00000203, False, None),            # Read address byte 2
        (0x00000204, False, None),            # Read address byte 3
        (0xabcd1234, False, None),            # Read from absolute address
        (0x00000205, False, None),            # Read second instruction
        (0x00000206, False, None),            # Read address byte 0
        (0x00000207, False, None),            # Read address byte 1
        (0x00000208, False, None),            # Read address byte 2
        (0x00000209, False, None),            # Read address byte 3
        (0x50515253, True, 0x55),             # Write to absolute address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_x(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # ldx #1
    # lda $abcd1234,x
    # inx
    # sta $01020304,x
    # wai
    write_bytes(0x200, [0xa2, 0x01])
    write_bytes(0x202, [0xbd, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x207, [0xe8])
    write_bytes(0x208, [0x9d, 0x04, 0x03, 0x02, 0x01])

    write_byte(0x20d, 0xcb)
    write_byte(0xabcd1235, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # ldx $abcd1234,x
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xabcd1235, False, None),            # Read from address
        (0x00000207, False, None),            # inx
        (0x00000208, False, None),            # sta $01020304,x
        (0x00000208, False, None),            # store reg
        (0x00000209, False, None),            # addr 0
        (0x0000020a, False, None),            # addr 1
        (0x0000020b, False, None),            # addr 2
        (0x0000020c, False, None),            # addr 3
        (0x01020306, False, None),            # Write to address
        (0x01020306, True, 0xaa),             # Write to address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_y(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # ldy #1
    # lda $abcd1234,y
    # iny
    # sta $01020304,y
    # wai
    write_bytes(0x200, [0xa0, 0x01])
    write_bytes(0x202, [0xb9, 0x34, 0x12, 0xcd, 0xab])
    write_bytes(0x207, [0xc8])
    write_bytes(0x208, [0x99, 0x04, 0x03, 0x02, 0x01])

    write_byte(0x20d, 0xcb)
    write_byte(0xabcd1235, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # ldx $abcd1234,x
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xabcd1235, False, None),            # Read from address
        (0x00000207, False, None),            # inx
        (0x00000208, False, None),            # sta $01020304,x
        (0x00000208, False, None),            # store reg
        (0x00000209, False, None),            # addr 0
        (0x0000020a, False, None),            # addr 1
        (0x0000020b, False, None),            # addr 2
        (0x0000020c, False, None),            # addr 3
        (0x01020306, False, None),            # Write to address
        (0x01020306, True, 0xaa),             # Write to address
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_absolute_x_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # ldx #1
    # jmp ($deadbeef,x)
    write_bytes(0x200, [0xa2, 0x01])
    write_bytes(0x202, [0x7c, 0xef, 0xbe, 0xad, 0xde])
    write_byte(0xbeefb055, 0xcb)

    write_dword(0xdeadbeef + 1, 0xbeefb055)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #1
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # jmp ($deadbeef,x)
        (0x00000203, False, None),            # addr 0
        (0x00000204, False, None),            # addr 1
        (0x00000205, False, None),            # addr 2
        (0x00000206, False, None),            # addr 3
        (0xdeadbef0, False, None),            # addr 0
        (0xdeadbef1, False, None),            # addr 1
        (0xdeadbef2, False, None),            # addr 2
        (0xdeadbef3, False, None),            # addr 3
        (0xbeefb055, False, None),            # target
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)


@cocotb.test
async def test_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # jmp ($deadbeef)
    write_bytes(0x200, [0x6c, 0xef, 0xbe, 0xad, 0xde])
    write_byte(0xbeefb055, 0xcb)

    write_dword(0xdeadbeef, 0xbeefb055)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # jmp ($deadbeef)
        (0x00000201, False, None),            # addr 0
        (0x00000202, False, None),            # addr 1
        (0x00000203, False, None),            # addr 2
        (0x00000204, False, None),            # addr 3
        (0xdeadbeef, False, None),            # addr 0
        (0xdeadbef0, False, None),            # addr 1
        (0xdeadbef1, False, None),            # addr 2
        (0xdeadbef2, False, None),            # addr 3
        (0xbeefb055, False, None),            # target
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_indexed_indirect(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # ldy #1
    # lda ($04),y
    # iny
    # inc
    # sta ($04),y
    write_bytes(0x200, [0xa0, 0x02])
    write_bytes(0x202, [0xb1, 0x04])
    write_bytes(0x204, [0xc8])
    write_bytes(0x205, [0x1a])
    write_bytes(0x206, [0x91, 0x04])
    write_byte(0x208, 0xcb)

    write_dword(0x04, 0xabcdbeef)
    write_dword(0xabcdbeef+2, 0xaa)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldy #0
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # lda ($04),y
        (0x00000203, False, None),            # ZP index
        (0x00000004, False, None),            # zp addr 0
        (0x00000005, False, None),            # zp addr 1
        (0x00000006, False, None),            # zp addr 2
        (0x00000007, False, None),            # zp addr 3
        (0xabcdbef1, False, None),            # fetch data
        (0x00000204, False, None),            # iny
        (0x00000205, False, None),            # iny
        (0x00000205, False, None),            # inc
        (0x00000206, False, None),            # inc
        (0x00000206, False, None),            # sta ($04),y
        (0x00000207, False, None),            # ZP index
        (0x00000004, False, None),            # zp addr 0
        (0x00000005, False, None),            # zp addr 1
        (0x00000006, False, None),            # zp addr 2
        (0x00000007, False, None),            # zp addr 3
        (0xabcdbef2, False, None),            # store takes extra cycle
        (0xabcdbef2, True, 0xAB),             # write data
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_indirect_indexed(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)


    # ldx #4
    # lda ($04,x)
    # inc
    # ldx #8
    # sta ($04,x)
    write_bytes(0x200, [0xa2, 0x04])
    write_bytes(0x202, [0xa1, 0x04])
    write_bytes(0x204, [0x1a])
    write_bytes(0x205, [0xa2, 0x08])
    write_bytes(0x207, [0x81, 0x04])
    write_byte(0x209, 0xcb)

    write_dword(0x08, 0xfeedf00d)
    write_dword(0x0c, 0xf00d600d)

    write_byte(0xfeedf00d, 0x69)

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),            # ldx #4
        (0x00000201, False, None),            # Immediate
        (0x00000202, False, None),            # lda ($04,x)
        (0x00000203, False, None),            # ZP index
        (0x00000004, False, None),            # Compute ZP index
        (0x00000008, False, None),            # zp addr 0
        (0x00000009, False, None),            # zp addr 1
        (0x0000000a, False, None),            # zp addr 2
        (0x0000000b, False, None),            # zp addr 3
        (0xfeedf00d, False, None),            # fetch data
        (0x00000204, False, None),            # iny
        (0x00000205, False, None),            # iny
        (0x00000205, False, None),            # inc
        (0x00000206, False, None),            # inc
        (0x00000207, False, None),            # sta ($04),y
        (0x00000208, False, None),            # ZP index
        (0x00000004, False, None),            # Compute ZP index
        (0x0000000c, False, None),            # zp addr 0
        (0x0000000d, False, None),            # zp addr 1
        (0x0000000e, False, None),            # zp addr 2
        (0x0000000f, False, None),            # zp addr 3
        (0xf00d600d, True, 0x6a),             # write data
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_jsr(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)

    # @0x200
    # ldx #$0
    # txs
    # jsr $12345678
    # wai
    #
    # @0x1234
    # rts
    write_bytes(0x200, [0xa2, 0xff, 0x9a, 0x20, 0x78, 0x56, 0x34, 0x12, 0xcb])
    write_bytes(0x12345678, [0x60])

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),  # ldx #$00
        (0x00000201, False, None),  # Immediate
        (0x00000202, False, None),  # txs
        (0x00000203, False, None),  # second cycle of txs
        (0x00000203, False, None),  # jsr $12345678
        (0x00000204, False, None),  # first byte of address
        (0x000001ff, True,  0x00),  # 24-31
        (0x000001fe, True,  0x00),  # 16-23
        (0x000001fd, True,  0x02),  # 8-15
        (0x000001fc, True,  0x05),  # 7-0
        (0x00000205, False, None),  # second byte of address
        (0x00000206, False, None),  # third byte of address
        (0x00000207, False, None),  # fourth byte of address
        (0x00000208, False, None),  # receive last byte of address
        (0x12345678, False, None),  # rts
        (0x12345679, False, None),  # rts
        (0x000001fb, False, None),  # current stack while we add 1 to it
        (0x000001fc, False, None),  # 7-0
        (0x000001fd, False, None),  # 15-8
        (0x000001fe, False, None),  # 23-16
        (0x000001ff, False, None),  # 31-24
        (0x1234567c, False, None),  # Updating PC before jump
        (0x00000208, False, None),  # WAI
        (0x00000209, False, None),  # second wai
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_rti(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)
    write_dword(0xfffffffc, 0x300)


    # @0x200
    # ldx #$ff
    # txs
    # brk
    # wai
    # @0x300
    # rti

    write_bytes(0x200, [0xa2, 0xff, 0x9a, 0x00, 0x00, 0xcb]) # BRK is technically a 2 byte instruction
    write_bytes(0x300, [0x40])

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00000200, False, None),  # ldx #$ff
        (0x00000201, False, None),  # immediate
        (0x00000202, False, None),  # txs
        (0x00000203, False, None),  # txs
        (0x00000203, False, None),  # brk
        (0x00000204, False, None),  # brk
        (0x000001ff, True,  0x00),  # brk   31-24
        (0x000001fe, True,  0x00),  # brk   13-16
        (0x000001fd, True,  0x02),  # brk   15-08
        (0x000001fc, True,  0x05),  # brk   07-00
        (0x000001fb, True,  0xb4),  # brk   flags
        (0xfffffffc, False, None),  # vector
        (0xfffffffd, False, None),  # vector
        (0xfffffffe, False, None),  # vector
        (0xffffffff, False, None),  # vector 
        (0x00000300, False, None),  # rti
        (0x00000301, False, None),  # rti
        (0x000001fa, False, None),  # rti
        (0x000001fb, False, None),  # rti
        (0x000001fc, False, None),  # rti
        (0x000001fd, False, None),  # rti
        (0x000001fe, False, None),  # rti
        (0x000001ff, False, None),  # rti
        (0x00000205, False, None),  # wai
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_irq(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)
    write_dword(0xfffffffc, 0x300)

    # @0x200
    # cli
    # wai
    # wai
    # 0x300
    # rti

    write_bytes(0x200, [0x58, 0xcb, 0xcb])
    write_bytes(0x300, [0x40])

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    await FallingEdge(dut.RDY_O)

    await RisingEdge(dut.clk)
    await RisingEdge(dut.clk)

    dut.IRQ.value = 1

    while True:
        await RisingEdge(dut.clk)
        if int(dut.state.value) == 0x08:
            break

    dut.IRQ.value = 0

    await Timer(300, "ns")

    assert int(dut.RDY_O.value) == 0

@cocotb.test
async def test_bra_always(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0xfffffd)

    # L1:   bra L1  0x80 0x03
    #       nop
    #       nop
    # L2:   bra L2

    write_bytes(0xfffffd, [0x80, 0x3, 0x00, 0x00, 0x00, 0x80, 0xfe])

    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    count = 0

    async def count_address():
        nonlocal count
        while True:
            await RisingEdge(dut.clk)
            if int(dut.AB.value) == 0x01000004:
                count+=1

    cocotb.start_soon(count_address())

    await Timer(1, "us")

    assert count > 0

@cocotb.test
async def test_bra_never(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0xfffffd)

    # lda #$00
    # bne <anywhere>
    # wai

    write_bytes(0xfffffd, [0xa9, 0x00, 0xd0, 0x7f, 0xd0, 0x80, 0xd0, 0xfe, 0xcb])


    dut.RDY.value = Immediate(1)

    dut.reset.value = Immediate(1)
    for _ in range(10):
        await RisingEdge(dut.clk)
    dut.reset.value = 0

    expected_cpu_outputs = [
        None,   # ignore reset sequence
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (0x00fffffd, False, None),
        (0x00fffffe, False, None),
        (0x00ffffff, False, None),
        (0x01000000, False, None),
        (0x01000001, False, None),
        (0x01000002, False, None),
        (0x01000003, False, None),
        (0x01000004, False, None),
        (0x01000005, False, None),
        (0x01000006, False, None),
        (0x01000006, False, None),
        (0x01000006, False, None),
    ]

    await check_instruction_sequence(dut, expected_cpu_outputs)

@cocotb.test
async def test_adc(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff8, 0x200)


    def zp_indirect():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffffff)

        address_le = struct.pack("<I", address)

        # lda #$0
        # adc ($00)
        # wai
        write_bytes(0x200, [0xa9, 0x00, 0x72, 0x00, 0xcb])

        write_bytes(0x00, address_le)
        write_byte(address, value)

        return value
    
    def absolute():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffffff)

        address_le = struct.pack("<I", address)

        # lda $#0
        # adc $address
        # wai

        write_bytes(0x200, [0xa9, 0x00, 0x6d])
        write_bytes(0x203, address_le)
        write_bytes(0x207, [0xcb])

        write_byte(address, value)

        return value

    def indirect_x():
        value = random.randint(0,255)
        address = random.randint(0x300, 0xffffff00)
        x_offset = random.randint(0,250)

        address_le = struct.pack("<I", address)

        # lda #$0
        # ldx #x_offset
        # adc ($00,x)
        # wai

        write_bytes(0x200, [0xa9, 0x00, 0xa2, x_offset, 0x61, 0x00, 0xcb])
        write_bytes(x_offset, address_le)
        write_byte(address, value)

        return value

    tests = [zp_indirect, absolute, indirect_x]

    for test in tests:
        value = test()

        dut.RDY.value = Immediate(1)

        dut.reset.value = Immediate(1)
        for _ in range(10):
            await RisingEdge(dut.clk)
        dut.reset.value = 0

        await FallingEdge(dut.RDY_O)

        assert value == int(dut.A.value)