import cocotb

import logging

from chacha_helpers import chacha_block

import struct


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

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())


    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()

    data_in = [0x65787061, 0x6e642033, 0x322d6279, 0x7465206b,
               0, 0, 0, 0,
               0, 0, 0, 0,
               0, 0, 0, 0]
    
    data_out = chacha_block(data_in)

    tb.dut.i_counter.value = 0
    tb.dut.i_nonce.value = 0
    tb.dut.i_key.value = 0
    tb.dut.i_ready.value = 1
    tb.dut.i_valid.value = 1
    await RisingEdge(tb.dut.i_clk)
    tb.dut.i_valid.value = 0

    await RisingEdge(tb.dut.o_valid)
    state = tb.dut.o_state.value.integer

    state_bytes = int.to_bytes(state, 64, "little")
    state_words = struct.unpack("16I", state_bytes)

    fail = False
    for i, (sim_val, dut_val) in enumerate(zip(data_out, state_words)):
        if sim_val != dut_val:
            tb.log.info(f"{i}: {sim_val:x} -> {dut_val:x}")
            fail = True

    assert not fail