from migen import *
from misoc.interconnect.csr import *
from misoc.interconnect import stream
from misoc.cores.liteeth_mini.mac.crc import LiteEthMACCRCEngine

from src.gateware.cxp_pipeline import *

class EOP_Marker(Module):
    def __init__(self):
        self.sink = stream.Endpoint(word_layout_dchar)
        self.source = stream.Endpoint(word_layout_dchar)

        # # #

        last_stb = Signal()
        self.sync += [
            If((~self.source.stb | self.source.ack),
                self.source.stb.eq(self.sink.stb),
                self.source.payload.eq(self.sink.payload),
            ),
            last_stb.eq(self.sink.stb),
        ]
        self.comb += [
            self.sink.ack.eq(~self.source.stb | self.source.ack),
            self.source.eop.eq(~self.sink.stb & last_stb),
        ]

class Stream_MetaData_Extractor(Module):
    def __init__(self):
        self.sink = stream.Endpoint(word_layout_dchar)
        self.source = stream.Endpoint(word_layout_dchar)

        # # #

        # GOAL:
        # 0) accept four sinks??
        # 1) decode SPH, SPT
        # 2) decode Image header, line break
        # 3) verify the crc before phrasing image data downstream

        # HOW??
        # combine streams?
        # phrase linedata to sys CD
        # check stream data tag
        # only need to support mono16 for now

        self.stream_id = Signal(char_width)
        self.pak_tag = Signal(char_width)
        self.stream_pak_size = Signal(char_width * 2)

        self.submodules.fsm = fsm = FSM(reset_state="WAIT_HEADER")

        fsm.act(
            "WAIT_HEADER",
            NextValue(self.stream_id, self.stream_id.reset),
            NextValue(self.pak_tag, self.pak_tag.reset),
            NextValue(self.stream_pak_size, self.stream_pak_size.reset),
            self.sink.ack.eq(1),
            If(
                self.sink.stb,
                NextValue(self.stream_id, self.sink.dchar),
                NextState("GET_PAK_TAG"),
            ),
        )

        fsm.act(
            "GET_PAK_TAG",
            If(
                self.sink.stb,
                self.sink.ack.eq(1),
                NextValue(self.pak_tag, self.sink.dchar),
                NextState("GET_PAK_SIZE_0"),
            ),
        )

        fsm.act(
            "GET_PAK_SIZE_0",
            self.sink.ack.eq(1),
            If(
                self.sink.stb,
                NextValue(self.stream_pak_size[8:], self.sink.dchar),
                NextState("GET_PAK_SIZE_1"),
            ),
        )

        fsm.act(
            "GET_PAK_SIZE_1",
            self.sink.ack.eq(1),
            If(
                self.sink.stb,
                NextValue(self.stream_pak_size[:8], self.sink.dchar),
                NextState("STORE_BUFFER"),
            ),
        )

        fsm.act(
            "STORE_BUFFER",
            self.sink.connect(self.source),
            # both serve the same function but using the pak size I can remove eop injecter and save 1 cycle
            If(self.sink.stb,
                NextValue(self.stream_pak_size, self.stream_pak_size - 1),
                If(self.stream_pak_size == 1,
                    NextState("WAIT_HEADER"),
                )
            ),
            # If((self.sink.stb & self.sink.eop),
            #     NextState("WAIT_HEADER"),
            # )
        )

class Frame_Decoder(Module):
    def __init__(self):
        self.sink = stream.Endpoint(word_layout_dchar)
        self.source = stream.Endpoint(word_layout_dchar)

        # # #

        # TODO: decode Image header, line break

class Pixel_Decoder(Module):
    def __init__(self, pixel_format="mono16"):
        assert pixel_format == "mono16"

        self.sink = stream.Endpoint(word_layout_dchar)
        self.source = stream.Endpoint(word_layout_dchar)

        # # #

        # TODO: support mono16 for now?

class Streams_Dispatcher(Module):
    def __init__(self, downconns):
        n_downconn = len(downconns)
        self.submodules.mux = mux = stream.Multiplexer(word_layout_dchar, n_downconn)

        for i, c in enumerate(downconns):
            # if i == 0:
            self.comb += [
                # no backpressure
                c.source.ack.eq(1),
                c.source.connect(getattr(mux, "sink"+str(i)))
            ]

        self.source = stream.Endpoint(word_layout_dchar)

        self.submodules.fsm = fsm = FSM(reset_state="WAIT_HEADER")

        # TODO: add different downstream
        # stream_id = Signal()
        # case = dict((i, mux.source.connect(b.sink)) for i, b in enumerate(buffers))
        fsm.act(
            "WAIT_HEADER",
            mux.source.connect(self.source),
            If(mux.source.eop,
                NextState("SWITCH_CONN"),
            ),
        )

        read_mask = Signal(max=n_downconn)
        self.comb += mux.sel.eq(read_mask)
        fsm.act(
            "SWITCH_CONN",
            If(read_mask == n_downconn - 1,
                NextValue(read_mask, read_mask.reset),
            ).Else(
                NextValue(read_mask, read_mask + 1),
            ),
            NextState("WAIT_HEADER"),
        )


def reverse_bytes(s):
    assert len(s) % 8 == 0
    char = [s[i*8:(i+1)*8] for i in range(len(s)//8)]
    return Cat(char[::-1])

@ResetInserter()
@CEInserter()
class CXPCRC32(Module):
    # Section 9.2.2.2 (CXP-001-2021)
    width = 32
    polynom = 0x04C11DB7
    seed = 2**width - 1
    check = 0x00000000

    def __init__(self, data_width):
        self.data = Signal(data_width)
        self.value = Signal(self.width)
        self.error = Signal()

        # # #

        self.submodules.engine = LiteEthMACCRCEngine(
            data_width, self.width, self.polynom
        )
        reg = Signal(self.width, reset=self.seed)
        self.sync += reg.eq(self.engine.next)
        self.comb += [
            # the CRC Engine use Big Endian, need to reverse the bytes
            self.engine.data.eq(reverse_bytes(self.data)),
            self.engine.last.eq(reg),
            self.value.eq(reverse_bytes(reg[::-1])),
            self.error.eq(reg != self.check),
        ]

class CXPCRC32Inserter(Module):
    def __init__(self):
        self.sink = stream.Endpoint(word_layout)
        self.source = stream.Endpoint(word_layout)

        # # #
        
        self.submodules.crc = crc = CXPCRC32(word_dw)
        self.submodules.fsm = fsm = FSM(reset_state="IDLE")

        fsm.act("IDLE",
            crc.reset.eq(1),
            self.sink.ack.eq(1),
            If(self.sink.stb,
                self.sink.ack.eq(0),
                NextState("COPY"),
            )
        )
        fsm.act("COPY",
            crc.ce.eq(self.sink.stb & self.source.ack),
            crc.data.eq(self.sink.data),
            self.sink.connect(self.source),
            self.source.eop.eq(0),
            If(self.sink.stb & self.sink.eop & self.source.ack,
                NextState("INSERT"),
            )
        )        
        fsm.act("INSERT",
            self.source.stb.eq(1),
            self.source.eop.eq(1),
            self.source.data.eq(crc.value),
            If(self.source.ack, NextState("IDLE"))
        )

class StreamPacket_Wrapper(Module):
    def __init__(self):
        self.sink = stream.Endpoint(word_layout)
        self.source = stream.Endpoint(word_layout)

        # # #

        self.submodules.fsm = fsm = FSM(reset_state="IDLE")
        
        fsm.act("IDLE",
            self.sink.ack.eq(1),
            If(self.sink.stb,
                self.sink.ack.eq(0),
                NextState("INSERT_HEADER_0"),
            )
        )

        fsm.act("INSERT_HEADER_0",
            self.sink.ack.eq(0),
            self.source.stb.eq(1),
            self.source.data.eq(Replicate(KCode["pak_start"], 4)),
            self.source.k.eq(Replicate(1, 4)),
            If(self.source.ack, NextState("INSERT_HEADER_1")),
        )

        fsm.act("INSERT_HEADER_1",
            self.sink.ack.eq(0),
            self.source.stb.eq(1),
            self.source.data.eq(Replicate(C(0x01, char_width), 4)),
            self.source.k.eq(Replicate(0, 4)),
            If(self.source.ack, NextState("COPY")),
        )

        fsm.act("COPY",
            self.sink.connect(self.source),
            self.source.eop.eq(0),
            If(self.sink.stb & self.sink.eop & self.source.ack,
                NextState("INSERT_FOOTER"),
            ),
        )

        fsm.act("INSERT_FOOTER",
            self.sink.ack.eq(0),
            self.source.stb.eq(1),
            self.source.data.eq(Replicate(KCode["pak_end"], 4)),
            self.source.k.eq(Replicate(1, 4)),
            # Simulate RX don't have eop tagged
            # self.source.eop.eq(1),
            If(self.source.ack, NextState("IDLE")),
        )
        
# With KCode & 0x01*4
class StreamData_Generator(Module):
    def __init__(self):
        # should be big enough for all test
        self.submodules.buffer = buffer = stream.SyncFIFO(word_layout, 32)
        self.submodules.crc_inserter = crc_inserter = CXPCRC32Inserter()
        self.submodules.wrapper = wrapper = StreamPacket_Wrapper()

        # # #
        pipeline = [buffer, crc_inserter, wrapper]
        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

# For verifying crc in stream data packet
class Double_Stream_Buffer(Module):
    # default size is 2 kBtyes - Section 9.5.2 (CXP-001-2021)
    def __init__(self, size=16000):
        # detect and tag end of packet for crc
        # self.submodules.eop_marker = eop_marker = EOP_Marker()
        # self.sink = eop_marker.sink
        self.sink = stream.Endpoint(word_layout_dchar)

        self.submodules.crc = crc = CXPCRC32(word_dw)
        self.comb += crc.data.eq(self.sink.data)

        self.submodules.fsm = fsm = FSM(reset_state="RESET")

        write_mask = Signal()
        self.submodules.line_buffer0 = line_buffer0 = ResetInserter()(stream.SyncFIFO(word_layout_dchar, 2**bits_for(size//word_dw)))
        self.submodules.line_buffer1 = line_buffer1 = ResetInserter()(stream.SyncFIFO(word_layout_dchar, 2**bits_for(size//word_dw)))
        fsm.act("RESET",
            Case(write_mask,
                {
                     0: line_buffer0.reset.eq(1),
                     1: line_buffer1.reset.eq(1),
                }
            ),
            crc.reset.eq(1),
            NextState("CHECKING"),
        )

        fsm.act("CHECKING",
            self.sink.ack.eq(1),
            If(self.sink.stb,
                crc.ce.eq(1),
                If(self.sink.eop,
                    # discard the crc at the end
                    NextState("SWITCH_BUFFER")
                ).Else(
                    If(write_mask == 0,
                        self.sink.connect(line_buffer0.sink),
                    ).Else(
                        self.sink.connect(line_buffer1.sink),
                    ),
                )
            )
        )

        # only valid data will be passed to downstream
        fsm.act("SWITCH_BUFFER",
            If(~crc.error,
                NextValue(write_mask, ~write_mask),
            ),
            NextState("RESET"),
        )

        self.submodules.mux = mux = stream.Multiplexer(word_layout_dchar, 2)
        self.comb += [
            line_buffer0.source.connect(mux.sink0),
            line_buffer1.source.connect(mux.sink1),
            mux.sel.eq(~write_mask),
        ]
        self.source = mux.source

        # DEBUG:
        self.write = Signal()
        self.error = Signal()
        self.comb += [
            self.write.eq(write_mask),
            self.error.eq(crc.error),
        ]

        # # to add eop in the same cycle
        # # the tricks relies on the fact source lags sink one cycle
        # # but fsm .connect by default use combinational logic which the same cycle rising/falling edge check immpossible
        # fsm.act("CHECKING",
        #     NextValue(self.source.payload.raw_bits(), 0),
        #     NextValue(self.source.stb, 0),
        #     If(self.sink.stb,
        #         crc.ce.eq(1),
        #         If(self.sink.eop,
        #             NextState("RESET")
        #         ).Else(
        #             NextValue(self.source.payload.raw_bits(), self.sink.payload.raw_bits()),
        #             NextValue(self.source.stb, 1),
        #         )
        #     )
        # )
        
        # last_eop = Signal()
        # self.comb += self.source.eop.eq(~last_eop & self.sink.eop)