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

import random

PRIME = 2**130-5

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

    value_a = random.randint(1,2**(130+16))

    # value_a = PRIME + 1000000
    tb.dut.i_valid.value = 1
    tb.dut.i_val.value = value_a
    await RisingEdge(tb.dut.i_clk)
    tb.dut.i_valid.value = 0
    tb.dut.i_val.value = 0

    await RisingEdge(tb.dut.o_valid)
    value = tb.dut.o_result.value.integer

    print(value_a % PRIME)
    print(value)

    await Timer(1, "us")