import cocotb
from cocotb.handle import LogicArray

from cocotb.clock import Clock
from cocotb.triggers import ReadOnly, NextTimeStep, RisingEdge

import logging

import random


logger = logging.getLogger()

logger.setLevel(logging.INFO)

CLK_PERIOD = 5


SETS = 64
WAYS = 4

data_arrays = [{}, {}, {}, {}]
meta_arrays = [{}, {}, {}, {}]

@cocotb.test
async def test_sanity(dut):
    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())

    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)


    for index in range(SETS):
        for way in range(WAYS):
            data = random.randbytes(64)
            meta = random.randint(0, 2**22-1)

            data_arrays[way][index] = data
            meta_arrays[way][index] = meta

            dut.i_cpu_write_data.value = LogicArray.from_bytes(data, byteorder="little")
            dut.i_cpu_write_meta.value = meta
            dut.i_cpu_write_index.value = index
            dut.i_cpu_write_valid.value = 1 << way

            await RisingEdge(dut.i_clk)

    dut.i_cpu_write_valid.value = 0

    await RisingEdge(dut.i_clk)

    for index in range(SETS):
        dut.i_cpu_read_index.value = index
        dut.i_cpu_read_valid.value = 1

        await RisingEdge(dut.i_clk)
        await ReadOnly()

        raw_data = dut.o_cpu_read_data.value
        raw_meta = dut.o_cpu_read_meta.value

        meta = [int(m) for m in raw_meta]
        data_bytes = [v.to_bytes(byteorder="little") for v in raw_data]

        expected_data = [data_arrays[way][index] for way in range(WAYS)]
        expected_meta = [meta_arrays[way][index] for way in range(WAYS)]

        if data_bytes != expected_data:
            logger.info("Data Error")

        if meta != expected_meta:
            logger.info("Meta Error")
            
        await NextTimeStep()

    dut.i_cpu_read_valid.value = 0

    for index in range(SETS):
        for way in range(WAYS):
            data = random.randbytes(64)
            meta = random.randint(0, 2**22-1)

            data_arrays[way][index] = data
            meta_arrays[way][index] = meta

            dut.i_snoop_write_data.value = LogicArray.from_bytes(data, byteorder="little")
            dut.i_snoop_write_meta.value = meta
            dut.i_snoop_write_index.value = index
            dut.i_snoop_write_valid.value = 1 << way

            await RisingEdge(dut.i_clk)

    dut.i_snoop_write_valid.value = 0

    await RisingEdge(dut.i_clk)

    for index in range(SETS):
        dut.i_snoop_read_index.value = index
        dut.i_snoop_read_valid.value = 1

        await RisingEdge(dut.i_clk)
        await ReadOnly()

        raw_data = dut.o_snoop_read_data.value
        raw_meta = dut.o_snoop_read_meta.value

        meta = [int(m) for m in raw_meta]
        data_bytes = [v.to_bytes(byteorder="little") for v in raw_data]

        expected_data = [data_arrays[way][index] for way in range(WAYS)]
        expected_meta = [meta_arrays[way][index] for way in range(WAYS)]

        if data_bytes != expected_data:
            logger.info("Data Error")

        if meta != expected_meta:
            logger.info("Meta Error")
            
        await NextTimeStep()

    dut.i_snoop_read_valid.value = 0


@cocotb.test
async def test_random_access(dut):
    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())

    ITERS = 1024

    for _ in range(ITERS):
        cpu_write_way = random.randint(0, WAYS-1)
        cpu_write_set = random.randint(0, SETS-1)

        while True:
            snoop_write_way = random.randint(0, WAYS-1)
            snoop_write_set = random.randint(0, SETS-1)
            if snoop_write_way != cpu_write_way and snoop_write_set != cpu_write_set:
                break


        cpu_write_data = random.randbytes(64)
        cpu_write_meta = random.randint(0, 2**22-1)

        snoop_write_data = random.randbytes(64)
        snoop_write_meta = random.randint(0, 2**22-1)

        data_arrays[cpu_write_way][cpu_write_set] = cpu_write_data
        meta_arrays[cpu_write_way][cpu_write_set] = cpu_write_meta

        data_arrays[snoop_write_way][snoop_write_set] = snoop_write_data
        meta_arrays[snoop_write_way][snoop_write_set] = snoop_write_meta

        dut.i_cpu_write_data.value = LogicArray.from_bytes(cpu_write_data, byteorder="little")
        dut.i_cpu_write_meta.value = cpu_write_meta
        dut.i_cpu_write_index.value = cpu_write_set
        dut.i_cpu_write_valid.value = 1 << cpu_write_way

        dut.i_snoop_write_data.value = LogicArray.from_bytes(snoop_write_data, byteorder="little")
        dut.i_snoop_write_meta.value = snoop_write_meta
        dut.i_snoop_write_index.value = snoop_write_set
        dut.i_snoop_write_valid.value = 1 << snoop_write_way



        cpu_read_way = random.randint(0, WAYS-1)
        cpu_read_set = random.randint(0, SETS-1)

        snoop_read_way = random.randint(0, WAYS-1)
        snoop_read_set = random.randint(0, SETS-1)

        dut.i_cpu_read_index.value = cpu_read_set
        dut.i_snoop_read_index.value = snoop_read_set


        dut.i_cpu_read_valid.value = 1
        dut.i_snoop_read_valid.value = 1

        await RisingEdge(dut.i_clk)
        await ReadOnly()

        cpu_data = dut.o_cpu_read_data.value[cpu_read_way].to_bytes(byteorder="little")
        cpu_meta = int(dut.o_cpu_read_meta.value[cpu_read_way])

        snoop_data = dut.o_snoop_read_data.value[snoop_read_way].to_bytes(byteorder="little")
        snoop_meta = int(dut.o_snoop_read_meta.value[snoop_read_way])

        cpu_expected_data = data_arrays[cpu_read_way][cpu_read_set]
        cpu_expected_meta = meta_arrays[cpu_read_way][cpu_read_set]

        snoop_expected_data = data_arrays[snoop_read_way][snoop_read_set]
        snoop_expected_meta = meta_arrays[snoop_read_way][snoop_read_set]

        if cpu_data != cpu_expected_data:
            logger.error("CPU Data Error")

        if cpu_meta != cpu_expected_meta:
            logger.info("CPU Meta Error")

        if snoop_data != snoop_expected_data:
            logger.error("snoop Data Error")

        if snoop_meta != snoop_expected_meta:
            logger.info("snoop Meta Error")

        await NextTimeStep()