from migen.fhdl.std import *
from migen.bank.description import *
from migen.genlib.fifo import SyncFIFOBuffered
from migen.genlib.cdc import MultiReg

from artiqlib.rtio.rbus import get_fine_ts_width

class _RTIOBankO(Module):
	def __init__(self, rbus, counter_width, fine_ts_width, fifo_depth, counter_init):
		self.sel = Signal(max=len(rbus))
		self.timestamp = Signal(counter_width+fine_ts_width)
		self.value = Signal(2)
		self.writable = Signal()
		self.we = Signal()
		self.underflow = Signal()
		self.level = Signal(bits_for(fifo_depth))

		###

		counter = Signal(counter_width, reset=counter_init)
		self.sync += [
			counter.eq(counter + 1),
			If(self.we & self.writable,
				If(self.timestamp[fine_ts_width:] < counter + 2, self.underflow.eq(1))
			)
		]

		fifos = []
		for n, chif in enumerate(rbus):
			fifo = SyncFIFOBuffered([
				("timestamp", counter_width+fine_ts_width), ("value", 2)],
				fifo_depth)
			self.submodules += fifo
			fifos.append(fifo)

			# FIFO write
			self.comb += [
				fifo.din.timestamp.eq(self.timestamp),
				fifo.din.value.eq(self.value),
				fifo.we.eq(self.we & (self.sel == n))
			]

			# FIFO read
			self.comb += [
				chif.o_stb.eq(fifo.readable &
					(fifo.dout.timestamp[fine_ts_width:] == counter)),
				chif.o_value.eq(fifo.dout.value),
				fifo.re.eq(chif.o_stb)
			]
			if fine_ts_width:
				self.comb += chif.o_fine_ts.eq(fifo.dout.timestamp[:fine_ts_width])

		selfifo = Array(fifos)[self.sel]
		self.comb += self.writable.eq(selfifo.writable), self.level.eq(selfifo.level)

class _RTIOBankI(Module):
	def __init__(self, rbus, counter_width, fine_ts_width, fifo_depth):
		self.sel = Signal(max=len(rbus))
		self.timestamp = Signal(counter_width+fine_ts_width)
		self.value = Signal()
		self.readable = Signal()
		self.re = Signal()
		self.overflow = Signal()

		###

		counter = Signal(counter_width)
		self.sync += counter.eq(counter + 1)

		timestamps = []
		values = []
		readables = []
		overflows = []
		for n, chif in enumerate(rbus):
			if hasattr(chif, "oe"):
				sensitivity = Signal(2)
				self.sync += If(~chif.oe & chif.o_stb,
					sensitivity.eq(chif.o_value))

				fifo = SyncFIFOBuffered([
					("timestamp", counter_width+fine_ts_width), ("value", 1)],
					fifo_depth)
				self.submodules += fifo
				
				# FIFO write
				if fine_ts_width:
					full_ts = Cat(chif.i_fine_ts, counter)
				else:
					full_ts = counter
				self.comb += [
					fifo.din.timestamp.eq(full_ts),
					fifo.din.value.eq(chif.i_value),
					fifo.we.eq(~chif.oe & chif.i_stb &
						((chif.i_value & sensitivity[0]) | (~chif.i_value & sensitivity[1])))
				]

				# FIFO read
				timestamps.append(fifo.dout.timestamp)
				values.append(fifo.dout.value)
				readables.append(fifo.readable)
				self.comb += fifo.re.eq(self.re & (self.sel == n))
				
				overflow = Signal()
				self.sync += If(fifo.we & ~fifo.writable, overflow.eq(1))
				overflows.append(overflow)
			else:
				timestamps.append(0)
				values.append(0)
				readables.append(0)
				overflows.append(0)

		self.comb += [
			self.timestamp.eq(Array(timestamps)[self.sel]),
			self.value.eq(Array(values)[self.sel]),
			self.readable.eq(Array(readables)[self.sel]),
			self.overflow.eq(Array(overflows)[self.sel])
		]

class RTIO(Module, AutoCSR):
	def __init__(self, phy, counter_width=32, ofifo_depth=8, ififo_depth=8):
		fine_ts_width = get_fine_ts_width(phy.rbus)

		# Submodules
		self.submodules.bank_o = InsertReset(_RTIOBankO(phy.rbus,
			counter_width, fine_ts_width, ofifo_depth,
			phy.loopback_latency))
		self.submodules.bank_i = InsertReset(_RTIOBankI(phy.rbus,
			counter_width, fine_ts_width, ofifo_depth))

		# CSRs
		self._r_reset = CSRStorage(reset=1)
		self._r_chan_sel = CSRStorage(flen(self.bank_o.sel))
		
		self._r_oe = CSR()

		self._r_o_timestamp = CSRStorage(counter_width+fine_ts_width)
		self._r_o_value = CSRStorage(2)
		self._r_o_writable = CSRStatus()
		self._r_o_we = CSR()
		self._r_o_underflow = CSRStatus()
		self._r_o_level = CSRStatus(bits_for(ofifo_depth))

		self._r_i_timestamp = CSRStatus(counter_width+fine_ts_width)
		self._r_i_value = CSRStatus()
		self._r_i_readable = CSRStatus()
		self._r_i_re = CSR()
		self._r_i_overflow = CSRStatus()

		# OE
		oes = []
		for n, chif in enumerate(phy.rbus):
			if hasattr(chif, "oe"):
				self.sync += \
					If(self._r_oe.re & (self._r_chan_sel.storage == n),
						chif.oe.eq(self._r_oe.r)
					)
				oes.append(chif.oe)
			else:
				oes.append(1)
		self.comb += self._r_oe.w.eq(Array(oes)[self._r_chan_sel.storage])

		# Output/Gate
		self.comb += [
			self.bank_o.reset.eq(self._r_reset.storage),
			self.bank_o.sel.eq(self._r_chan_sel.storage),
			self.bank_o.timestamp.eq(self._r_o_timestamp.storage),
			self.bank_o.value.eq(self._r_o_value.storage),
			self._r_o_writable.status.eq(self.bank_o.writable),
			self.bank_o.we.eq(self._r_o_we.re),
			self._r_o_underflow.status.eq(self.bank_o.underflow),
			self._r_o_level.status.eq(self.bank_o.level)
		]

		# Input
		self.comb += [
			self.bank_i.reset.eq(self._r_reset.storage),
			self.bank_i.sel.eq(self._r_chan_sel.storage),
			self._r_i_timestamp.status.eq(self.bank_i.timestamp),
			self._r_i_value.status.eq(self.bank_i.value),
			self._r_i_readable.status.eq(self.bank_i.readable),
			self.bank_i.re.eq(self._r_i_re.re),
			self._r_i_overflow.status.eq(self.bank_i.overflow)
		]