#
# This file is part of LUNA.
#
# Copyright (c) 2023 Great Scott Gadgets <info@greatscottgadgets.com>
# SPDX-License-Identifier: BSD-3-Clause

from amaranth       import *
from amaranth.utils import log2_int

from ..vendor.amaranth_soc        import wishbone
from ..vendor.amaranth_soc.memory import MemoryMap
from ..vendor.amaranth_soc.periph import ConstantMap

from ..vendor.lambdasoc.periph    import Peripheral


__all__ = ["SRAMPeripheral"]

# TODO this should live in wishbone/ !
class SRAMPeripheral(Peripheral, Elaboratable):
    """SRAM storage peripheral.

    Parameters
    ----------
    size : int
        Memory size in bytes.
    data_width : int
        Bus data width.
    granularity : int
        Bus granularity.
    writable : bool
        Memory is writable.

    Attributes
    ----------
    bus : :class:`amaranth_soc.wishbone.Interface`
        Wishbone bus interface.
    """
    # TODO raise bus.err if read-only and a bus write is attempted.
    def __init__(self, *, size, data_width=32, granularity=8, writable=True, init=None):
        super().__init__()

        if not isinstance(size, int) or size <= 0 or size & size-1:
            raise ValueError("Size must be an integer power of two, not {!r}"
                             .format(size))
        if size < data_width // granularity:
            raise ValueError("Size {} cannot be lesser than the data width/granularity ratio "
                             "of {} ({} / {})"
                              .format(size, data_width // granularity, data_width, granularity))

        self._mem = Memory(depth=(size * granularity) // data_width, width=data_width, init=init)

        self.bus = wishbone.Interface(addr_width=log2_int(self._mem.depth),
                                      data_width=self._mem.width, granularity=granularity,
                                      features={"cti", "bte"})

        map = MemoryMap(addr_width=log2_int(size), data_width=granularity, name=self.name)
        map.add_resource(self._mem, name="mem", size=size)
        self.bus.memory_map = map

        self.size        = size
        self.granularity = granularity
        self.writable    = writable

    @property
    def init(self):
        return self._mem.init

    @init.setter
    def init(self, init):
        self._mem.init = init

    @property
    def constant_map(self):
        return ConstantMap(
            SIZE = self.size,
        )

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

        incr = Signal.like(self.bus.adr)

        with m.Switch(self.bus.bte):
            with m.Case(wishbone.BurstTypeExt.LINEAR):
                m.d.comb += incr.eq(self.bus.adr + 1)
            with m.Case(wishbone.BurstTypeExt.WRAP_4):
                m.d.comb += incr[:2].eq(self.bus.adr[:2] + 1)
                m.d.comb += incr[2:].eq(self.bus.adr[2:])
            with m.Case(wishbone.BurstTypeExt.WRAP_8):
                m.d.comb += incr[:3].eq(self.bus.adr[:3] + 1)
                m.d.comb += incr[3:].eq(self.bus.adr[3:])
            with m.Case(wishbone.BurstTypeExt.WRAP_16):
                m.d.comb += incr[:4].eq(self.bus.adr[:4] + 1)
                m.d.comb += incr[4:].eq(self.bus.adr[4:])

        m.submodules.mem_rp = mem_rp = self._mem.read_port()
        m.d.comb += self.bus.dat_r.eq(mem_rp.data)

        # We'll ACK below
        # with m.If(self.bus.ack):
        #     m.d.sync += self.bus.ack.eq(0)

        with m.If(self.bus.cyc & self.bus.stb):
            m.d.sync += self.bus.ack.eq(1)
            with m.If((self.bus.cti == wishbone.CycleType.INCR_BURST) & self.bus.ack):
                m.d.comb += mem_rp.addr.eq(incr)
            with m.Else():
                m.d.comb += mem_rp.addr.eq(self.bus.adr)

        if self.writable:
            m.submodules.mem_wp = mem_wp = self._mem.write_port(granularity=self.granularity)
            m.d.comb += mem_wp.addr.eq(mem_rp.addr)
            m.d.comb += mem_wp.data.eq(self.bus.dat_w)
            with m.If(self.bus.cyc & self.bus.stb & self.bus.we):
                m.d.comb += mem_wp.en.eq(self.bus.sel)

        # From: lunasoc.gateware.soc.memory.WishboneRAM
        # We can handle any transaction request in a single cycle, when our RAM handles
        # the read or write. Accordingly, we'll ACK the cycle after any request.
        m.d.sync += self.bus.ack.eq(
            self.bus.cyc &
            self.bus.stb &
            ~self.bus.ack
        )

        return m
