import cocotb

import logging

from chacha_helpers import chacha_block, QR

import struct

import random


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

CLK_PERIOD = 4

CONSTANT = [0x65787061, 0x6e642033, 0x322d6279, 0x7465206b]


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 = [1, 0, 0, 0,
               0, 0, 0, 0,
               0, 0, 0, 0,
               0, 0, 0, 0]
    
    # data[0], data[4], data[8],  data[12] = QR(data[0], data[4], data[8],  data[12])
    # data[1], data[5], data[9],  data[13] = QR(data[1], data[5], data[9],  data[13])
    # data[2], data[6], data[10], data[14] = QR(data[2], data[6], data[10], data[14])
    # data[3], data[7], data[11], data[15] = QR(data[3], data[7], data[11], data[15])

    data[0], data[5], data[10], data[15] = QR(data[0], data[5], data[10], data[15])
    data[1], data[6], data[11], data[12] = QR(data[1], data[6], data[11], data[12])
    data[2], data[7], data[8],  data[13] = QR(data[2], data[7], data[8],  data[13])
    data[3], data[4], data[9],  data[14] = QR(data[3], data[4], data[9],  data[14])

    for val in data:
        tb.log.info(f"{val:x}")

    # tb.log.info(data_out)
    
    tb.dut.i_ready.value = 1
    
    tb.dut.i_state.value = 1
    # tb.dut.i_state.value = 0x657870616e642033322d62797465206b

    tb.dut.i_valid.value = 1
    await RisingEdge(tb.dut.i_clk)
    tb.dut.i_valid.value = 1

    await Timer(1, "us")