from migen.fhdl.std import *
from migen.bus.transactions import *
from migen.bus import wishbone
from migen.genlib.misc import timeline
from migen.genlib.record import Record
from migen.bank.description import AutoCSR, CSRStorage, CSRStatus

_FAST_READ = 0x0b
_DIOFR = 0xbb
_QIOFR = 0xeb


def _format_cmd(cmd, spi_width):
    """
    `cmd` is the read instruction. Since everything is transmitted on all
    dq lines (cmd, adr and data), extend/interleave cmd to full pads.dq
    width even if dq1-dq3 are don't care during the command phase:
    For example, for N25Q128, 0xeb is the quad i/o fast read, and
    extended to 4 bits (dq1,dq2,dq3 high) is: 0xfffefeff
    """
    c = 2**(8*spi_width)-1
    for b in range(8):
        if not (cmd>>b)%2:
            c &= ~(1<<(b*spi_width))
    return c


class SpiFlash(Module, AutoCSR):
    def __init__(self, pads, dummy=15, div=2, with_bitbang=True):
        """
        Simple SPI flash, e.g. N25Q128 on the LX9 Microboard.

        Supports multi-bit pseudo-parallel reads (aka Dual or Quad I/O Fast
        Read). Only supports mode0 (cpol=0, cpha=0).
        Optionally supports software bitbanging (for write, erase, or other commands).
        """
        self.bus = bus = wishbone.Interface()
        spi_width = flen(pads.dq)
        if with_bitbang:
            self.bitbang = CSRStorage(4)
            self.miso = CSRStatus()
            self.bitbang_en = CSRStorage()

        ###

        cs_n = Signal(reset=1)
        clk = Signal()
        dq_oe = Signal()
        wbone_width = flen(bus.dat_r)


        read_cmd_params = {
            4: (_format_cmd(_QIOFR, 4), 4*8),
            2: (_format_cmd(_DIOFR, 2), 2*8),
            1: (_format_cmd(_FAST_READ, 1), 1*8)
        }
        read_cmd, cmd_width = read_cmd_params[spi_width]
        addr_width = 24

        pads.cs_n.reset = 1

        dq = TSTriple(spi_width)
        self.specials.dq = dq.get_tristate(pads.dq)

        sr = Signal(max(cmd_width, addr_width, wbone_width))
        dqs = Replicate(1, spi_width-1)

        self.comb += bus.dat_r.eq(sr)

        hw_read_logic = [
            pads.clk.eq(clk),
            pads.cs_n.eq(cs_n),
            dq.o.eq(sr[-spi_width:]),
            dq.oe.eq(dq_oe)
        ]

        if with_bitbang:
            bitbang_logic = [
                pads.clk.eq(self.bitbang.storage[1]),
                pads.cs_n.eq(self.bitbang.storage[2]),
                dq.o.eq(Cat(self.bitbang.storage[0], dqs)),
                If(self.bitbang.storage[3],
                    dq.oe.eq(0)
                ).Else(
                    dq.oe.eq(1)
                ),
                If(self.bitbang.storage[1],
                    self.miso.status.eq(dq.i[1])
                )
            ]

            self.comb += \
                If(self.bitbang_en.storage,
                    bitbang_logic
                ).Else(
                    hw_read_logic
                )
        else:
            self.comb += hw_read_logic

        if div < 2:
            raise ValueError("Unsupported value \'{}\' for div parameter for SpiFlash core".format(div))
        else:
            i = Signal(max=div)
            dqi = Signal(spi_width)
            self.sync += [
                If(i == div//2 - 1,
                    clk.eq(1),
                    dqi.eq(dq.i),
                ),
                If(i == div - 1,
                    i.eq(0),
                    clk.eq(0),
                    sr.eq(Cat(dqi, sr[:-spi_width]))
                ).Else(
                    i.eq(i + 1),
                ),
            ]

        # spi is byte-addressed, prefix by zeros
        z = Replicate(0, log2_int(wbone_width//8))

        seq = [
            (cmd_width//spi_width*div,
                [dq_oe.eq(1), cs_n.eq(0), sr[-cmd_width:].eq(read_cmd)]),
            (addr_width//spi_width*div,
                [sr[-addr_width:].eq(Cat(z, bus.adr))]),
            ((dummy + wbone_width//spi_width)*div,
                [dq_oe.eq(0)]),
            (1,
                [bus.ack.eq(1), cs_n.eq(1)]),
            (div, # tSHSL!
                [bus.ack.eq(0)]),
            (0,
                []),
        ]

        # accumulate timeline deltas
        t, tseq = 0, []
        for dt, a in seq:
            tseq.append((t, a))
            t += dt

        self.sync += timeline(bus.cyc & bus.stb & (i == div - 1), tseq)


class SpiFlashTB(Module):
    def __init__(self):
        self.submodules.master = wishbone.Initiator(self.gen_reads())
        self.pads = Record([("cs_n", 1), ("clk", 1), ("dq", 4)])
        self.submodules.slave = SpiFlash(self.pads)
        self.submodules.tap = wishbone.Tap(self.slave.bus)
        self.submodules.intercon = wishbone.InterconnectPointToPoint(
                self.master.bus, self.slave.bus)
        self.cycle = 0

    def gen_reads(self):
        for a in range(10):
            t = TRead(a)
            yield t
            print("read {} in {} cycles(s)".format(t.data, t.latency))

    def do_simulation(self, selfp):
        if selfp.pads.cs_n:
            self.cycle = 0
        else:
            self.cycle += 1
            if not selfp.slave.dq.oe:
                selfp.slave.dq.i = self.cycle & 0xf
    do_simulation.passive = True

if __name__ == "__main__":
    from migen.sim.generic import run_simulation
    from migen.fhdl import verilog

    pads = Record([("cs_n", 1), ("clk", 1), ("dq", 4)])
    s = SpiFlash(pads)
    print(verilog.convert(s, ios={pads.clk, pads.cs_n, pads.dq, s.bus.adr,
        s.bus.dat_r, s.bus.cyc, s.bus.ack, s.bus.stb}))

    run_simulation(SpiFlashTB(), vcd_name="spiflash.vcd")