from amaranth import *
from amaranth.lib.wiring import Component, In
from amaranth.utils import exact_log2

from amaranth_soc import wishbone
from amaranth_soc.memory import MemoryMap

class GranularityConverter(Component):
    def __init__(self, target_bus, *, granularity = 8):
        assert target_bus.data_width == target_bus.granularity, 'Target bus must have a granularity equal to width.' # This can be relaxed later.
        assert target_bus.data_width % granularity == 0, 'Target bus width must be a multiple of granularity'

        ratio = target_bus.data_width // granularity
        ratio_bits = exact_log2(ratio)

        wb_sig = wishbone.Signature(addr_width = target_bus.addr_width, data_width=target_bus.data_width, granularity = granularity)

        super().__init__({"wb_bus": In(wb_sig)})

        self.wb_bus.memory_map = MemoryMap(addr_width=target_bus.addr_width + ratio_bits,
                                           data_width=granularity)

        target_bus.memory_map.freeze()

        # MemoryMap doesn't support adding windows with a larger granularity than the parent.
        # Add and remap all resources directly instead:

        for resource in target_bus.memory_map.all_resources():
            self.wb_bus.memory_map.add_resource(resource.resource, name = tuple(part for name in resource.path for part in name), addr = resource.start * ratio, size = (resource.end - resource.start) * ratio)

        self._target_bus = target_bus

    def elaborate(self, platform):
        target = self._target_bus
        initiator  = self.wb_bus

        m = Module()

        m.d.comb += [
            target.adr.eq(initiator.adr),
            target.dat_w.eq(initiator.dat_w),
            initiator.dat_r.eq(target.dat_r),
            target.sel.eq(initiator.sel.all()), # With .all(), smaller writes will be silently ignored.
            target.cyc.eq(initiator.cyc),
            target.stb.eq(initiator.stb),
            target.we.eq(initiator.we),
            initiator.ack.eq(target.ack),
        ]

        return m