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

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_absolute(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)),            # Load PC with vector
        (0x12345678, False, int(dut.regfile.value)),            # Read first instruction
    ]

    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