from nmigen import *
from nmigen_stdio import spiflash
from nmigen_soc import csr, wishbone
from nmigen.utils import *


__all__ = ["SPIFlashCore"]


class SPIFlashCore(Elaboratable):
    def __init__(self, spi_pins, *, spi_protocol, read_type, 
                 addr_width, sys_clk_freq, freq, 
                 data_width=32, divisor_bits=32,
                 bus_type="wishbone", granularity=None, **kwargs):
        self.spi_pins = spi_pins
        self.div = int(sys_clk_freq/freq) - 1
        self.div_bits = divisor_bits

        if read_type not in ["slow", "fast"]:
            raise ValueError("Type of the read operation must be one of "
                             "\"slow\" or \"fast\", not {}".format(read_type))
        if read_type == "slow":
            self.spi_io = spiflash.SPIFlashSlowReader(
                protocol=spi_protocol, 
                addr_width=addr_width,
                data_width=data_width,
                divisor_bits=divisor_bits,
                divisor=self.div,
                pins=self.spi_pins
            )
        elif read_type == "fast":
            self.spi_io = spiflash.SPIFlashFastReader(
                protocol=spi_protocol, 
                addr_width=addr_width,
                data_width=data_width,
                divisor_bits=divisor_bits,
                divisor=self.div,
                pins=self.spi_pins
            )

        if bus_type not in ["wishbone"]:
            raise ValueError("Bus type must be "
                             "\"wishbone\", not {}".format(bus_type))
        self.bus_type = bus_type
        if bus_type == "wishbone":
            if granularity is None:
                granularity = data_width
            self.bus = wishbone.Interface(
                addr_width=addr_width - log2_int(data_width//granularity),
                data_width=data_width,
                granularity=granularity,
                **kwargs
            )

        bus_dw = data_width
        if granularity is None:
            granularity = data_width
        self.granularity = bus_gr = granularity

        with csr.Bank(name="spi", addr_width=max(1, log2_int(bus_dw//bus_gr)),
                      data_width=bus_gr, type="mux") as self.csr:
            self.csr.r += [
                csr.Register(
                    "control", "rw", width=bus_dw,
                    fields=[csr.Field("divisor", "rw", width=self.div_bits, reset_value=self.div)]
                )
            ]
        self.wb2csr = csr.WishboneCSRBridge(self.csr.mux.bus, data_width=data_width)
        self.csr_bus = self.wb2csr.wb_bus

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

        m.submodules.spi = spi_io = self.spi_io
        m.submodules += self.csr, self.wb2csr

        # CSR - Divisor
        m.d.comb += spi_io.divisor.eq(self.csr.r.control.f.divisor.s)

        # non-CSR - Return Memory
        if self.bus_type == "wishbone":
            m.d.comb += [
                spi_io.ack.eq(self.bus.cyc & self.bus.stb & ~self.bus.we
                           & spi_io.rdy),
                spi_io.addr.eq(self.bus.adr),
                self.bus.ack.eq(spi_io.r_rdy)
            ]
            for i, sel_index in enumerate(self.bus.sel):
                m.d.comb += [
                    self.bus.dat_r[i*self.granularity:(i+1)*self.granularity]
                    .eq(spi_io.r_data[i*self.granularity:(i+1)*self.granularity]),
                ]

        return m