from amaranth import *
from amaranth_soc.wishbone import *
from amaranth_soc.memory import *

from math import log2, ceil

# TODO impl select


# We sub-class wishbone.Interface here because it needs to be a bus object to be added as a window to Wishbone stuff
class ROM(Elaboratable, Interface):
    def __init__(self, *, size_bytes=4096, data=None):
        #self.size = len(data)
        addr_width = ceil(log2(size_bytes >> 2))
        self.data = Memory(width=32, depth=size_bytes >> 2, init=data)
        self.r = self.data.read_port()

        # Need to init Interface
        Interface.__init__(self, addr_width=addr_width, data_width=32, granularity=8)

        # This is effectively a "window", and it has a certain set of resources
        # 12 = log2(4096)
        memory_map = MemoryMap(addr_width=addr_width + 2, data_width=8)
        # TODO need to unify how I deal with size
        # In this case, one resource, which is out memory
        memory_map.add_resource(self.data, name="rom_data", size=size_bytes)

        self.memory_map = memory_map

    # Connects memory port signals to wishbone interface
    def elaborate(self, platform):
        # Stolen from https://vivonomicon.com/2020/04/14/learning-fpga-design-with-nmigen/
        m = Module()
        # Register the read port submodule.
        m.submodules.r = self.r

        # 'ack' signal should rest at 0.
        m.d.sync += self.ack.eq( 0 )
        # Simulated reads only take one cycle, but only acknowledge
        # them after 'cyc' and 'stb' are asserted.
        with m.If( self.cyc ):
            m.d.sync += self.ack.eq( self.stb )

        # Set 'dat_r' bus signal to the value in the
        # requested 'data' array index.
        m.d.comb += [
            self.r.addr.eq( self.adr ),
            self.dat_r.eq( self.r.data )
        ]

        # End of simulated memory module.
        return m


# TODO support read segmentation or whatever it's called, where you read/write certian bytes from memory
# Otherwise we can't store individual bytes, and this will wreck shit in weird ways.
class RAM(Elaboratable, Interface):
    def __init__(self, *, size_bytes=4096):
        addr_width = ceil(log2(size_bytes >> 2))
        self.addr_width = addr_width
        self.data = Memory(width=32, depth=size_bytes >> 2)
        self.r = self.data.read_port()
        self.w = self.data.write_port()

        # Need to init Interface
        Interface.__init__(self, addr_width=addr_width, data_width=32, granularity=8)

        # This is effectively a "window", and it has a certain set of resources
        # 12 = log2(4096)
        memory_map = MemoryMap(addr_width=addr_width + 2, data_width=8)
        # TODO need to unify how I deal with size
        # In this case, one resource, which is out memory
        memory_map.add_resource(self.data, name="ram_data", size=size_bytes)

        self.memory_map = memory_map

    # Connects memory port signals to wishbone interface
    def elaborate(self, platform):
        # Stolen from https://vivonomicon.com/2020/04/14/learning-fpga-design-with-nmigen/
        m = Module()
        # Register the read port submodule.
        m.submodules.r = self.r
        m.submodules.w = self.w

        # TODO not sure if this is the idiomatic way
        self.r.en = Const(1)
        # Default to not writing data
        m.d.sync += self.w.en.eq(0)


        # Use read data to populate un-written wishbone data for sub-32-bit reads
        for i in range(4):
            with m.If(self.sel.bit_select(i, 1)):
                m.d.sync += self.w.data.word_select(i, 8).eq(self.dat_w.word_select(i, 8))
            with m.Else():
                m.d.sync += self.w.data.word_select(i, 8).eq(self.r.data.word_select(i, 8))

        # 'ack' signal should rest at 0.
        m.d.sync += self.ack.eq(0)
        # Simulated reads only take one cycle, but only acknowledge
        # them after 'cyc' and 'stb' are asserted.
        with m.If( self.cyc & self.stb):
            m.d.sync += self.ack.eq(1)

            # Write to address if we are writing
            with m.If(self.we):
                m.d.sync += self.w.en.eq(1)

        # Set 'dat_r' bus signal to the value in the
        # requested 'data' array index.
        m.d.comb += [
            self.r.addr.eq( self.adr ),
            self.dat_r.eq( self.r.data ),
            self.w.addr.eq(self.adr),
            #self.w.data.eq(self.dat_w),
        ]

        # End of simulated memory module.
        return m