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, reset= 0b1_100_011_00)

            write_ptr = Signal(range(len_storage))
            read_ptr  = Signal(range(len_storage), reset=2)

            # 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
                                ))




            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())