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):
        # to construct correct crc and ack/stb signal
        self.submodules.arbiter = arbiter = Stream_Arbiter(n_downconn)
        self.submodules.broadcaster = broadcaster = Stream_Broadcaster(1)
        self.submodules.buffer = buffer = Buffer(word_layout_dchar)
        self.sinks = []

        for i in range(n_downconn):
            # generating the packet
            dchar_decoder = Duplicated_Char_Decoder()
            eop_marker = EOP_Marker()

            pipeline = [dchar_decoder, eop_marker]
            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.comb += pipeline[-1].source.connect(arbiter.sinks[i])

        # self.comb += arbiter.source.ack.eq(1)
         
        self.comb += [
            arbiter.source.connect(broadcaster.sink),
            broadcaster.sources[0].connect(buffer.sink),
        ]
        self.comb += buffer.source.ack.eq(1)



dut = Frame_Pipeline(CXP_CHANNELS)


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

    print("=================TEST========================")
    # yield dut.arbiter.active_channels.eq((2**active_ch) - 1)
    yield dut.arbiter.active_channels.eq((2**active_ch) - 1)
    sinks = dut.sinks

    for p in packets:
        for i in range(CXP_CHANNELS):
            if p["sink_id"] == i:
                yield sinks[p["sink_id"]].data.eq(p["data"])
                yield sinks[p["sink_id"]].k.eq(p["k"])
                yield sinks[p["sink_id"]].stb.eq(1)
                if "eop" in p:
                    yield sinks[p["sink_id"]].eop.eq(1)
                else:
                    yield sinks[p["sink_id"]].eop.eq(0)
            else:
                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

    # extra clk cycles
    for _ in range(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, with_header=True, with_crc=True):
    paks = []

    stream_id = 0
    pix_len = 10
    for p in range(20):
        n_conn = p % n_downconn
        if with_header:
            frame_pak_len = pix_len if with_crc else pix_len - 1
            frame = [
                {
                    "sink_id": n_conn,
                    "data": Replicate(C(stream_id, char_width), 4),
                    "k": Replicate(0, 4),
                },
                {
                    "sink_id": n_conn,
                    "data": Replicate(C(0, char_width), 4),
                    "k": Replicate(0, 4),
                },
                {
                    "sink_id": n_conn,
                    "data": Replicate(C(frame_pak_len, 2 * char_width)[8:], 4),
                    "k": Replicate(0, 4),
                },
                {
                    "sink_id": n_conn,
                    "data": Replicate(C(frame_pak_len, 2 * char_width)[:8], 4),
                    "k": Replicate(0, 4),
                },
            ]
        else:
            frame = []

        for n_pix in range(pix_len):
            if n_pix < pix_len - 1:
                frame.append(
                    {
                        "sink_id": n_conn,
                        # "data": C((n_conn + 1) << 28 | p << 24 | n_pix, word_width),
                        "data": C( p << 24 | n_pix, word_width),
                        "k": Replicate(0, 4),
                    }
                )
            else:
                frame.append(
                    {
                        "sink_id": n_conn,
                        # "data": C((n_conn + 1) << 28 | p << 24 | n_pix, word_width),
                        "data": C( p << 24 | n_pix, word_width),
                        "k": Replicate(0, 4),
                    }
                )
                frame.append(
                    {
                        "sink_id": n_conn,
                        "data": Replicate(KCode["pak_end"], 4),
                        "k": Replicate(1, 4),
                    }
                )

        paks += frame

    yield from packet_sim(paks)


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