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


AXI_BURST_LEN = 16


class BlockAddressGenerator(Module):
    def __init__(self, membus_stream, data_width):
        address_width = len(membus_stream.addr)
        alignment_bits = log2_int(AXI_BURST_LEN*data_width//8)
        aligned_address_width = address_width - alignment_bits

        self.base_address = Signal(aligned_address_width)
        self.length = Signal(aligned_address_width)
        self.start = Signal()
        self.busy = Signal()

        current_address = Signal(aligned_address_width)
        remaining = Signal(aligned_address_width)

        self.comb += [
            membus_stream.addr.eq(Cat(C(0, 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(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")
                )
            )
        )


class ADCWriter(Module):
    def __init__(self, membus_stream, data_width, pads):
        self.busy = Signal()

        self.comb += [
            membus_stream.id.eq(0),
            membus_stream.valid.eq(1),
            #self.sink.ack.eq(membus.w.ready),
            membus_stream.data.eq(0x12345678deadbeef),
            membus_stream.strb.eq(2**(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)
                )
            )
        ]


class ADC(Module, AutoCSR):
    def __init__(self, membus, pads):
        data_width = len(membus.w.data)
        address_width = len(membus.aw.addr)
        alignment_bits = log2_int(AXI_BURST_LEN*data_width//8)

        self.base_address = CSRStorage(address_width, alignment_bits=alignment_bits)
        self.length = CSRStorage(address_width, alignment_bits=alignment_bits)
        self.start = CSR()
        self.busy = CSRStatus()
        self.bus_error = CSRStatus()
        
        address_generator = BlockAddressGenerator(membus.aw, 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, data_width, pads)
        self.submodules += adc_writer

        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))
        ]