from migen import *
from misoc.interconnect.csr import *
from misoc.interconnect import stream

from sim_generator import CXPCRC32Inserter
from src.gateware.cxp_frame_pipeline import *
from src.gateware.cxp_pipeline import *

CXP_CHANNELS = 2

class Frame_Pipeline(Module):
    def __init__(self, n_downconn, n_buffer=2):
        # to construct correct crc and ack/stb signal
        self.sinks = []
        self.downconn_source = []
        for i in range(n_downconn):
            # generating the packet
            buffer = stream.SyncFIFO(word_layout, 32)
            crc_inserter = CXPCRC32Inserter(False)
            dchar_decoder = Duplicated_Char_Decoder()

            pipeline = [buffer, crc_inserter, dchar_decoder]
            self.submodules += pipeline

            for s, d in zip(pipeline, pipeline[1:]):
                self.comb += s.source.connect(d.sink)
            self.sinks.append(pipeline[0].sink)
            self.downconn_source.append(pipeline[-1].source)


        self.buffer_sinks = []
        for i in range(n_buffer):
            # it should be mem block not "cycle buffer"
            buf = DChar_Dropper()
            self.submodules += buf
            self.sync += buf.source.ack.eq(1) # no backpressure for sim

            self.buffer_sinks.append(buf.sink)
        
        self.submodules.router = Frame_Packet_Router(self.downconn_source, self.buffer_sinks,packet_size=16384)




dut = Frame_Pipeline(CXP_CHANNELS)


def packet_sim(packets=[], active_ch=2):
    assert len(packets) == CXP_CHANNELS
    assert active_ch <= CXP_CHANNELS

    print("=================TEST========================")
    yield dut.router.active_src_mask.eq(active_ch - 1)
    sinks = dut.sinks

    cyc = len(packets[0])
    for c in range(cyc):
        for i, pak in enumerate(packets):
            yield sinks[i].data.eq(pak[c]["data"])
            yield sinks[i].k.eq(pak[c]["k"])
            yield sinks[i].stb.eq(1)
            if "eop" in pak[c]:
                yield sinks[i].eop.eq(1)
            else:
                yield sinks[i].eop.eq(0)

        yield

    # extra clk cycles
    for _ in range(cyc, cyc + 100):
        for i in range(CXP_CHANNELS):
            yield sinks[i].data.eq(0)
            yield sinks[i].k.eq(0)
            yield sinks[i].stb.eq(0)
            yield sinks[i].eop.eq(0)
            yield
    assert True


def testbench(n_downconn):
    paks = []
    for i in range(n_downconn):
        stream_id = i

        frame_packets = [
            # frame #0
            [
                {"data": C((i+1)<< 28 | 0x1111, word_dw), "k": Replicate(1, 4)},
                {"data": C((i+1)<< 28 | 0x2222, word_dw), "k": Replicate(0, 4)},
                {"data": C((i+1)<< 28 | 0x3333, word_dw), "k": Replicate(0, 4)},
                {"data": C((i+1)<< 28 | 0x4444, word_dw), "k": Replicate(0, 4), "eop":0},
            ],
            # frame #1
            [
                {"data": C((i+1)<< 28 | 0x5555, word_dw), "k": Replicate(1, 4)},
                {"data": C((i+1)<< 28 | 0x6666, word_dw), "k": Replicate(0, 4)},
                {"data": C((i+1)<< 28 | 0x7777, word_dw), "k": Replicate(0, 4)},
                {"data": C((i+1)<< 28 | 0x8888, word_dw), "k": Replicate(0, 4), "eop":0},
            ],
        ]

        total_pak = []
        for frame in frame_packets:
            total_pak += [
                {"data": Replicate(C(stream_id, char_width), 4), "k": Replicate(0, 4)},
                {"data": Replicate(C(0, char_width), 4), "k": Replicate(0, 4)},
                {
                    "data": Replicate(C(len(frame), 2 * char_width)[8:], 4),
                    "k": Replicate(0, 4),
                },
                {
                    "data": Replicate(C(len(frame), 2 * char_width)[:8], 4),
                    "k": Replicate(0, 4),
                },
            ]
            total_pak += frame
        
        paks.append(
            total_pak
        )

    yield from packet_sim(paks)


run_simulation(dut, testbench(CXP_CHANNELS), vcd_name="sim-cxp.vcd")