import cocotb
from cocotb.handle import Immediate

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

from collections import defaultdict

import struct
import random

import os

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 do_asm_test(dut, filename):
    cocotb.start_soon(Clock(dut.clk, CLK_PERIOD, unit="ns").start())
    cocotb.start_soon(handle_memory(dut))

    path = os.path.dirname(os.path.abspath(__file__))

    base_addr = 0xfffff000

    with open(f"{path}/asm_source/{filename}", "rb") as file:
        for i, val in enumerate(file.read()):
            write_byte(base_addr+i, val)

    dut.RDY.value = Immediate(1)

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

    await with_timeout(FallingEdge(dut.RDY_O), 10, "us")

    assert memory[0] == 1

@cocotb.test
async def test_lda(dut):
    await do_asm_test(dut, "lda_test")

@cocotb.test
async def test_jsr(dut):
    await do_asm_test(dut, "jsr_test")