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 *

from types import SimpleNamespace

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 Streams_Crossbar(Module):
    def __init__(self, downconn_sources, stream_sinks):
        n_downconn = len(downconn_sources)
        self.active_conn= C(n_downconn) 
        # TODO: change self.active_conns to signal and link it to rx_ready of GTX lanes

        # # #

        self.submodules.mux = mux = stream.Multiplexer(word_layout_dchar, n_downconn)

        for i, downconn in enumerate(downconn_sources):
            self.comb += downconn.source.connect(getattr(mux, "sink"+str(i)))

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

        self.stream_id = Signal(char_width)
        case = dict((i, mux.source.connect(b.sink)) for i, b in enumerate(stream_sinks))
        fsm.act(
            "WAIT_HEADER",
            NextValue(self.stream_id, mux.source.dchar),
            If(mux.source.stb,
                NextState("COPY"),
            ),
        )

        fsm.act(
            "COPY",
            Case(self.stream_id, case),
            If(mux.source.eop,
                NextState("SWITCH_CONN"),
            ),
        )

        # Section 9.5.5 (CXP-001-2021)
        # When Multiple connections are active, stream packets are transmitted in
        # ascending order of Connection ID. And one connection shall be transmitting data at a time.
        read_mask = Signal(max=n_downconn)
        self.comb += mux.sel.eq(read_mask)
        fsm.act(
            "SWITCH_CONN",
            # assuming downconn_sources have ascending Connection ID 
            If(read_mask == self.active_conn - 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),
        ]

# For verifying crc in stream data packet
class Double_Stream_Buffer(Module):
    def __init__(self, size):
        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)


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

        # # #

        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",
            self.sink.ack.eq(1),
            If(
                self.sink.stb,
                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"),
                )
            ),
        )

class Frame_Extractor(Module):
    def __init__(self, pixel_format="mono16"):
        assert pixel_format in ["mono16"]
        pixel_format = {
            "mono16": C(0x0105, 2*char_width)
        }
        self.format_error = Signal()
        self.decode_err = Signal()

        self.new_frame = Signal()
        self.new_line = Signal()

        n_metadata_chars = 23
        img_header_layout = [
            ("stream_id", char_width),
            ("source_tag", 2*char_width), 
            ("x_size", 3*char_width), 
            ("x_offset", 3*char_width), 
            ("y_size", 3*char_width), 
            ("y_offset", 3*char_width), 
            ("l_size", 3*char_width), # number of data words per image line 
            ("pixel_format", 2*char_width), 
            ("tap_geo", 2*char_width), 
            ("flag", char_width), 
        ]
        assert layout_len(img_header_layout) == n_metadata_chars*char_width  


        # # #

        # TODO: decode Image header, line break
        self.sink = stream.Endpoint(word_layout_dchar)
        self.source = stream.Endpoint(word_layout_dchar)

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

        # DEBUG: remove this
        self.fsm_state = Signal()
        self.comb += self.fsm_state.eq(fsm.ongoing("IDLE"))

        fsm.act("IDLE",
            self.sink.ack.eq(1),
            If((self.sink.stb & (self.sink.dchar == KCode["stream_marker"]) & (self.sink.dchar_k == 1)),
                NextState("DECODE"),
            )
        )

        fsm.act("COPY",
            # until for new line or new frame
            If((self.sink.stb & (self.sink.dchar == KCode["stream_marker"]) & (self.sink.dchar_k == 1)),
                self.sink.ack.eq(1),
                NextState("DECODE"),
            ).Else(
                self.sink.connect(self.source),
            )
        )

        type = {
            "new_frame": 0x01,
            "line_break": 0x02,
        }

        cnt = Signal(max=n_metadata_chars)
        fsm.act("DECODE",
            self.sink.ack.eq(1),
            If(self.sink.stb,
                Case(self.sink.dchar, {
                    type["new_frame"]: [
                        self.new_frame.eq(1),
                        NextValue(cnt, cnt.reset),
                        NextState("GET_FRAME_DATA"),
                    ],
                    type["line_break"]: [
                        self.new_line.eq(1),
                        NextState("COPY"),
                    ],
                    "default": [
                         self.decode_err.eq(1),
                         # discard all data until valid frame
                         NextState("IDLE"),
                    ],
                }),
            )
        )

        packet_buffer = Signal(layout_len(img_header_layout))
        case = dict(
            (i, NextValue(packet_buffer[8*i:8*(i+1)], self.sink.dchar))
            for i in range(n_metadata_chars)
        )
        fsm.act("GET_FRAME_DATA",
            self.sink.ack.eq(1),
            If(self.sink.stb,
                Case(cnt, case),
                If(cnt == n_metadata_chars - 1,
                    NextState("COPY"),
                    NextValue(cnt, cnt.reset),
                ).Else(
                    NextValue(cnt, cnt + 1),
                ),
            ),
        )

        # dissect packet 
        self.metadata = SimpleNamespace()
        idx = 0
        for name, size in img_header_layout:
            # CXP use MSB even when sending duplicate chars
            setattr(self.metadata, name, reverse_bytes(packet_buffer[idx:idx+size]))
            idx += size


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 Stream_Pipeline(Module):
    # optimal stream packet size is 2 kBtyes - Section 9.5.2 (CXP-001-2021)
    def __init__(self, size=16000):

        self.submodules.double_buffer = double_buffer = Double_Stream_Buffer(size)
        self.submodules.parser = parser = Stream_Parser()
        self.submodules.frame_extractor = frame_extractor = Frame_Extractor()

        pipeline = [double_buffer, parser, frame_extractor]
        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
        
        # no backpressure for sim purposes
        self.sync += self.source.ack.eq(1)