# SPDX-License-Identifier: BSD-3-Clause
#
# This file is part of Torii-USB.
#
# Copyright (c) 2020 Great Scott Gadgets <info@greatscottgadgets.com>
# Copyright (c) 2020 Florent Kermarrec <florent@enjoy-digital.fr>
#
# Code based on ``usb3_pipe``.

''' Scrambling and descrambling for USB3. '''

from torii.hdl import Cat, Elaboratable, Module, Signal

from ...stream import USBRawSuperSpeedStream
from .coding   import COM, stream_word_matches_symbol

#
# Scrambling modules.
# See [USB3.2r1: Appendix B].
#

class ScramblerLFSR(Elaboratable):
	''' Scrambler LFSR.

	Linear feedback shift register used for USB3 scrambling.
	Polynomial: X^16 + X^5 + X^4 + X^3 + 1

	See [USB3.2: Appendix B]

	Attributes
	----------
	clear: Signal(), input
		Strobe; when high, resets the LFSR to its initial value.
	advance: Signal(), input
		Strobe; when high, the LFSR advances on each clock cycle.
	value: Signal(32), output
		The current value of the LFSR.

	Parameters
	----------
	initial_value: 32-bit int, optional
		The initial value for the LFSR. Optional; defaults to all 1's, per the USB3 spec.
	'''
	def __init__(self, initial_value = 0xffff):
		self._initial_value = initial_value

		#
		# I/O port
		#
		self.clear   = Signal()
		self.advance = Signal()
		self.value   = Signal(32)

	def elaborate(self, platform):
		m = Module()

		next_value       = Signal(16)
		current_value    = Signal(16, reset = self._initial_value)

		def xor_bits(*indices):
			bits = Cat(current_value[i] for i in indices)
			return bits.xor()

		# Compute the next value in our internal LFSR state...
		m.d.comb += next_value.eq(Cat(
			xor_bits(0, 6, 8, 10),               # 0
			xor_bits(1, 7, 9, 11),               # 1
			xor_bits(2, 8, 10, 12),              # 2
			xor_bits(3, 6, 8, 9, 10, 11, 13),    # 3
			xor_bits(4, 6, 7, 8, 9, 11, 12, 14), # 4
			xor_bits(5, 6, 7, 9, 12, 13, 15),    # 5
			xor_bits(0, 6, 7, 8, 10, 13, 14),    # 6
			xor_bits(1, 7, 8, 9, 11, 14, 15),    # 7
			xor_bits(0, 2, 8, 9, 10, 12, 15),    # 8
			xor_bits(1, 3, 9, 10, 11, 13),       # 9
			xor_bits(0, 2, 4, 10, 11, 12, 14),   # 10
			xor_bits(1, 3, 5, 11, 12, 13, 15),   # 11
			xor_bits(2, 4, 6, 12, 13, 14),       # 12
			xor_bits(3, 5, 7, 13, 14, 15),       # 13
			xor_bits(4, 6, 8, 14, 15),           # 14
			xor_bits(5, 7, 9, 15)                # 15
		))

		# Compute the LFSR's current output.
		m.d.comb += self.value.eq(Cat(
			current_value[15],
			current_value[14],
			current_value[13],
			current_value[12],
			current_value[11],
			current_value[10],
			current_value[9],
			current_value[8],
			current_value[7],
			current_value[6],
			current_value[5],
			xor_bits(4,  15),
			xor_bits(3,  14, 15),
			xor_bits(2,  13, 14, 15),
			xor_bits(1,  12, 13, 14),
			xor_bits(0,  11, 12, 13),
			xor_bits(10, 11, 12, 15),
			xor_bits(9,  10, 11, 14),
			xor_bits(8,  9 , 10, 13),
			xor_bits(7,  8 , 9,  12),
			xor_bits(6,  7 , 8,  11),
			xor_bits(5,  6 , 7,  10),
			xor_bits(4,  5 , 6,  9,  15),
			xor_bits(3,  4 , 5,  8,  14),
			xor_bits(2,  3 , 4,  7,  13, 15),
			xor_bits(1,  2 , 3,  6,  12, 14),
			xor_bits(0,  1 , 2,  5,  11, 13, 15),
			xor_bits(0,  1 , 4,  10, 12, 14),
			xor_bits(0,  3 , 9,  11, 13),
			xor_bits(2,  8 , 10, 12),
			xor_bits(1,  7 , 9,  11),
			xor_bits(0,  6 , 8,  10)
		))

		# If we have a reset, clear our LFSR.
		with m.If(self.clear):
			m.d.ss += current_value.eq(self._initial_value)

		# Otherwise, advance when desired.
		with m.Elif(self.advance):
			m.d.ss += current_value.eq(next_value)

		return m

class Scrambler(Elaboratable):
	''' USB3-compliant data scrambler.

	Scrambles the transmitted data stream to reduce EMI.

	Attributes
	----------
	clear: Signal(), input
		Strobe; when high, resets the scrambler to the start of its sequence.
	enable: Signal(), input
		When high, data scrambling is enabled. When low, data is passed through without scrambling.
	sink: USBRawSuperSpeedStream(), input stream
		The stream containing data to be scrambled.
	sink: USBRawSuperSpeedStream(), output stream
		The stream containing data the scrambled output.

	Parameters
	----------
	initial_value: 32-bit int, optional
		The initial value for the LFSR. Optional.
	'''
	def __init__(self, initial_value = 0x7dbd):
		self._initial_value = initial_value

		#
		# I/O port
		#
		self.clear  = Signal()
		self.enable = Signal()
		self.hold   = Signal()

		self.sink   = USBRawSuperSpeedStream()
		self.source = USBRawSuperSpeedStream()

		# Debug signaling.
		self.lfsr_state = Signal.like(self.source.data)

	def elaborate(self, platform):
		m = Module()

		sink   = self.sink
		source = self.source

		# Detect when we're sending a comma; which should reset our scrambling LFSR.
		comma_present = stream_word_matches_symbol(sink, 0, symbol = COM)

		# Create our inner LFSR, which should advance whenever our input streams do.
		m.submodules.lfsr = lfsr = ScramblerLFSR(initial_value = self._initial_value)
		m.d.comb += [
			lfsr.clear.eq(self.clear | comma_present),
			lfsr.advance.eq(sink.valid & source.ready & ~self.hold)
		]

		# Pass through non-scrambled signals directly.
		m.d.comb += [
			source.ctrl.eq(sink.ctrl),
			source.valid.eq(sink.valid),
			sink.ready.eq(source.ready)
		]

		# If we have any non-control words, scramble them by overriding our data assignment above
		# with the relevant data word XOR'd with our LFSR value. Note that control words are -never-
		# scrambled, per [USB3.2: Appendix B]
		for i in range(4):
			is_data_code = ~sink.ctrl[i]
			lfsr_word    = lfsr.value.word_select(i, 8)

			with m.If(self.enable & is_data_code):
				m.d.comb += source.data.word_select(i, 8).eq(sink.data.word_select(i, 8) ^ lfsr_word)
			with m.Else():
				m.d.comb += source.data.word_select(i, 8).eq(sink.data.word_select(i, 8))

		# Connect up our debug outputs.
		m.d.comb += [
			self.lfsr_state.eq(lfsr.value)
		]

		return m

class Descrambler(Scrambler):
	''' USB3-compliant data descrambler.

	This module descrambles the received data stream. K-codes are not affected.
	This module automatically resets itself whenever a COM alignment character is seen.

	Attributes
	----------
	enable: Signal(), input
		When high, data scrambling is enabled. When low, data is passed through without scrambling.
	sink: USBRawSuperSpeedStream(), input stream
		The stream containing data to be descrambled.
	source: USBRawSuperSpeedStream(), output stream
		The stream containing data the descrambled output.

	Parameters
	----------
	initial_value: 32-bit int, optional
		The initial value for the LFSR. Optional.

	'''
	def __init__(self, initial_value = 0xffff):
		self._initial_value = initial_value
		super().__init__(initial_value = initial_value)
