from migen import *
from migen.genlib.fsm import FSM
from migen.genlib.fifo import SyncFIFOBuffered
from migen_axi.interconnect import axi
from misoc.interconnect.csr import *


AXI_ADDRESS_WIDTH = 32
AXI_DATA_WIDTH = 64
AXI_BURST_LEN = 16
AXI_ALIGNMENT_BITS = log2_int(AXI_BURST_LEN*AXI_DATA_WIDTH//8)
AXI_ALIGNED_ADDRESS_WIDTH = AXI_ADDRESS_WIDTH - AXI_ALIGNMENT_BITS


class BlockAddressGenerator(Module):
    def __init__(self, membus_stream, data_width):
        self.base_address = Signal(AXI_ALIGNED_ADDRESS_WIDTH)
        self.length = Signal(AXI_ALIGNED_ADDRESS_WIDTH)
        self.start = Signal()
        self.busy = Signal()

        current_address = Signal(AXI_ALIGNED_ADDRESS_WIDTH)
        remaining = Signal(AXI_ALIGNED_ADDRESS_WIDTH)

        self.comb += [
            membus_stream.addr.eq(Cat(C(0, AXI_ALIGNMENT_BITS), current_address)),
            membus_stream.id.eq(0),                         # Same ID for all transactions to forbid reordering.
            membus_stream.burst.eq(axi.Burst.incr.value),
            membus_stream.len.eq(AXI_BURST_LEN-1),          # Number of transfers in burst (0->1 transfer, 1->2 transfers...).
            membus_stream.size.eq(log2_int(AXI_DATA_WIDTH//8)), # Width of burst: 3 = 8 bytes = 64 bits.
            membus_stream.cache.eq(0xf),
        ]

        fsm = FSM()
        self.submodules += fsm
        fsm.act("IDLE",
            If(self.start,
                NextValue(current_address, self.base_address),
                NextValue(remaining, self.length),
                NextState("RUNNING")
            )
        )
        fsm.act("RUNNING",
            self.busy.eq(1),
            membus_stream.valid.eq(1),
            If(membus_stream.ready,
                NextValue(current_address, current_address + 1),
                NextValue(remaining, remaining - 1),
                If(remaining == 1,
                    NextState("IDLE")
                )
            )
        )


def convert_signal_endianness(signal):
    assert len(signal) % 8 == 0
    nbytes = len(signal)//8
    signal_bytes = []
    for i in range(nbytes):
        signal_bytes.append(signal[8*i:8*(i+1)])
    return Cat(*reversed(signal_bytes))


class ADCWriter(Module):
    def __init__(self, membus_stream, pads):
        self.length = Signal(AXI_ALIGNED_ADDRESS_WIDTH) # in AXI bursts
        self.start = Signal()
        self.overflow = Signal()
        self.busy = Signal()

        fifo = SyncFIFOBuffered(64, 512)
        self.submodules += fifo

        # FIFO write
        adc_a = Signal(16)
        adc_b = Signal(16)
        self.sync += [
            adc_a.eq(pads.data_a),
            adc_b.eq(pads.data_b),
        ]

        fifo_inbuf = Signal(64)
        fifo_inbuf_sel = Signal()
        self.sync += [
            fifo_inbuf_sel.eq(~fifo_inbuf_sel),
            If(fifo_inbuf_sel,
                fifo_inbuf[32:].eq(Cat(adc_a, adc_b)),
            ).Else(
                fifo_inbuf[:32].eq(Cat(adc_a, adc_b)),
            )
        ]
        self.comb += fifo.din.eq(fifo_inbuf),

        assert AXI_DATA_WIDTH == len(fifo_inbuf)
        remaining = Signal(AXI_ADDRESS_WIDTH - log2_int(AXI_DATA_WIDTH//8)) # in AXI_DATA_WIDTH words
        self.sync += [
            If(remaining != 0, remaining.eq(remaining - 1)),
            If(self.start, remaining.eq(self.length << log2_int(AXI_BURST_LEN))),
        ]
        self.comb += fifo.we.eq((remaining != 0) & ~fifo_inbuf_sel)

        self.comb += self.overflow.eq(fifo.we & ~fifo.writable)

        # FIFO read
        self.comb += [
            membus_stream.id.eq(0),
            membus_stream.valid.eq(fifo.readable),
            fifo.re.eq(membus_stream.ready),
            membus_stream.data.eq(convert_signal_endianness(fifo.dout)),
            membus_stream.strb.eq(2**(AXI_DATA_WIDTH//8)-1),
        ]
        beat_count = Signal(max=AXI_BURST_LEN)
        self.sync += [
            If(membus_stream.valid & membus_stream.ready,
                membus_stream.last.eq(0),
                If(membus_stream.last,
                    beat_count.eq(0)
                ).Else(
                    If(beat_count == AXI_BURST_LEN-2, membus_stream.last.eq(1)),
                    beat_count.eq(beat_count + 1)
                )
            )
        ]

        # Busy generation
        remaining_sys = Signal(AXI_ADDRESS_WIDTH - log2_int(AXI_DATA_WIDTH//8))
        self.sync += [
           If(self.start, remaining_sys.eq(self.length << log2_int(AXI_BURST_LEN))),
           If(fifo.readable & fifo.re, remaining_sys.eq(remaining_sys - 1))
        ]
        self.comb += self.busy.eq(remaining_sys != 0)


class ADC(Module, AutoCSR):
    def __init__(self, membus, pads):
        self.base_address = CSRStorage(AXI_ADDRESS_WIDTH, alignment_bits=AXI_ALIGNMENT_BITS)
        self.length = CSRStorage(AXI_ADDRESS_WIDTH, alignment_bits=AXI_ALIGNMENT_BITS)
        self.start = CSR()
        self.busy = CSRStatus()
        self.bus_error = CSRStatus()
        self.overflow = CSRStatus()

        address_generator = BlockAddressGenerator(membus.aw, AXI_DATA_WIDTH)
        self.submodules += address_generator
        self.comb += [
            address_generator.base_address.eq(self.base_address.storage),
            address_generator.length.eq(self.length.storage),
            address_generator.start.eq(self.start.re)
        ]

        adc_writer = ADCWriter(membus.w, pads)
        self.submodules += adc_writer
        self.comb += [
            adc_writer.length.eq(self.length.storage),
            adc_writer.start.eq(self.start.re)
        ]
        self.sync += [
            If(self.start.re, self.overflow.status.eq(0)),
            If(adc_writer.overflow, self.overflow.status.eq(1))
        ]

        self.comb += self.busy.status.eq(address_generator.busy | adc_writer.busy)

        self.comb += membus.b.ready.eq(1)
        self.sync += [
            If(self.start.re, self.bus_error.status.eq(0)),
            If(membus.b.valid & membus.b.ready & (membus.b.resp != axi.Response.okay),
                self.bus_error.status.eq(1))
        ]