# SPDX-License-Identifier: BSD-3-Clause
#
# This file is part of Torii-USB.
#
# Copyright (c) 2020 Great Scott Gadgets <info@greatscottgadgets.com>

''' Full-device test harnesses for USB2. '''

from usb_construct.types import USBPacketID, USBStandardRequests

from ..interface.utmi    import UTMIInterface
from .                   import ToriiUSBGatewareTestCase
from .contrib            import usb_packet

class USBDeviceTest(ToriiUSBGatewareTestCase):
	''' Test case strap for UTMI-connected devices. '''

	# Use only the USB clock domain; ignore the sync one unless overridden.
	SYNC_CLOCK_FREQUENCY = None
	USB_CLOCK_FREQUENCY  = 60e6

	# The name of the argument to the DUT that will accept our UTMI bus.
	UTMI_BUS_ARGUMENT = 'bus'

	# The maximum number of NAKs we should allow in a loop before we bail out.
	# Keeping this reasonable prevents the simulation from running indefinitely.
	MAX_NAKS = 100

	def instantiate_dut(self):
		self.utmi    = UTMIInterface()

		# Vitals about the device.
		self.address             = 0
		self.max_packet_size_ep0 = 64

		# Always pass in the UTMI bus.
		arguments = self.FRAGMENT_ARGUMENTS.copy()
		arguments[self.UTMI_BUS_ARGUMENT] = self.utmi

		dut = self.FRAGMENT_UNDER_TEST(**arguments)
		self.provision_dut(dut)

		return dut

	def provision_dut(self, dut):
		'''
		Hook that allows us to add any desired properties to the DUT before simulation.

		This method is called before initial elaboration; so functions that modify devices
		before elaboration can be used.
		'''
		pass

	def provide_byte(self, byte):
		''' Provides a given byte on the UTMI receive data for one cycle. '''
		yield self.utmi.rx_data.eq(byte)
		yield

	def start_packet(self, *, set_rx_valid = True):
		''' Starts a UTMI packet receive. '''
		yield self.utmi.rx_active.eq(1)

		if set_rx_valid:
			yield self.utmi.rx_valid.eq(1)

		yield

	def end_packet(self):
		''' Starts a UTMI packet receive. '''
		yield self.utmi.rx_active.eq(0)
		yield self.utmi.rx_valid.eq(0)
		yield

	def provide_packet(self, *octets, cycle_after = True):
		''' Provides an entire packet transaction at once. '''

		yield from self.start_packet()

		for b in octets:
			yield from self.provide_byte(b)

		yield from self.end_packet()

		if cycle_after:
			yield

	@staticmethod
	def bits_to_octets(bits):
		'''
		Converts a string of bits to octets.

		For compatibility with antmicro's USB-testbench functions.
		'''

		bits   = bits[:]
		octets = []

		if (len(bits) % 8) != 0:
			raise ValueError('The length of bits must be divisible by 8')

		while bits:
			# Grab an octet worth of bits..
			octet_bits = bits[0:8][::-1]
			bits = bits[8:]

			# ... and parse it.
			octet = int(octet_bits, 2)
			octets.append(octet)

		return octets

	def provide_bits(self, bits, cycle_after = True):
		''' Provides an entire packet transaction at once; accepts bits. '''

		octets = self.bits_to_octets(bits)
		yield from self.provide_packet(*octets, cycle_after = cycle_after)

	def send_token(self, pid, *, endpoint = 0, address = None):
		'''
		Issues a token packet to the simulated USB device.

		Parameters
		----------
		pid
			The PID of the packet to be sent.

		endpoint
			The endpoint on which the token should be sent.

		address
			The address of the device to be targeted.
			If omitted, the most recently set-address'd address is used.

		'''

		# Grab the raw bits that make up our token from the Antmicro library...
		address = address if address else self.address
		bits    = usb_packet.token_packet(pid, address, endpoint)

		# ... and issue them ourselves.
		yield from self.provide_bits(bits)

	def send_data(self, pid, *octets):
		'''
		Sends a data packet to the simulated USB device.

		Parameters
		----------
		pid
			The PID to send the provided packet with.

		*data
			The data to be sent.

		'''

		bits = usb_packet.data_packet(pid, octets)
		yield from self.provide_bits(bits)

	def send_handshake(self, pid):
		'''
		Issues a handshake packet to the simulated USB device.

		Parameters
		----------
		pid
			The PID of the packet to be sent.

		'''

		# Ensure we have an USBPacketID-wrapped PID.
		pid = USBPacketID(pid)
		yield from self.provide_packet(pid.byte())

	def receive_packet(self, as_bytes = True, timeout = 1000):
		''' Receives a collection of data from the USB bus. '''

		data = []

		# Mark us as ready to receive data.
		yield self.utmi.tx_ready.eq(1)
		yield

		# Wait for the TX line to be valid.
		yield from self.wait_until(self.utmi.tx_valid, timeout = timeout)

		# For as long as tx_valid is asserted, capture data.
		while (yield self.utmi.tx_valid):
			data.append((yield self.utmi.tx_data))
			yield

		# Mark us as ready to receive data.
		yield self.utmi.tx_ready.eq(0)

		return bytes(data) if as_bytes else data

	#
	# More complex transaction helpers.
	#

	def interpacket_delay(self):
		''' Waits for a period appropriate between each packet. '''

		# We'll use a shorter value than real interpacket delays here,
		# to speed up simulation. This can be tuned longer if necessary.
		# A spec-valued length would be 10 for a FS device, or 1 for a HS.
		yield from self.advance_cycles(1)

	def out_transaction(self, *octets, endpoint = 0, token_pid = USBPacketID.OUT,
		data_pid = USBPacketID.DATA0, expect_handshake = None):
		'''
		Performs an OUT transaction.

		Parameters
		----------
		*octets
			The data to send.

		endpoint
			The endpoint on which to send the relevant data.

		token_pid
			The token PID to send. Defaults to OUT.

		data_pid
			The data PID to send. Defaults to DATA0.

		expect_handshake
			If provided, we'll assert that a given handshake is provided.

		Returns
		-------
			The handshake received.

		'''

		# Issue the token...
		yield from self.send_token(token_pid, endpoint = endpoint)
		yield from self.interpacket_delay()

		# ... issue our DATA packet...
		yield from self.send_data(data_pid, *octets)

		# ... and receive our handshake.
		data = yield from self.receive_packet()

		# If we expect the SETUP packet to be ACK'd, validate that.
		if expect_handshake:
			self.assertEqual(USBPacketID.from_byte(data), expect_handshake)

		yield from self.interpacket_delay()
		return USBPacketID.from_byte(data)

	def out_transfer(self, *octets, endpoint = 0, data_pid = USBPacketID.DATA0, max_packet_size = 64):
		'''
		Performs an OUT transfer.

		Parameters
		----------
		*octets
			The data to send.

		endpoint
			The endpoint on which to send the relevant data.

		data_pid
			The first data PID to send. Defaults to DATA0.

		max_packet_size
			The maximum packet size for the current endpoint.

		Returns
		-------
		The final handshake received.

		'''

		to_send  = octets[:]

		# If we're going to send all max-packet-length sized packet,
		# we'll add a ZLP.
		send_zlp = (len(octets) % max_packet_size) == 0

		while to_send:

			# Pull a packet out of the send queue...
			packet = to_send[0:max_packet_size]

			# ... and try to send it.
			handshake = yield from self.out_transaction(endpoint = endpoint, data_pid = data_pid, *packet)

			# If we're stalled, abort immediately.
			if handshake == USBPacketID.STALL:
				return USBPacketID.STALL

			# If our packet is NAK'd, don't continue.
			if handshake == USBPacketID.NAK:
				continue

			# Otherwise, advance in the stream...
			to_send = to_send[max_packet_size:]

			# ... and toggle DATA PIDs.
			data_pid = USBPacketID.DATA1 if (data_pid == USBPacketID.DATA0) else USBPacketID.DATA0

		# If we're going to send a ZLP, send it.
		if send_zlp:
			handshake = USBPacketID.NAK
			while handshake == USBPacketID.NAK:
				handshake = yield from self.out_transaction(endpoint = endpoint, data_pid = data_pid)

		return handshake

	def in_transaction(self, endpoint = 0, data_pid = None, handshake = USBPacketID.ACK):
		'''
		Performs an IN transaction.

		Parameters
		----------
		endpoint
			The endpoint on which to fetch the relevant data.

		data_pid
			The data PID to expect, or None if we don't are.

		handshake
			The response we should give after receiving data.

		Returns
		-------
		handshake
			The handshake retrieved in response.

		data
			A list of octets received.

		'''

		# Issue the IN token...
		yield from self.send_token(USBPacketID.IN, endpoint = endpoint)

		# ... receive our response...
		result = yield from self.receive_packet()
		try:
			pid, *data, crc_low, crc_high = result

		# If we can't unpack the response, just return the PID.
		except ValueError:
			return USBPacketID.from_int(result[0]), None

		# ... validate its CRC...
		expected_crc = usb_packet.crc16(data)
		self.assertEqual([crc_low, crc_high], expected_crc)

		# ... validate pid toggling, if desired...
		if data_pid:
			self.assertEqual(pid, data_pid.byte())

		# ... and issue our handshake.
		yield from self.interpacket_delay()
		yield from self.send_handshake(handshake)

		yield from self.interpacket_delay()
		return USBPacketID.from_int(pid), data

	def in_transfer(self, endpoint = 0, data_pid = None, handshake = USBPacketID.ACK):
		'''
		Performs an IN transfer.

		Parameters
		----------
		endpoint
			The endpoint on which to fetch the relevant data.

		data_pid
			The data PID to expect, or None if we don't are.

		handshake
			The response we should give after receiving data.

		Returns
		-------
		handshake
			The last data handshake received.

		data
			A list of octets received.

		'''

		naks = 0
		data = []

		while True:
			pid, packet = yield from self.in_transaction(
				endpoint = 0, data_pid = data_pid, handshake = USBPacketID.ACK)

			# If we were NAK'd, try again.
			if pid == USBPacketID.NAK:
				naks += 1
				self.assertLess(naks, self.MAX_NAKS)

				continue

			# If we were stalled, abort here.
			if pid == USBPacketID.STALL:
				return pid, None

			# Add the given packet to our transaction.
			data.extend(packet)

			# Swap to expecting the next packet ID.
			data_pid = USBPacketID.DATA1 if (data_pid == USBPacketID.DATA0) else USBPacketID.DATA0

			# If this is a short packet, stop receiving.
			if len(packet) < self.max_packet_size_ep0:
				break

		return pid, data

	def setup_transaction(self, request_type, request, value = 0, index = 0, length = 0):
		'''
		Sends a SETUP transaction. All arguments match their SETUP packet definitions.

		Returns the handshake received.
		'''

		def split(arg):
			''' Convenience function that splits a setup parameter. '''
			high = arg >> 8
			low  = arg & 0xFF

			return low, high

		# ... issue our DATA token...
		response = yield from self.out_transaction(
			request_type, request, *split(value), *split(index), *split(length),
			token_pid = USBPacketID.SETUP, expect_handshake = USBPacketID.ACK
		)

		return response

	def control_interphase_delay(self):
		''' Waits for a period appropriate between each delays. '''

		# This is shorter than would be normal, in order to speed up simulation.
		# If necessary, this can be extended arbitrarily.
		yield from self.advance_cycles(1)

	def control_request_in(self, request_type, request, value = 0, index = 0, length = 0):
		'''
		Performs an IN control request, and returns the results.

		Arguments match the SETUP packets.

		Returns
		-------
		handshake
			The handshake value returned.

		data
			A list of octets returned.

		'''

		naks = 0

		# If we don't have a data phase, treat this identically to an OUT control request.
		if length == 0:
			response = yield from self.control_request_out(request_type, request,
				value = value, index = index)
			return response

		#
		# Issue the Setup phase.
		#
		yield from self.setup_transaction(request_type, request, value, index, length)
		yield from self.control_interphase_delay()

		#
		# Issue the Data phase.
		#
		pid, packet = yield from self.in_transfer(data_pid = USBPacketID.DATA1)
		if pid == USBPacketID.STALL:
			return pid, None

		#
		# Issue the Status phase.
		#
		yield from self.control_interphase_delay()

		# Read until we get a status other than a NAK.
		handshake = USBPacketID.NAK
		while handshake == USBPacketID.NAK:
			naks += 1
			self.assertLess(naks, self.MAX_NAKS)

			handshake = yield from self.out_transaction(data_pid = USBPacketID.DATA1)

		# Finally, return our handshake and our data.
		return handshake, packet

	def control_request_out(self, request_type, request, value = 0, index = 0, data = ()):
		''' Performs an OUT control request, and returns the results.

		Arguments match the SETUP packets.
		'''

		naks = 0

		#
		# Issue the Setup phase.
		#
		yield from self.setup_transaction(request_type, request, value, index, len(data))
		yield from self.control_interphase_delay()

		#
		# Issue the Data phase, if we have one.
		#
		if data:
			pid = yield from self.out_transfer(data_pid = USBPacketID.DATA1, *data)
			if pid == USBPacketID.STALL:
				return pid

		#
		# Issue the Status phase.
		#
		yield from self.control_interphase_delay()

		# We'll perform an IN transfer, and expect a ZLP back.
		pid = USBPacketID.NAK
		while pid == USBPacketID.NAK:
			naks += 1
			self.assertLess(naks, self.MAX_NAKS)

			pid, data = yield from self.in_transaction(data_pid = USBPacketID.DATA1)

		# If this wasn't stalled, ensure we get a ZLP back.
		if pid != USBPacketID.STALL:
			self.assertEqual(data, [])

		# Finally, return our handshake.
		return pid

	def get_descriptor(self, descriptor_type, index = 0, length = 64):
		'''
		Performs a GET_DESCRIPTOR request; fetching a USB descriptor.

		Parameters
		----------
		descriptor_type
			The descriptor type number to fetch.

		index
			The descriptor index to fetch.

		length
			The requested length to read.

		'''

		# The type is stored in the MSB of the value; and the index in the LSB.
		value = descriptor_type << 8 | index
		descriptor = yield from self.control_request_in(0x80,
				USBStandardRequests.GET_DESCRIPTOR, value = value, length = length)

		return descriptor

	def set_address(self, new_address, update_address = True):
		'''
		Performs a SET_ADDRESS request; setting the device address.

		Parameters
		----------
		new_address
			The address to apply.

		'''

		response_pid = yield from self.control_request_out(0,
			USBStandardRequests.SET_ADDRESS, value = new_address)

		if update_address and (response_pid == USBPacketID.DATA1):
			self.address = new_address

		return response_pid

	def set_configuration(self, number):
		'''
		Performs a SET_CONFIGURATION request; configuring the device.

		Parameters
		----------
		number
			The configuration to select.

		'''

		response_pid = yield from self.control_request_out(0,
			USBStandardRequests.SET_CONFIGURATION, value = number)
		return response_pid

	def get_configuration(self):
		''' Performs a GET_CONFIGURATION request; reading the device's configuration. '''
		response = yield from self.control_request_in(0x80,
			USBStandardRequests.GET_CONFIGURATION, length = 1)
		return response
