from nmigen import *
from nmigen.utils import *
from nmigen.hdl.rec import *

"""Common RTIO Interface"""

# CRI write happens in 3 cycles:
# 1. set timestamp and channel
# 2. set other payload elements and issue write command
# 3. check status

commands = {
    "nop": 0,
    "write": 1,
    # i_status should have the "wait for status" bit set until
    # an event is available, or timestamp is reached.
    "read": 2,
    # targets must assert o_buffer_space_valid in response
    # to this command
    "get_buffer_space": 3
}


layout = [
    ("cmd", 2, DIR_FANOUT),
    # 8  MSBs of chan_sel = routing destination
    # 16 LSBs of chan_sel = channel within the destination
    ("chan_sel", 24, DIR_FANOUT),

    ("o_timestamp", 64, DIR_FANOUT),
    ("o_data", 512, DIR_FANOUT),
    ("o_address", 8, DIR_FANOUT),
    # o_status bits:
    # <0:wait> <1:underflow> <2:destination unreachable>
    ("o_status", 3, DIR_FANIN),

    # pessimistic estimate of the number of outputs events that can be
    # written without waiting.
    # this feature may be omitted on systems without DRTIO.
    ("o_buffer_space_valid", 1, DIR_FANIN),
    ("o_buffer_space", 16, DIR_FANIN),

    ("i_timeout", 64, DIR_FANOUT),
    ("i_data", 32, DIR_FANIN),
    ("i_timestamp", 64, DIR_FANIN),
    # i_status bits:
    # <0:wait for event (command timeout)> <1:overflow> <2:wait for status>
    # <3:destination unreachable>
    # <0> and <1> are mutually exclusive. <1> has higher priority.
    ("i_status", 4, DIR_FANIN),
]

class Interface(Record):
    def __init__(self, **kwargs):
        super().__init__(layout, **kwargs)

# Skip KernelInitiator for now

class CRIDecoder(Elaboratable):
    def __init__(self, slaves=2, master=None, mode="async", enable_routing=False):
        if isinstance(slaves, int):
            slaves = [Interface() for _ in range(slaves)]
        if master is None:
            master = Interface()
        self.slaves = slaves
        self.master = master
        self.mode = mode
        self.enable_routing = enable_routing
    
    def elaborate(self, platform):
        m = Module()

        # routing
        if self.enable_routing:
            destination_unreachable = Interface()
            m.d.comb += destination_unreachable.o_status.eq(4)
            m.d.comb += destination_unreachable.i_status.eq(8)
            self.slaves = self.slaves[:]
            self.slaves.append(destination_unreachable)
            target_len = 2 ** (len(slaves) - 1).bit_length()
            self.slaves += [destination_unreachable] * (target_len - len(slaves))

        slave_bits = bits_for(len(self.slaves) - 1)
        selected = Signal(slave_bits)

        if self.enable_routing:
            routing_table = Memory(slave_bits, 256)

            if self.mode == "async":
                rtp_decoder_rdport = routing_table.read_port()
                rtp_decoder_wrport = routing_table.write_port()
            elif self.mode == "sync":
                rtp_decoder_rdport = routing_table.read_port(clock_domain="rtio")
                rtp_decoder_wrport = routing_tables.write_port(clock_domain="rtio")
            else:
                raise ValueError
            m.submodules.rtp_decoder_rdport = rtp_decoder_rdport
            m.submodules.rtp_decoder_wrport = rtp_decoder_wrport
            m.d.comb += rtp_decoder_rdport.addr.eq(self.master.chan_sel[16:])
            m.d.comb += selected.eq(rtp_decoder_rdport.data)
        else:
            m.d.sync += selected.eq(self.master.chan_sel[16:])

        # master -> slave
        for n, slave in enumerate(self.slaves):
            for name, size, direction in layout:
                if direction == DIR_FANOUT and name != "cmd":
                    m.d.comb += getattr(slave, name).eq(getattr(self.master, name))
            with m.If(selected == n):
                m.d.comb += slave.cmd.eq(self.master.cmd)

        # slave -> master
        with m.Switch(selected):
            for n, slave in enumerate(self.slaves):
                with m.Case(n):
                    for name, size, direction in layout:
                        if direction == DIR_FANIN:
                            m.d.comb += getattr(self.master, name).eq(getattr(slave, name))

        return m

# TODO: CRISwitch

# TODO: CRIInterconnectShared

# TODO: RoutingTableAccess