from migen import *
from sync_serdes import MultiLineRX, MultiLineTX
from multi_coders import MultiEncoder, CrossbarDecoder
from migen.build.platforms.sinara import kasli, efc
from eem_helpers import generate_pads
from kasli_crg import TransceiverCRG
from uart import UART
from migen.genlib.fifo import SyncFIFO


class BiDirectionalIO(Module):
    def __init__(self, i_pads, o_pads):
        self.i = Signal(4)
        self.o = Signal(4)
        self.t = Signal(4)

        for i in range(4):
            self.specials += Instance("OBUFTDS",
                i_I=self.i[i],
                o_O=o_pads[i].p,
                o_OB=o_pads[i].n,
                # Always chain the 3-states input to serializer
                # Vivado will complain otherwise
                i_T=self.t[i],
            )
        
        for i in range(4):
            self.specials += Instance("IBUFDS",
                i_I=i_pads[i].p,
                i_IB=i_pads[i].n,
                o_O=self.o[i],
            )


class ClockOut(Module):
    def __init__(self, clk_out):
        self.clk = Signal()
        # clk_buffer = Signal()
        # clk_se = Signal()

        # self.specials += Instance("BUFIO",
        #     i_I=self.clk,
        #     o_O=clk_buffer,
        # )

        # self.specials += Instance("ODDR",
        #     i_C=clk_buffer, i_CE=1, i_D1=0, i_D2=1, o_Q=clk_se),

        self.specials += Instance("OBUFDS",
            i_I=self.clk,
            o_O=clk_out.p,
            o_OB=clk_out.n,
        )


class Master(Module):
    def __init__(self, i_pads, o_pads):
        self.sysclk = ClockSignal("sys")
        self.submodules.tx = MultiLineTX()
        self.submodules.rx = MultiLineRX()
        self.submodules.channel = BiDirectionalIO(i_pads, o_pads)

        self.submodules.encoder = MultiEncoder(lsb_first=False)
        decoders = [ CrossbarDecoder(lsb_first=False) for _ in range(2) ]
        self.submodules += decoders

        self.comb += [
            # Transmitter to SERDES
            self.channel.i.eq(self.tx.ser_out),
            self.channel.t.eq(self.tx.t_out),

            # SERDES to receiver
            self.rx.ser_in_no_dly.eq(self.channel.o),

            # Connect encoders & decoders by default
            # Overrule the encoder connection during alignment
            self.tx.txdata.eq(Cat(self.encoder.output[0], self.encoder.output[1])),
            decoders[0].raw_input.eq(self.rx.rxdata[:10]),
            decoders[1].raw_input.eq(self.rx.rxdata[10:]),
        ]

        tx_fsm = FSM(reset_state="SEND_TRAINING")
        self.submodules += tx_fsm

        tx_fsm.act("SEND_TRAINING",
            self.tx.txdata.eq(0b00100001000010000100),
            # Keep sending the training sequence unless
            # an identifier is received
            NextState("SEND_TRAINING"),
            # If(self.rx.rxdata == 0b11111111111111111111,
            #     NextState("WAIT_IDENT_END"),
            # ),
        )

        tx_fsm.act("WAIT_IDENT_END",
            self.tx.txdata.eq(0),
            If(self.rx.rxdata != 0b11111111111111111111,
                NextState("SEND_ZERO"),
            ),
        )

        send_zero_duration = Signal(3)

        tx_fsm.act("SEND_ZERO",
            self.tx.txdata.eq(0),
            If(send_zero_duration == 0b111,
                NextState("SEND_PULSE"),
            ).Else(
                NextValue(send_zero_duration, send_zero_duration + 1),
            ),
        )

        tx_fsm.act("SEND_PULSE",
            self.tx.txdata.eq(0b11111111111111111111),
            self.encoder.start.eq(1),
            # Slave decoder start will be triggered by this state
            NextState("WAIT_GROUP_ALIGN"),
        )

        data = [ Signal(8) for _ in range(2) ]

        tx_fsm.act("WAIT_GROUP_ALIGN",
            # Wait for the identifier from the slave
            # TODO: Align the master receiver after
            If(self.rx.rxdata == 0b11111111111111111111,
                NextValue(data[0], 0x80),
                NextValue(data[1], 0x7F),
                NextState("TERMINATE"),
            ),
        )

        tx_fsm.act("TERMINATE",
            self.encoder.d[0].eq(0x89),
            self.encoder.d[1].eq(0x75),
            NextValue(data[0], data[0] + 1),
            NextValue(data[1], data[1] - 1),
            NextState("TERMINATE"),
        )


class Satellite(Module):
    def __init__(self, i_pads, o_pads, sys_clk_freq):
        self.uart_rx = Signal()
        self.uart_tx = Signal()

        self.submodules.uart = UART(round((115200/sys_clk_freq)*2**32))
        self.comb += [
            self.uart.phy_rx.eq(self.uart_rx),
            self.uart_tx.eq(self.uart.phy_tx),
        ]

        # Attach FIFO to UART TX, send rate is too slow w.r.t sysclk
        self.submodules.tx_fifo = SyncFIFO(8, 64)

        self.comb += [
            # UART TX path
            self.uart.tx_data.eq(self.tx_fifo.dout),
            self.uart.tx_stb.eq(self.tx_fifo.readable),
            self.tx_fifo.re.eq(self.uart.tx_ack),
        ]

        self.submodules.tx = MultiLineTX()
        self.submodules.rx = MultiLineRX()
        self.submodules.channel = BiDirectionalIO(i_pads, o_pads)

        self.submodules.encoder = MultiEncoder(lsb_first=False)
        decoders = [ CrossbarDecoder(lsb_first=False) for _ in range(2) ]
        self.submodules += decoders

        self.comb += [
            # Transmitter to SERDES
            self.channel.i.eq(self.tx.ser_out),
            self.channel.t.eq(self.tx.t_out),

            # SERDES to receiver
            self.rx.ser_in_no_dly.eq(self.channel.o),

            # Immediately start alignment for RX
            self.rx.start.eq(1),

            # Connect encoders & decoders by default
            # Overrule the encoder connection during alignment
            self.tx.txdata.eq(Cat(self.encoder.output[0], self.encoder.output[1])),
            decoders[0].raw_input.eq(self.rx.rxdata[:10]),
            decoders[1].raw_input.eq(self.rx.rxdata[10:]),

            # Start decoder after delay_done is set
            decoders[0].start.eq(self.rx.delay_done),
            decoders[1].start.eq(self.rx.delay_done),
        ]

        rx_fsm = FSM(reset_state="WAIT_ALIGN_DELAY")
        self.submodules += rx_fsm

        log_buffer = SyncFIFO(20, 128)
        self.submodules += log_buffer

        rx_fsm.act("WAIT_ALIGN_DELAY",
            If(self.rx.align_done & log_buffer.writable,
                log_buffer.we.eq(1),
                log_buffer.din.eq(self.rx.rxdata),
            ),
            If(~log_buffer.writable,
                NextState("DUMP_LOG_UPPER"),
            ),
            If(self.rx.delay_done,
                NextState("LOG_TRAFFIC"),
            ),
        )

        rx_fsm.act("LOG_TRAFFIC",
            If(log_buffer.writable,
                log_buffer.we.eq(1),
                log_buffer.din[0:8].eq(decoders[0].d),
                log_buffer.din[10:18].eq(decoders[1].d),
                # log_buffer.din.eq(self.rx.rxdata),
            ).Else(
                NextState("DUMP_LOG_UPPER"),
            ),
        )

        rx_fsm.act("DUMP_LOG_UPPER",
            If(log_buffer.readable,
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(log_buffer.dout[18:20]),
                    NextState("DUMP_LOG_LOWER"),
                ),
            ).Else(
                NextState("TERMINATE")
            ),
        )

        rx_fsm.act("DUMP_LOG_LOWER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(log_buffer.dout[10:18]),
                # log_buffer.re.eq(1),
                NextState("DUMP_LOG_UPPER_SECOND"),
            ),
        )

        rx_fsm.act("DUMP_LOG_UPPER_SECOND",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(log_buffer.dout[8:10]),
                NextState("DUMP_LOG_LOWER_SECOND"),
            ),
        )

        rx_fsm.act("DUMP_LOG_LOWER_SECOND",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(log_buffer.dout[:8]),
                log_buffer.re.eq(1),
                NextState("DUMP_LOG_UPPER"),
            ),
        )

        rx_fsm.act("TERMINATE",
            NextState("TERMINATE"),
        )

        tx_fsm = FSM(reset_state="WAIT_RX_INTRA_ALIGN")
        self.submodules += tx_fsm

        tx_fsm.act("WAIT_RX_INTRA_ALIGN",
            # Note: (Redundant) Only send something when aligned
            self.tx.txdata.eq(0),
            If(self.rx.align_done,
                NextState("SEND_INTRA_ALIGN_IDENT"),
            ),
        )

        # Master receiver is not aligned yet
        # We need to account for the possibility of group delay
        # and SERDES misalignment
        # 4 cycles of full 1s seems sufficiently detectable
        ident_timer = Signal(2)

        tx_fsm.act("SEND_INTRA_ALIGN_IDENT",
            self.tx.txdata.eq(0b11111111111111111111),
            If(ident_timer == 0b11,
                NextValue(ident_timer, 0),
                NextState("WAIT_RX_GROUP_ALIGN"),
            ).Else(
                NextValue(ident_timer, ident_timer + 1)
            ),
        )

        tx_fsm.act("WAIT_RX_GROUP_ALIGN",
            self.tx.txdata.eq(0),
            If(self.rx.delay_done,
                NextState("SEND_RX_DELAY_IDENT"),
            ),
        )

        tx_fsm.act("SEND_RX_DELAY_IDENT",
            self.tx.txdata.eq(0b11111111111111111111),
            If(ident_timer == 0b11,
                NextValue(ident_timer, 0),
                NextState("TERMINATE"),
            ).Else(
                NextValue(ident_timer, ident_timer + 1)
            ),
        )

        tx_fsm.act("TERMINATE",
            # TODO: Release the TX serdes to the encoders
            self.tx.txdata.eq(0),
            NextState("TERMINATE"),
        )


if __name__ == "__main__":
    import argparse
    import functools
    import os

    parser = argparse.ArgumentParser()
    parser.add_argument("platform")
    parser.add_argument("variant")
    args = parser.parse_args()

    platform_dict = {
        "kasli": kasli.Platform(hw_rev="v2.0"),
        "efc": efc.Platform(),
    }

    sysclk_name = {
        "kasli": "clk125_gtp",
        "efc": "gtp_clk",
    }

    platform = platform_dict[args.platform]
    
    # Generate pads for the I/O blocks
    for eem in range(2):
        generate_pads(platform, eem)
    data_eem = 0
    clk_eem = 1

    clk_pad = platform.request("dio{}".format(clk_eem), 0)

    # Manage clock source
    if args.variant == "master":
        sysclk = platform.request(sysclk_name[args.platform])
        # Fan out clock through EEM1 pair 0
        # clkout = ClockOut(clk_pad)
    elif args.variant == "satellite":
        # Receive clock from EEM1 pair 0
        sysclk = clk_pad

    if args.variant == "master": 
        i_pads = [
            platform.request("dio{}".format(data_eem), i) for i in range(4)
        ]
        o_pads = [
            platform.request("dio{}".format(data_eem), i+4) for i in range(4)
        ]
    elif args.variant == "satellite":
        i_pads = [
            platform.request("dio{}".format(data_eem), i) for i in range(4)
        ]
        o_pads = [
            platform.request("dio{}".format(data_eem), i+4) for i in range(4)
        ]
    else:
        raise ValueError("variant {} not implemented".format(args.variant))

    crg = TransceiverCRG(platform, sysclk, gte=(args.variant == "master"))

    variant = {
        "master": functools.partial(Master, i_pads, o_pads),
        "satellite": functools.partial(Satellite, i_pads, o_pads, crg.sys_clk_freq),
    }
    module_cls = variant[args.variant]

    top = module_cls()
    
    # Wire up UART core to the pads
    if args.variant == "satellite":
        uart_pads = platform.request("serial")
        top.comb += [
            top.uart_rx.eq(uart_pads.rx),
            uart_pads.tx.eq(top.uart_tx),
        ]

    top.submodules += crg
    # if args.variant == "master":
    #     top.submodules += clkout
    #     top.comb += clkout.clk.eq(top.sysclk)

    output_dir = "{}_{}_build".format(args.platform, args.variant)
    platform.build(top, build_dir=output_dir)