from migen.fhdl.std import *
from migen.actorlib.structuring import *
from migen.genlib.fsm import FSM, NextState

from misoclib.com.liteusb.common import *

class LiteUSBPacketizer(Module):
	def __init__(self):
		self.sink = sink = Sink(user_layout)
		self.source = source = Source(phy_layout)

		# Packet description
		#   - preamble : 4 bytes
		#   - dst      : 1 byte
		#   - length   : 4 bytes
		#   - payload
		header = [
			# preamble
			0x5A,
			0xA5,
			0x5A,
			0xA5,
			# dst
			sink.dst,
			# length
			sink.length[24:32],
			sink.length[16:24],
			sink.length[8:16],
			sink.length[0:8],
		]

		header_unpack = Unpack(len(header), phy_layout)
		self.submodules += header_unpack

		for i, byte in enumerate(header):
			chunk = getattr(header_unpack.sink.payload, "chunk" + str(i))
			self.comb += chunk.d.eq(byte)

		fsm = FSM()
		self.submodules += fsm

		fsm.act("WAIT_SOP",
			If(sink.stb & sink.sop, NextState("SEND_HEADER"))
		)

		fsm.act("SEND_HEADER",
			header_unpack.sink.stb.eq(1),
			source.stb.eq(1),
			source.d.eq(header_unpack.source.d),
			header_unpack.source.ack.eq(source.ack),
			If(header_unpack.sink.ack, NextState("SEND_DATA"))
		)

		fsm.act("SEND_DATA",
			source.stb.eq(sink.stb),
			source.d.eq(sink.d),
			sink.ack.eq(source.ack),
			If(source.ack & sink.eop, NextState("WAIT_SOP"))
		)

#
# TB
#
src_data = [
	(0x01, 4,
		[0x0, 0x1, 0x2, 0x3]
	),
	(0x16, 8,
		[0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7]
	),
	(0x22, 16,
		[0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF]
	),
]

class PacketizerSourceModel(Module, Source, RandRun):
	def __init__(self, data):
		Source.__init__(self, user_layout, True)
		RandRun.__init__(self, 25)
		self.data = data

		self._stb = 0
		self._sop = 0
		self._eop = 0
		self._frame_cnt = 0
		self._payload_cnt = 0

	def do_simulation(self, selfp):
		RandRun.do_simulation(self, selfp)
		dst, length, payload = self.data[self._frame_cnt]

		if selfp.stb and selfp.ack:
			if self._payload_cnt == length-1:
				self._frame_cnt += 1
				self._payload_cnt = 0
			else:
				self._payload_cnt += 1
			if self.run:
				self._stb = 1
			else:
				self._stb = 0

		if self.run and not self._stb:
			self._stb = 1

		self._sop = int((self._payload_cnt == 0))
		self._eop = int((self._payload_cnt == length-1))

		selfp.stb = self._stb
		selfp.sop = self._sop & self._stb
		selfp.eop = self._eop & self._stb
		selfp.dst = dst
		selfp.length = length
		selfp.d = payload[self._payload_cnt]

		if self._frame_cnt == len(self.data):
			raise StopSimulation

class PacketizerSinkModel(Module, Sink, RandRun):
	def __init__(self):
		Sink.__init__(self, phy_layout)
		RandRun.__init__(self, 25)

	def do_simulation(self, selfp):
		RandRun.do_simulation(self, selfp)
		if self.run:
			selfp.ack = 1
		else:
			selfp.ack = 0


class TB(Module):
	def __init__(self):
		self.submodules.source = PacketizerSourceModel(src_data)
		self.submodules.dut = LiteUSBPacketizer()
		self.submodules.sink = PacketizerSinkModel()

		self.comb +=[
			self.source.connect(self.dut.sink),
			self.dut.source.connect(self.sink),
		]

def main():
	from migen.sim.generic import run_simulation
	run_simulation(TB(), ncycles=400, vcd_name="tb_packetizer.vcd")

if __name__ == "__main__":
	main()