import logging
import cocotb
from Vai import VaiDriver, VaiReceiver, VaiMonitor
from cocotb.clock import Clock
from cocotb.queue import Queue
from cocotb.triggers import RisingEdge, Timer
from Crypto.Cipher import AES
import vsc


# Reset coroutine
async def reset_dut(reset_n, duration_ns):
    reset_n.value = 0
    await Timer(duration_ns, units="ns")
    reset_n.value = 1


# Stimuli model class
@vsc.randobj
class constraints():
    def __init__(self):
        self.key = vsc.rand_bit_t(128)
        self.data = vsc.rand_bit_t(128)

    @vsc.constraint
    def c(self):
        self.data >= 0 and self.data <= 2**128-1
        vsc.dist(self.key, [
            vsc.weight(0,  25),
            vsc.weight((1,2**128-2), 50),
            vsc.weight((2**128-1), 25)])


# Stimuli covergroup
@vsc.covergroup
class covergroup():
    def __init__(self):
        self.with_sample(
            mode = vsc.bit_t(1),
            key = vsc.bit_t(128)
        )

        self.enc = vsc.coverpoint(self.mode, bins=dict(
            enc = vsc.bin(0)))

        self.dec = vsc.coverpoint(self.mode, bins=dict(
            dec = vsc.bin(1)))

        self.key0 = vsc.coverpoint(self.key, bins=dict(
            key0 = vsc.bin(0)))

        self.keyF = vsc.coverpoint(self.key, bins=dict(
            keyF = vsc.bin(2**128-1)))

        self.encXkey0 = vsc.cross([self.enc, self.key0])
        self.encXkeyF = vsc.cross([self.enc, self.keyF])

        self.decXkey0 = vsc.cross([self.dec, self.key0])
        self.decXkeyF = vsc.cross([self.dec, self.keyF])


async def cg_sample(cg, queue):
    while True:
        _data = await queue.get()
        cg.sample(_data[0].value, _data[1].value)


@cocotb.test(skip=False)
async def test_aes_enc(dut):
    """ Test AES encryption """

    clkedge = RisingEdge(dut.clk_i)

    # Connect reset
    reset = dut.reset_i

    _input = [dut.mode_i, dut.key_i, dut.data_i]
    _output = dut.data_o
    # DUT input side
    vai_driver = VaiDriver(dut.clk_i, _input, dut.valid_i, dut.accept_o)
    vai_in_queue = cocotb.queue.Queue()
    vai_in_monitor = VaiMonitor(dut.clk_i, _input, dut.valid_i, dut.accept_o, vai_in_queue)
    # DUT output side
    vai_receiver = VaiReceiver(dut.clk_i, dut.data_o, dut.valid_o, dut.accept_i)
    vai_out_monitor = VaiMonitor(dut.clk_i, _output, dut.valid_o, dut.accept_i)

    cr = constraints()
    cg = covergroup()
    cocotb.start_soon(cg_sample(cg, vai_in_queue))

    # Drive input defaults (setimmediatevalue to avoid x asserts)
    dut.mode_i.setimmediatevalue(0)
    dut.key_i.setimmediatevalue(0)
    dut.data_i.setimmediatevalue(0)
    dut.valid_i.setimmediatevalue(0)
    dut.accept_i.setimmediatevalue(0)

    clock = Clock(dut.clk_i, 10, units="ns")  # Create a 10 ns period clock
    cocotb.start_soon(clock.start())  # Start the clock

    # Execution will block until reset_dut has completed
    dut._log.info("Hold reset")
    await reset_dut(reset, 100)
    dut._log.info("Released reset")

    # Test 10 AES calculations
    for i in range(10):
        # Get now random stimuli
        cr.randomize()
        _key = cr.key
        _data = cr.data
        await clkedge
        # Drive AES inputs
        await vai_driver.send([0, _key, _data])
        # Calc reference data
        _aes = AES.new(_key.to_bytes(16, 'big'), AES.MODE_ECB)
        _ref = _aes.encrypt(_data.to_bytes(16, 'big'))
        # Get DUT output data
        _rec = await vai_receiver.receive()
        # Equivalence check
        assert _rec.buff == _ref, \
            f"Encrypt error, got 0x{_rec.buff.hex()}, expected 0x{_ref.hex()}"


@cocotb.test(skip=False)
async def test_aes_dec(dut):
    """ Test AES decryption """

    clkedge = RisingEdge(dut.clk_i)

    # Connect reset
    reset = dut.reset_i

    _input = [dut.mode_i, dut.key_i, dut.data_i]
    _output = dut.data_o
    # DUT input side
    vai_driver = VaiDriver(dut.clk_i, _input, dut.valid_i, dut.accept_o)
    vai_in_queue = cocotb.queue.Queue()
    vai_in_monitor = VaiMonitor(dut.clk_i, _input, dut.valid_i, dut.accept_o, vai_in_queue)
    # DUT output side
    vai_receiver = VaiReceiver(dut.clk_i, dut.data_o, dut.valid_o, dut.accept_i)
    vai_out_monitor = VaiMonitor(dut.clk_i, _output, dut.valid_o, dut.accept_i)

    cr = constraints()
    cg = covergroup()
    cocotb.start_soon(cg_sample(cg, vai_in_queue))

    # Drive input defaults (setimmediatevalue to avoid x asserts)
    dut.mode_i.setimmediatevalue(0)
    dut.key_i.setimmediatevalue(0)
    dut.data_i.setimmediatevalue(0)
    dut.valid_i.setimmediatevalue(0)
    dut.accept_i.setimmediatevalue(0)

    clock = Clock(dut.clk_i, 10, units="ns")  # Create a 10 ns period clock
    cocotb.start_soon(clock.start())  # Start the clock

    # Execution will block until reset_dut has completed
    await reset_dut(reset, 100)
    dut._log.info("Released reset")

    # Test 10 AES calculations
    for i in range(10):
        # Get now random stimuli
        cr.randomize()
        _key = cr.key
        _data = cr.data
        await clkedge
        # Drive AES inputs
        await vai_driver.send([1, _key, _data])
        # Calc reference data
        _aes = AES.new(_key.to_bytes(16, 'big'), AES.MODE_ECB)
        _ref = _aes.decrypt(_data.to_bytes(16, 'big'))
        # Get DUT output data
        _rec = await vai_receiver.receive()
        # Equivalence check
        assert _rec.buff == _ref, \
            f"Decrypt error, got 0x{_rec.buff.hex()}, expected 0x{_ref.hex()}"

    with open('results/tb_aes_fcover.txt', 'w', encoding='utf-8') as f:
        f.write(vsc.get_coverage_report())