import logging


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

from cocotbext.axi import AxiStreamBus, AxiStreamSource

CLK_PERIOD = 4


class TB:
    def __init__(self, dut):
        self.dut = dut

        self.log = logging.getLogger("cocotb.tb")
        self.log.setLevel(logging.INFO)

        cocotb.start_soon(Clock(self.dut.i_clk, CLK_PERIOD, units="ns").start())

        self.s_data_axis = AxiStreamSource(AxiStreamBus.from_prefix(dut, ""), dut.i_clk, dut.i_rst)

    async def cycle_reset(self):
        await self._cycle_reset(self.dut.i_rst, self.dut.i_clk)

    async def _cycle_reset(self, rst, clk):
        rst.setimmediatevalue(0)
        await RisingEdge(clk)
        await RisingEdge(clk)
        rst.value = 1
        await RisingEdge(clk)
        await RisingEdge(clk)
        rst.value = 0
        await RisingEdge(clk)
        await RisingEdge(clk)

@cocotb.test
async def test_sanity(dut):
    tb = TB(dut)

    await tb.cycle_reset()

    s = 0x1bf54941aff6bf4afdb20dfb8a800301
    r = 0xa806d542fe52447f336d555778bed685
    r_masked = 0x0806d5400e52447c036d555408bed685

    result = 0xa927010caf8b2bc2c6365130c11d06a8

    msg = b"Cryptographic Forum Research Group"


    tb.dut.i_otk.value = ((r << 128) | s)
    tb.dut.i_otk_valid.value = 1
    await RisingEdge(tb.dut.i_clk)
    tb.dut.i_otk_valid.value = 0
    await RisingEdge(tb.dut.i_clk)

    dut_s = tb.dut.u_dut.poly1305_s.value.integer
    dut_r = tb.dut.u_dut.poly1305_r.value.integer

    assert dut_s == s
    assert dut_r == r_masked

    await tb.s_data_axis.send(msg)

    await RisingEdge(tb.dut.o_tag_valid)
    tag = tb.dut.o_tag.value.integer
    
    tb.log.info(f"tag: {tag:x}")

    assert tag == result

    await Timer(1, "us")