from amaranth import Module, Signal, unsigned
from amaranth.lib.wiring import Component, Out, In, connect, flipped, Signature
from amaranth.lib.data import StructLayout, View
from amaranth.lib.memory import Memory
from amaranth_soc import wishbone


class CacheAddr(StructLayout):
    def __init__(self):
        super().__init__({
            "index": unsigned(8),
            "tag": unsigned(22)
        })


class TagData(StructLayout):
    def __init__(self):
        super().__init__({
            "tag": unsigned(22),
            "valid": unsigned(1)
        })


class WishboneMinimalICache(Component):
    def __init__(self):
        sig = {
            "cpu": In(wishbone.Signature(addr_width=30,
                                         data_width=32,
                                         granularity=8)),
            "en": In(1),
            "inval": In(Signature({
                "req": Out(1),
                "resp": In(1)
            })),
            "system": Out(wishbone.Signature(addr_width=30,
                                             data_width=32,
                                             granularity=8))
        }

        self.tag_len = 22
        self.index_len = 8

        self.data = Memory(shape=32, depth=2**self.index_len, init=[])
        self.tags = Memory(shape=TagData().as_shape(), depth=2**self.index_len,
                           init=[])
        super().__init__(sig)

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

        m.submodules.data = self.data
        m.submodules.tags = self.tags

        data_port_w = self.data.write_port()
        tags_port_w = self.tags.write_port()
        data_port_r = self.data.read_port(transparent_for=(data_port_w,))
        tags_port_r = self.tags.read_port(transparent_for=(tags_port_w,))

        cpu_uncached_path = wishbone.Interface(addr_width=30,
                                               data_width=32,
                                               granularity=8)
        cpu_cached_path = wishbone.Interface(addr_width=30,
                                             data_width=32,
                                             granularity=8)

        connect(m, flipped(cpu_uncached_path), flipped(self.cpu))
        connect(m, flipped(cpu_cached_path), flipped(self.cpu))

        with m.If(self.en):
            m.d.comb += [
                cpu_uncached_path.cyc.eq(0),
                cpu_uncached_path.stb.eq(0),
                cpu_cached_path.cyc.eq(self.cpu.cyc),
                cpu_cached_path.stb.eq(self.cpu.stb),
                self.cpu.ack.eq(cpu_cached_path.ack),
                self.cpu.dat_r.eq(cpu_cached_path.dat_r),
            ]
        with m.Else():
            m.d.comb += [
                cpu_uncached_path.cyc.eq(self.cpu.cyc),
                cpu_uncached_path.stb.eq(self.cpu.stb),
                cpu_cached_path.cyc.eq(0),
                cpu_cached_path.stb.eq(0),
                self.cpu.ack.eq(cpu_uncached_path.ack),
                self.cpu.dat_r.eq(cpu_uncached_path.dat_r),
            ]

            connect(m, flipped(self.system), cpu_uncached_path)

        cache_in = View(CacheAddr(), cpu_cached_path.adr)
        tag_in = View(TagData(), tags_port_w.data)
        tag_out = View(TagData(), tags_port_r.data)

        m.d.comb += cpu_cached_path.dat_r.eq(data_port_r.data)
        curr_line = Signal(self.index_len)

        with m.FSM(init="FLUSH"):
            with m.State("IDLE"):
                with m.If(self.inval.req == 1):
                    m.next = "FLUSH"
                with m.Elif(cpu_cached_path.cyc & cpu_cached_path.stb &
                            ~cpu_cached_path.we):
                    m.d.comb += [
                        data_port_r.addr.eq(cache_in.index),
                        tags_port_r.addr.eq(cache_in.index),
                        data_port_r.en.eq(1),
                        tags_port_r.en.eq(1),
                    ]

                    m.next = "CHECK"

            with m.State("CHECK"):
                with m.If((tag_out.tag == cache_in.tag) & tag_out.valid):
                    m.d.comb += cpu_cached_path.ack.eq(1)
                    m.next = "IDLE"
                with m.Else():
                    # Get a head-start on fetching.
                    connect(m, flipped(self.system), cpu_cached_path)
                    m.next = "FETCH"

            with m.State("FETCH"):
                connect(m, flipped(self.system), cpu_cached_path)

                with m.If(cpu_cached_path.cyc & cpu_cached_path.stb &
                          cpu_cached_path.ack):
                    m.d.comb += [
                        data_port_w.addr.eq(cache_in.index),
                        tags_port_w.addr.eq(cache_in.index),
                        data_port_w.en.eq(1),
                        tags_port_w.en.eq(1),
                        tag_in.valid.eq(1),
                        tag_in.tag.eq(cache_in.tag),
                        data_port_w.data.eq(cpu_cached_path.dat_r)
                    ]
                    m.next = "IDLE"

            with m.State("FLUSH"):
                m.d.sync += curr_line.eq(curr_line + 1)

                m.d.comb += [
                    tag_in.valid.eq(0),
                    tags_port_w.addr.eq(curr_line),
                    tags_port_w.en.eq(1),
                ]

                with m.If((curr_line + 1)[0:8] == 0):
                    m.d.comb += self.inval.resp.eq(1)
                    m.next = "IDLE"
                    # m.state = "INIT"

        return m