from migen import *
from misoc.interconnect import stream

from sim_pipeline import *
from sim_generator import StreamData_Generator

from src.gateware.cxp_pipeline import *

class CXP_Links(Module):
    def __init__(self):
        # TODO: select the correct buffer to read from
        # NOTE: although there are double buffer in each connect, the reading must be faster than writing to avoid data loss

        self.downconns = []
        self.stream_buffers = []
        for i in range(2):
            downconn = Pipeline()
            setattr(self.submodules, "cxp_conn"+str(i), downconn)
            self.downconns.append(downconn)

            double_buffer = Double_Stream_Buffer()
            setattr(self.submodules, "stream_buffer"+str(i), double_buffer)
            self.stream_buffers.append(double_buffer)

            # no backpressure for sim purposes
            self.sync += double_buffer.source.ack.eq(1)
            

        self.submodules.crossbar = crossbar = Streams_Crossbar(self.downconns, self.stream_buffers)
        
        # TODO: add extractor

        # pipeline = [dispatcher, double_buffer]
        # for s, d in zip(pipeline, pipeline[1:]):
        #     self.comb += s.source.connect(d.sink)
        # self.source = pipeline[-1].source
        
        # # no backpressure
        # self.sync += self.source.ack.eq(1)


class Pipeline(Module):
    def __init__(self):
        self.submodules.generator = generator = StreamData_Generator()
        self.submodules.dchar_decoder = dchar_decoder = Duplicated_Char_Decoder()
        self.submodules.data_decoder = data_decoder = RX_Bootstrap()
        self.submodules.eop_marker = eop_marker = EOP_Marker()

        # # #

        pipeline = [generator, dchar_decoder, data_decoder, eop_marker]
        for s, d in zip(pipeline, pipeline[1:]):
            self.comb += s.source.connect(d.sink)
        self.sink = pipeline[0].sink
        self.source = pipeline[-1].source

        # self.comb += self.source.ack.eq(1)


dut = CXP_Links()

def check_case(packet=[]):
    print("=================TEST========================")

    downconns = dut.downconns
    stream_buffers = dut.stream_buffers
    ch = 0
    
    for i, p in enumerate(packet):
        for x in range(len(downconns)):
            if x == ch:
                yield downconns[x].sink.data.eq(p["data"])
                yield downconns[x].sink.k.eq(p["k"])
                yield downconns[x].sink.stb.eq(1)
            else:
                yield downconns[x].sink.data.eq(0)
                yield downconns[x].sink.k.eq(0)
                yield downconns[x].sink.stb.eq(0)
                yield downconns[x].sink.eop.eq(0)
        if "eop" in p:
            yield downconns[ch].sink.eop.eq(1)
            # compensate for delay
            # yield
            # yield downconns[ch].sink.data.eq(0)
            # yield downconns[ch].sink.k.eq(0)
            # yield downconns[ch].sink.stb.eq(0)
            # yield downconns[ch].sink.eop.eq(0)
            # yield
            # yield
            # yield
            ch = (ch + 1) % len(downconns)
        else:
            yield downconns[ch].sink.eop.eq(0)

        # check cycle result
        yield
        source = dut.stream_buffers[0].source
        print(
            f"\nCYCLE#{i} : source char = {yield source.data:#X} k = {yield source.k:#X} stb = {yield source.stb} ack = {yield source.ack} eop = {yield source.eop}"
            # f" source dchar = {yield source.dchar:#X} dchar_k = {yield source.dchar_k:#X}"
            f"\nCYCLE#{i} : read mask = {yield dut.crossbar .mux.sel}"
            # f"\nCYCLE#{i} : stream id = {yield decoder.stream_id:#X} pak_tag = {yield decoder.pak_tag:#X}"
            # f" stream_pak_size = {yield decoder.stream_pak_size:#X}"
        )
        # crc = downconns[1].generator.crc_inserter.crc
        # crc = dut.double_buffer.crc
        # print(
        #     f"CYCLE#{i} : crc error = {yield crc.error:#X} crc value = {yield crc.value:#X}"
        #     f" crc data = {yield crc.data:#X} engine next = {yield crc.engine.next:#X} ce = {yield crc.ce}"
        # )

            
    # extra clk cycles
    cyc = i + 1
    for i in range(cyc, cyc + 30):
        for x in range(len(downconns)):
            # yield won't reset every cycle
            yield downconns[x].sink.data.eq(0)
            yield downconns[x].sink.k.eq(0)
            yield downconns[x].sink.stb.eq(0)
            yield downconns[x].sink.eop.eq(0)
        yield
        print(
            f"\nCYCLE#{i} : source char = {yield source.data:#X} k = {yield source.k:#X} stb = {yield source.stb} ack = {yield source.ack} eop = {yield source.eop}"
            # f" source dchar = {yield source.dchar:#X} dchar_k = {yield source.dchar_k:#X}"
            f"\nCYCLE#{i} : read mask = {yield dut.crossbar .mux.sel}"
            # f"\nCYCLE#{i} : stream id = {yield decoder.stream_id:#X} pak_tag = {yield decoder.pak_tag:#X}"
            # f" stream_pak_size = {yield decoder.stream_pak_size:#X}"
        )
    assert True


def testbench():
    # stream_id = 0x01
    streams = [
        [
            {"data": 0x11111111, "k": Replicate(0, 4)},
            {"data": 0xB105F00D, "k": Replicate(0, 4)},
        ],
        [
            {"data": 0x22222222, "k": Replicate(0, 4)},
            {"data": 0xC00010FF, "k": Replicate(0, 4)},
        ],
        [
            {"data": 0x33333333, "k": Replicate(0, 4)},
            {"data": 0xC0A79AE5, "k": Replicate(0, 4)},
        ],
    ]

    packet = []
    for i, s in enumerate(streams):
        s[-1]["eop"] = 0
        packet += [
            {"data": Replicate(C(i % 2, char_width), 4), "k": Replicate(0, 4)},
            {"data": Replicate(C(i, char_width), 4), "k": Replicate(0, 4)},
            {
                "data": Replicate(C(len(s) >> 8 & 0xFF, char_width), 4),
                "k": Replicate(0, 4),
            },
            {"data": Replicate(C(len(s) & 0xFF, char_width), 4), "k": Replicate(0, 4)},
            *s,
        ]


    yield from check_case(packet)


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