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

from nmigen.cli import main


class GearboxBusLayout(Layout):
    def __init__(self, *, in_width, out_width):
        super().__init__([
            # DATA
            ("data_in",        unsigned(in_width)),      # FROM SOURCE
            ("data_out",       unsigned(out_width)),     # TO DEST

            # CONTROL
            ("valid_in",       1),                       # FROM SOURCE
            ("ready_out",      1),                       # TO SOURCE

            ("ready_in",       1),                       # FROM DEST
            ("valid_out",      1),                       # TO DEST

        ])

class GearboxBus(Record):
    def __init__(self, *, in_width, out_width):
        super().__init__(GearboxBusLayout(in_width=in_width, out_width=out_width))

class ArbitraryGearbox(Elaboratable):
        def __init__(self, *, in_width, out_width):
            self.in_width = in_width
            self.out_width = out_width


            self.bus = GearboxBus(in_width=in_width, out_width=out_width)

        def elaborate(self, platform):
            m = Module()
            loop = Signal(1)
            len_storage = self.in_width + self.out_width
            #storage   = Signal(len_storage, reset=0b001_010_011_100_101_110_111)
            #storage   = Signal(len_storage, reset= 0b111_110_101_100_011_010_001)


            storage   = Signal(len_storage)

            write_ptr = Signal(range(len_storage))
            read_ptr  = Signal(range(len_storage))

            # read index logic here
            with m.If(read_ptr + self.out_width >= len_storage):
                m.d.sync += read_ptr.eq(read_ptr + self.out_width - len_storage)
            with m.Else():
                m.d.sync += read_ptr.eq(read_ptr + self.out_width)

            # read-out case analysis here:

            with m.If(read_ptr + self.out_width <= len_storage):
                m.d.comb += self.bus.data_out.eq(storage.bit_select(read_ptr, self.out_width))
            with m.Else():
                with m.Switch(read_ptr):
                    for cand_rptr in range(len_storage - self.out_width + 1, len_storage):
                        with m.Case(cand_rptr):
                            m.d.comb += loop.eq(1)
                            non_wrapping_bitwidth = len_storage - cand_rptr
                            wrapping_bitwidth     = self.out_width + cand_rptr - len_storage
                            m.d.comb += self.bus.data_out.eq(Cat(
                                storage.bit_select(cand_rptr, non_wrapping_bitwidth), # the non-wrapping bits are less-significant
                                storage.bit_select(0,         wrapping_bitwidth)      # ...than the wrapping bit
                                ))


            # The buffer is composed of two different flavor of bits:

            # Invalid bits  CANNOT BE READ OUT     CAN BE WRITTEN TO
            # Valid bits    CAN BE READ OUT        CANNOT BE WRITTEN TO

            # and the location of those bits are given by the read_ptr and the write_ptr

            # We only allow a read operation (which is not idempotent as the read_ptr is advanced) if:
            # 1) the downstream interface is signaling ready
            # 2) there are out_width valid bits in front of the read_ptr


            # Likewise, we only allow a write operation (which is heavily non-idempotent as the write_ptr is advanced
            # and the buffer is modified) if:
            # 1) the upstream interface is signaling valid (we cannot allow ourselves to ingest invalid data!)
            # 2) there are in_width invalid bits in front of the write_ptr


            # There are multiple cases for read_ptr and write_ptr:

            # read_ptr < write_ptr. Here, the valid bits do not wrap, the invalid bits wrap.
            # The valid bits are:   inclusive [read_ptr, write_ptr) exclusive
            # the invalid bits are  inclusive [write_ptr, K) exclusive, union with inclusive [0, read_ptr) exclusive


            # write_ptr < read_ptr. Here, the valid bits wrap, and the invalid bits do not wrap.
            # The valid bits are:   inclusive [read_ptr, K) exclusive, union with inclusive [0, write_ptr) exclusive
            # the invalid bits are  inclusive [write_ptr, read_ptr) exclusive

            # Naturally, there is a tricky edge case which requires an extra bit (literally) of disambiguation.
            # If write_ptr == read_ptr we don't know if the buffer is entirely full or entirely empty.

            # If the last operation was a read, and we find write_ptr == read_ptr, we know it's empty
            # If the last operation was a write, and we find write_ptr == read_ptr, we know it's full
            

            return m




class DummyPlug(Elaboratable):

    #def __init__(self):



    def elaborate(self, platform):
        m = Module()

        m.submodules.table = table = ArbitraryGearbox(in_width=(9-3), out_width=3)
        counter = Signal(8)
        m.d.sync += counter.eq(counter+1)

        with m.If(counter == 3):
            m.d.comb += table.bus.data_in.eq(1)


        return m



if __name__ == '__main__':
    baka =DummyPlug()
    main(baka)
    #platform.build(DummyPlug())