import cocotb
from cocotb.handle import Immediate

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

from collections import defaultdict

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)

@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(0xfffffff4, 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
        (0xfffffff4, False, int(dut.regfile.value)),            # read vector byte 0
        (0xfffffff5, False, int(dut.regfile.value)),            # read vector byte 1
        (0xfffffff6, False, int(dut.regfile.value)),            # read vector byte 2
        (0xfffffff7, 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
    ]

    for expected_output in expected_cpu_outputs:
        await RisingEdge(dut.clk)

        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
        assert dut_do == expected_do

@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(0xfffffff4, 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
    ]

    for expected_output in expected_cpu_outputs:
        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_absolute_x(dut):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    write_dword(0xfffffff4, 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
    ]

    for expected_output in expected_cpu_outputs:
        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