# test_my_design.py (simple)

import cocotb
from cocotb.triggers import Timer, RisingEdge, FallingEdge, ClockCycles, Join, First
from cocotb.clock import Clock
from random import getrandbits
import struct
import itertools

class FrameData:
    def __init__(self):
        self.dat = []

async def fake_usb_sender(dut, expected_data):
    dut.DATA_COMPLETE.value = 0
    dut.RDADDR.value = 0
    # wait for the rising edge of READY before we can send any data
    await RisingEdge(dut.DATA_READY)

    # grab the number of bytes to send, and send them all
    # we just pretend to send a giant packet even though USB bulk 
    # wouldn't allow that
    num_bytes = int(dut.BUFFER_LENGTH.value)
    dut._log.info("USB sending {} bytes".format(num_bytes))
    read_address = 0
    all_data = []
    await RisingEdge(dut.RDCLK) # wait to clock the 0th data
    all_data.append(int(dut.RDDATA.value))
    dut.RDADDR.value = 0x7FF & (read_address + 1) 
    await RisingEdge(dut.RDCLK) # extra clock to get to 1st data
    read_address = read_address + 1
    dut.RDADDR.value = 0x7FF & (read_address + 1) 
    while(len(all_data) < num_bytes):
        await RisingEdge(dut.RDCLK) 
        all_data.append(int(dut.RDDATA.value))
        read_address = read_address + 1
        dut.RDADDR.value = 0x7FF & (read_address + 1) 
        assert dut.DATA_READY.value == 1, "DATA_READY fell before end of USB transmission"
    # Now that we're done, make sure the finish handshake works
    await RisingEdge(dut.RDCLK)
    dut.DATA_COMPLETE.value = 1
    await FallingEdge(dut.DATA_READY)
    await RisingEdge(dut.RDCLK)
    dut.DATA_COMPLETE.value = 0

    dut._log.info("USB data complete handshake finished")

    # return the sent data
    return all_data

async def simple_usb_loop_sender(dut):
    # Just handles the handshake pins
    # doesn't actually send the data
    USB_SEND_WAIT_TIME = 500
    dut.DATA_COMPLETE.value = 0
    dut.RDADDR.value = 0
    while True:
        # wait for the rising edge of READY 
        await RisingEdge(dut.DATA_READY)
        # Wait a few clocks to pretend we're sending data
        await ClockCycles(dut.RDCLK, USB_SEND_WAIT_TIME)
        # Perform handshake
        dut.DATA_COMPLETE.value = 1
        await FallingEdge(dut.DATA_READY)
        dut.DATA_COMPLETE.value = 0

async def complete_usb_loop_sender(dut, frame_data):
    dut.DATA_COMPLETE.value = 0
    dut.RDADDR.value = 0
    while True:
        # wait for the rising edge of READY 
        await RisingEdge(dut.DATA_READY)
        # figure out how many bytes to send 
        num_bytes = int(dut.BUFFER_LENGTH.value)
        dut._log.info("USB sending {} bytes".format(num_bytes))
        read_address = 0
        all_data = []
        await RisingEdge(dut.RDCLK) # wait to clock the 0th data
        all_data.append(int(dut.RDDATA.value))
        dut.RDADDR.value = 0x7FF & (read_address + 1) 
        await RisingEdge(dut.RDCLK) # extra clock to get to 1st data
        read_address = read_address + 1
        dut.RDADDR.value = 0x7FF & (read_address + 1) 
        while(len(all_data) < num_bytes):
            await RisingEdge(dut.RDCLK) 
            all_data.append(int(dut.RDDATA.value))
            read_address = read_address + 1
            dut.RDADDR.value = 0x7FF & (read_address + 1) 
            assert dut.DATA_READY.value == 1, "DATA_READY fell before end of USB transmission"
        dut.RDADDR.value = 0 # keep this zeroed for the start of the next buffer
        # copy the received data to our frame buffer
        frame_data.dat.extend(all_data)
        # Now that we're done, make sure the finish handshake works
        await RisingEdge(dut.RDCLK)
        dut.DATA_COMPLETE.value = 1
        await FallingEdge(dut.DATA_READY)
        await RisingEdge(dut.RDCLK)
        dut.DATA_COMPLETE.value = 0
        dut._log.info("USB data complete handshake finished")

async def fake_camera_data(dut, data_to_send, line_width):
    # will de-assert PIXEL_DE after every line_width bytes as if it is doing lines of camera data
    # PIXEL_VS will de-assert after all data_to_send has been sent out

    dut.PIXEL_VS.value = 0
    dut.PIXEL_DE.value = 0
    # Idle a few clocks here
    await ClockCycles(dut.WRCLK, 50)
    dut.PIXEL_VS.value = 1
    await ClockCycles(dut.WRCLK, 6)
    bytes_sent = 0
    for dat in data_to_send:
        dut.PIXEL_DE.value = 1
        dut.WRDATA.value = dat
        await RisingEdge(dut.WRCLK)
        bytes_sent = bytes_sent + 1
        if(bytes_sent == line_width):
            bytes_sent = 0
            dut.PIXEL_DE.value = 0
            await ClockCycles(dut.WRCLK, 6)

async def fake_camera_frame(dut, hsize, vsize):
    VBLANK = 100
    HBLANK = 20

    generated_frame = FrameData()
    dut.PIXEL_VS.value = 0
    dut.PIXEL_DE.value = 0
    dut.WRDATA.value = 0
    # Idle a few clocks here
    await ClockCycles(dut.WRCLK, 50)
    dut.PIXEL_VS.value = 1
    await ClockCycles(dut.WRCLK, VBLANK)
    for yy in range(vsize):
        await ClockCycles(dut.WRCLK, HBLANK)
        dut.PIXEL_DE.value = 1
        for vv in range(hsize):
            newdata = getrandbits(16)
            dut.WRDATA.value = newdata
            generated_frame.dat.append(newdata)
            await RisingEdge(dut.WRCLK)
        dut.PIXEL_DE.value = 0
    dut.PIXEL_VS.value = 0
    await ClockCycles(dut.WRCLK, VBLANK)

    return generated_frame

@cocotb.test()
async def simple_single_buffer(dut):
    """Try accessing the design."""

    cocotb.start_soon(Clock(dut.WRCLK, 50, units="ns").start()) # mipi clock at 20MHz
    cocotb.start_soon(Clock(dut.RDCLK, 12, units="ns").start()) # usb clock at 83.33MHz

    await ClockCycles(dut.WRCLK, 10)
    dut.ARESET.value = 1
    await ClockCycles(dut.WRCLK, 4)
    dut.ARESET.value = 0
    await ClockCycles(dut.WRCLK, 20)

    # generate random data to send
    test_data = [getrandbits(16) for _ in range(1030)] # a little longer than a buffer
    test_data_bytes = list(itertools.chain.from_iterable([struct.pack("<H",bb) for bb in test_data]))

    usb_coro = cocotb.start_soon(fake_usb_sender(dut, test_data[0:2048]))
    # send it!
    await fake_camera_data(dut, test_data, 400)

    # finish up the other task
    usb_bytes = await Join(usb_coro)

    # check that everything matches
    assert len(usb_bytes) == 2048, "USB did not send correct number of bytes"
    if usb_bytes != test_data_bytes[0:2048]:
        dut._log.warn("Data did not match")
        for i in range(len(usb_bytes)):
            assert usb_bytes[i] == test_data_bytes[i], "First failed at data {}: {} vs {}".format(i, usb_bytes[i], test_data_bytes[i])

@cocotb.test()
async def send_full_frame(dut):
    cocotb.start_soon(Clock(dut.WRCLK, 50, units="ns").start()) # mipi clock at 20MHz
    cocotb.start_soon(Clock(dut.RDCLK, 12, units="ns").start()) # usb clock at 83.33MHz

    # Initialize all inputs
    dut.WRDATA.value = 0
    dut.PIXEL_DE.value = 0
    dut.PIXEL_VS.value = 0

    usb_sent_data = FrameData()

    dut._log.info("Sending Frame 1:")

    # Start the USB receiver
    cocotb.start_soon(complete_usb_loop_sender(dut, usb_sent_data))

    # Trigger reset
    await ClockCycles(dut.WRCLK, 10)
    dut.ARESET.value = 1
    await ClockCycles(dut.WRCLK, 4)
    dut.ARESET.value = 0
    await ClockCycles(dut.WRCLK, 20)

    # Send a frame
    my_generated_frame = await fake_camera_frame(dut, 120, 120)
    frame_bytes = list(itertools.chain.from_iterable([struct.pack("<H",bb) for bb in my_generated_frame.dat]))
    # Wait for the final send to start
    await First(RisingEdge(dut.DATA_READY), ClockCycles(dut.WRCLK, 1000))
    # and wait for it to finish
    await First(FallingEdge(dut.DATA_READY), ClockCycles(dut.WRCLK, 1000))

    # And a few clocks after it
    await ClockCycles(dut.WRCLK, 20)
    # Check the data
    assert len(usb_sent_data.dat) == len(frame_bytes), "Frame lengths do not match"
    if usb_sent_data.dat != frame_bytes:
        dut._log.warn("Data did not match")
        for i in range(len(usb_sent_data.dat)):
            assert usb_sent_data.dat[i] == frame_bytes[i], "First failed at data {}: {} vs {}".format(i, usb_sent_data.dat[i], frame_bytes[i])

    
    dut._log.info("Sending Frame 2:")
    # Now send another!!
    usb_sent_data.dat = [] # clear the usb receiver
    my_generated_frame = await fake_camera_frame(dut, 120, 120)
    frame_bytes = list(itertools.chain.from_iterable([struct.pack("<H",bb) for bb in my_generated_frame.dat]))

    # Wait for the final send to start
    await First(RisingEdge(dut.DATA_READY), ClockCycles(dut.WRCLK, 1000))
    # and wait for it to finish
    await First(FallingEdge(dut.DATA_READY), ClockCycles(dut.WRCLK, 1000))  

    # And a few clocks after it
    await ClockCycles(dut.WRCLK, 20)
    # Check the data
    assert len(usb_sent_data.dat) == len(frame_bytes), "Frame lengths do not match"
    if usb_sent_data.dat != frame_bytes:
        dut._log.warn("Data did not match")
        for i in range(len(usb_sent_data.dat)):
            assert usb_sent_data.dat[i] == frame_bytes[i], "First failed at data {}: {} vs {}".format(i, usb_sent_data.dat[i], frame_bytes[i])