"""Auxiliary controller, common to satellite and master"""

from operator import attrgetter

from migen import *
from migen.fhdl.simplify import FullMemoryWE

from misoc.interconnect.csr import *
from migen_axi.interconnect import axi
from artiq.gateware.drtio.aux_controller import Transmitter, Receiver

max_packet = 1024

OUT_BURST_LEN = 10
IN_BURST_LEN = 4

class SRAM(Module):
    def __init__(self, mem_or_size, read_only=False, init=None, bus=None):

        # SRAM initialisation

        if bus is None:
            bus = axi.Interface()
        self.bus = bus
        bus_data_width = len(self.bus.r.data)
        if isinstance(mem_or_size, Memory):
            assert(mem_or_size.width <= bus_data_width)
            self.mem = mem_or_size
        else:
            self.mem = Memory(bus_data_width, mem_or_size//(bus_data_width//8), init=init)

        # memory
        port = self.mem.get_port(write_capable=not read_only, we_granularity=8)
        self.specials += self.mem, port

        ###

        ar, aw, w, r, b = attrgetter("ar", "aw", "w", "r", "b")(bus)

        # Dout : Data received from CPU, output by SRAM <- port.dat_r
        # Din : Data driven into SRAM, written into CPU <- port.dat_w
        self.dout_index = Signal.like(ar.len)

        self.r_addr_incr = axi.Incr(ar)
        self.w_addr_incr = axi.Incr(aw)

        ### Read

        self.comb += [
            r.data.eq(port.dat_r),
            port.adr.eq(self.r_addr_incr.addr)
        ]

        # read control
        self.submodules.read_fsm = read_fsm = FSM(reset_state="IDLE")
        read_fsm.act("IDLE",
            If(ar.valid,
                port.adr.eq(self.r_addr_incr.addr),
                ar.ready.eq(1),
                NextState("READ_START"),
            )
        )
        read_fsm.act("READ_START",
            r.resp.eq(axi.Response.okay.value),
            r.valid.eq(1),
            If(r.ready,
                NextState("READ"))
        )
        read_fsm.act("READ",
            If(r.last & r.ready, # that's a smart way of skipping "LAST" state    
                NextState("IDLE")
            )
        )

        self.sync += [
            If(read_fsm.ongoing("IDLE"),
                self.dout_index.eq(0),
                r.valid.eq(0),  # shall it be reset too on IDLE?
                ar.ready.eq(0),
                r.last.eq(0)
            ).Else(If(r.ready & read_fsm.ongoing("READ"),
                    self.dout_index.eq(self.dout_index+1),
                    If(self.dout_index==ar.len, r.last.eq(1)) # and update last
                )
            )
        ]

        ### Write

        if not read_only:
            self.comb += [
                port.dat_w.eq(w.data),
                port.adr.eq(self.w_addr_incr.addr),
            ]

            self.submodules.write_fsm = write_fsm = FSM(reset_state="IDLE")
            write_fsm.act("IDLE",
                w.ready.eq(0),
                aw.ready.eq(0),
                b.valid.eq(0),
                If(aw.valid,
                    NextState("AW_VALID_WAIT")
                )
            )
            write_fsm.act("AW_VALID_WAIT",  # wait for data
                aw.ready.eq(1),
                If(w.valid,
                    NextState("WRITE"),
                )
            )
            # write_fsm.act("DATA_WAIT",
            #     aw.valid.eq(0),
            #     If(self.din_ready,
            #         w.valid.eq(1),
            #         NextState("WRITE")
            #     )
            # )
            write_fsm.act("WRITE",
                w.ready.eq(1),
                If(w.ready & w.last,
                    NextState("WRITE_RESP")
                )
            )

            write_fsm.act("WRITE_RESP",
                port.we.eq(0),
                b.resp.eq(axi.Response.okay.value),
                b.valid.eq(1),
                If(b.ready,
                    NextState("IDLE")
                )
            )

            self.sync += If(w.ready & w.valid, port.we.eq(1))

            # self.sync += [
            #     If(write_fsm.ongoing("IDLE"),
            #         self.din_index.eq(0) 
            #     ), # but need to synchronise the address too)
            # ]


        # # generate write enable signal
        # if not read_only:
        #     #                          replace with?      stb -> w.strb    we->w.ready?     sel[i]-> r.valid
        #     self.comb += [port.we[i].eq(self.bus.cyc & self.bus.w.strb & self.bus.w.ready & self.bus.r.valid)
        #         for i in range(4)]
        # # address and data
        # self.comb += [
        #     self.bus.r.ready.eq(self.bus.r.valid), # AXI handshake?
        #     port.adr.eq(self.bus.ar.addr[:len(port.adr)]),
        #     self.bus.r.data.eq(port.dat_r)
        # ]
        # if not read_only:
        #     self.comb += port.dat_w.eq(self.bus.w.data),
        # # generate ack
        # self.sync += [
        #     self.bus.ack.eq(0),
        #     If(self.bus.cyc & self.bus.stb & ~self.bus.ack, self.bus.ack.eq(1))
        # ]

# TODO: FullMemoryWE should be applied by migen.build
@FullMemoryWE()
class DRTIOAuxControllerAxi(Module):
    def __init__(self, link_layer):
        self.bus = axi.Interface()
        self.submodules.transmitter = Transmitter(link_layer, len(self.bus.w.data))
        self.submodules.receiver = Receiver(link_layer, len(self.bus.w.data))

        # probably will need to make axi.SRAM based on wb code
        tx_sdram_if = SRAM(self.transmitter.mem, read_only=False)
        rx_sdram_if = SRAM(self.receiver.mem, read_only=True)
        wsb = log2_int(len(self.bus.w.data)//8)
        decoder = axi.AddressDecoder(self.bus,
            [(lambda a: a[log2_int(max_packet)-wsb] == 0, tx_sdram_if.bus),
             (lambda a: a[log2_int(max_packet)-wsb] == 1, rx_sdram_if.bus)],
            register=True)
        self.submodules += tx_sdram_if, rx_sdram_if, decoder

    def get_csrs(self):
        return self.transmitter.get_csrs() + self.receiver.get_csrs()