# SPDX-License-Identifier: BSD-2-Clause

from collections.abc import Generator, Iterable
from contextlib      import contextmanager
from itertools       import chain
from re              import search
from typing          import IO

from vcd             import VCDWriter
from vcd.gtkw        import GTKWSave
from vcd.writer      import Variable
from torii.hdl.ast   import Signal, SignalDict, Value
from torii.hdl.ir    import Fragment
from torii.sim._base import BaseEngine, BaseSignalState, BaseSimulation

from ._wasm_engine   import __version__, WASMInstance, WASMValue
from .wasmrtl        import WASMFragmentCompiler
from .wasmclock      import WASMClockProcess
from .wasmcoro       import WASMCoroProcess

__all__ = (
	'WASMSimEngine',
)

__version__ = __version__

class _VCDWriter:
	@staticmethod
	def decode_to_vcd(signal, value):
		return signal.decoder(value).expandtabs().replace(' ', '_')

	def __init__(
		self, fragment: Fragment, *, vcd_file: IO | str | None, gtkw_file: IO | str | None = None,
		traces: Iterable[Signal] = ()
	) -> None:

		self._close_vcd = False
		self._close_gtkw = False
		if isinstance(vcd_file, str):
			vcd_file = open(vcd_file, 'w')
			self._close_vcd = True

		if isinstance(gtkw_file, str):
			gtkw_file = open(gtkw_file, 'w')
			self._close_gtkw = True

		self.vcd_vars = SignalDict[Variable]()
		self.vcd_file = vcd_file
		self.vcd_writer = None
		if vcd_file is not None:
			self.vcd_writer = VCDWriter(vcd_file, timescale = '1 ps', comment = 'Generated by Torii')

		self.gtkw_names = SignalDict[tuple[str, ...]]()
		self.gtkw_file = gtkw_file
		self.gtkw_save = None
		if gtkw_file is not None:
			self.gtkw_save = GTKWSave(gtkw_file)

		self.traces = []

		signal_names = SignalDict[set[tuple[str, ...]]]()
		for subfragment, subfragment_name in fragment._assign_names_to_fragments(
			hierarchy = ('bench', 'top', )
		).items():
			for signal, signal_name in subfragment._assign_names_to_signals().items():
				if signal not in signal_names:
					signal_names[signal] = set[tuple[str, ...]]()
				signal_names[signal].add((*subfragment_name, signal_name))

		trace_names = SignalDict[set[tuple[str, ...]]]()
		assigned_names = set[str]()
		for trace in traces:
			trace_value = Value.cast(trace)
			for trace_signal in trace_value._rhs_signals():
				assert trace_signal is not None and isinstance(trace_signal, Signal)
				if trace_signal not in signal_names:
					if trace_signal.name not in assigned_names:
						name = trace_signal.name
					else:
						name = f'{trace_signal.name}${len(assigned_names)}'
						if name in assigned_names:
							raise RuntimeError(
								f'Unable to create unique signal name for {trace_signal.name}. '
								f'Tried {name}, but it already exists.'
							)
					trace_names[trace_signal] = { ('bench', name) }
					assigned_names.add(name)
				self.traces.append(trace_signal)

		if self.vcd_writer is None:
			return

		for signal, names in chain(signal_names.items(), trace_names.items()):
			if signal.decoder:
				var_type = 'string'
				var_size = 1
				var_init = self.decode_to_vcd(signal, signal.reset)
			else:
				var_type = 'wire'
				var_size = signal.width
				var_init = signal.reset

			for (*var_scope, var_name) in names:
				if search(r'[ \t\r\n]', var_name):
					raise NameError(f'Signal \'{".".join(var_scope)}.{var_name}\' contains a whitespace character')

				suffix = None
				while True:
					try:
						if suffix is None:
							var_name_suffix = var_name
						else:
							var_name_suffix = f'{var_name}${suffix}'
						if signal not in self.vcd_vars:
							vcd_var = self.vcd_writer.register_var(
								scope = var_scope, name = var_name_suffix,
								var_type = var_type, size = var_size, init = var_init)
							self.vcd_vars[signal] = vcd_var
						else:
							self.vcd_writer.register_alias(
								scope = var_scope, name = var_name_suffix,
								var = self.vcd_vars[signal])
						break
					except KeyError:
						suffix = (suffix or 0) + 1

				if signal not in self.gtkw_names:
					self.gtkw_names[signal] = (*var_scope, var_name_suffix)

	def update(self, timestamp, signal, value):
		vcd_var = self.vcd_vars.get(signal)
		if vcd_var is None:
			return

		if signal.decoder:
			var_value = self.decode_to_vcd(signal, value)
		else:
			var_value = value
		self.vcd_writer.change(vcd_var, timestamp, var_value)

	def close(self, timestamp):
		if self.vcd_writer is not None:
			self.vcd_writer.close(timestamp)

		if self.gtkw_save is not None:
			self.gtkw_save.dumpfile(self.vcd_file.name)
			self.gtkw_save.dumpfile_size(self.vcd_file.tell())

			self.gtkw_save.treeopen('top')
			for signal in self.traces:
				if len(signal) > 1 and not signal.decoder:
					suffix = f'[{len(signal) - 1}:0]'
				else:
					suffix = ''
				self.gtkw_save.trace('.'.join(self.gtkw_names[signal]) + suffix)

		if self._close_vcd:
			self.vcd_file.close()
		if self._close_gtkw:
			self.gtkw_file.close()

class _Timeline:
	def __init__(self) -> None:
		self.now = 0
		self.deadlines = dict()

	def reset(self):
		self.now = 0
		self.deadlines.clear()

	def at(self, run_at, process):
		if process in self.deadlines:
			raise ValueError(f'Process {process!r} already in deadline list')
		self.deadlines[process] = run_at

	def delay(self, delay_by, process):
		if delay_by is None:
			run_at = self.now
		else:
			run_at = self.now + delay_by
		self.at(run_at, process)

	def advance(self):
		nearest_processes = set()
		nearest_deadline = None
		for process, deadline in self.deadlines.items():
			if deadline is None:
				if nearest_deadline is not None:
					nearest_processes.clear()
				nearest_processes.add(process)
				nearest_deadline = self.now
				break
			elif nearest_deadline is None or deadline <= nearest_deadline:
				if deadline < self.now:
					raise ValueError(f'Deadline of {deadline} is behind our current state of {self.now}')

				if nearest_deadline is not None and deadline < nearest_deadline:
					nearest_processes.clear()
				nearest_processes.add(process)
				nearest_deadline = deadline

		if not nearest_processes:
			return False

		for process in nearest_processes:
			process.runnable = True
			del self.deadlines[process]
		self.now = nearest_deadline

		return True

class _WASMGlobal():
	def __init__(self, memory: WASMInstance, signal, offset, value) -> None:
		self._memory = memory
		self._signal = signal
		self._wasm   = WASMValue(memory, len(self._signal), offset, value)
		self._value  = value
		self.set(value)

	def set(self, value):
		self._value = value & ((1 << len(self._signal)) - 1)
		self._wasm.set(value)

	def get(self):
		self._wasm.get()

	def value(self):
		return self._value

	def update(self, value):
		self._value = value

	def __eq__(self, other):
		if isinstance(other, _WASMGlobal):
			return self.value() == other.value()
		else:
			return self.value() == int(other)

class _WASMSignalState(BaseSignalState):
	__slots__ = ('signal', 'curr', 'next', 'waiters', 'pending')

	def __init__(self, memory: WASMInstance, index, signal, pending) -> None:
		self.signal = signal
		self.pending = pending
		self.waiters = dict()
		self.curr = _WASMGlobal(memory, signal, index * 2, signal.reset)
		self.next = _WASMGlobal(memory, signal, index * 2 + 1, signal.reset)

	def set(self, value):
		self.next.update(value)
		self.pending.add(self)

	def update(self, value):
		raw_val = int(value)
		if self.next.value() == raw_val:
			return

		self.next.set(raw_val)
		self.pending.add(self)

	def commit(self):
		if self.curr == self.next:
			return False
		self.curr.set(self.next.value())

		awoken_any = False
		for process, trigger in self.waiters.items():
			if trigger is None or trigger == self.curr.value():
				process.runnable = awoken_any = True
		return awoken_any

class _WASMimulation(BaseSimulation):
	def __init__(self) -> None:
		self.timeline = _Timeline()
		self.signals  = SignalDict()
		self.slots    = []
		self.pending  = set()
		self.memory = WASMInstance()

	def set_slot(self, index, value):
		self.slots[index].set(value)

	def get_signal(self, signal):
		try:
			return self.signals[signal]
		except KeyError:
			index = len(self.slots)
			self.slots.append(_WASMSignalState(self.memory, index, signal, self.pending))
			self.signals[signal] = index
			return index

	def commit(self, changed = None):
		converged = True
		for signal_state in self.pending:
			if signal_state.commit():
				converged = False
		if changed is not None:
			changed.update(self.pending)
		self.pending.clear()
		return converged

	def wait_interval(self, process, interval):
		self.timeline.delay(interval, process)

	def add_trigger(self, process, signal, *, trigger = None):
		index = self.get_signal(signal)
		if process in self.slots[index].waiters and self.slots[index].waiters[process] != trigger:
			raise ValueError('Unable to add trigger for process!')
		self.slots[index].waiters[process] = trigger

	def remove_trigger(self, process, signal):
		index = self.get_signal(signal)
		if process not in self.slots[index].waiters:
			raise ValueError(f'Unable to remove trigger for process {process!r}, not in the slot list')
		del self.slots[index].waiters[process]

class WASMSimEngine(BaseEngine):
	def __init__(self, fragment: Fragment) -> None:
		self._state = _WASMimulation()
		self._timeline = self._state.timeline
		self._frag = fragment
		self._processes = WASMFragmentCompiler(self._state)(self._frag)
		self._vcd_writers = []

	def add_coroutine_process(self, process, *, default_cmd):
		self._processes.add(
			WASMCoroProcess(self._state, self._frag.domains, process, default_cmd = default_cmd)
		)

	def add_clock_process(self, clock, *, phase, period):
		self._processes.add(
			WASMClockProcess(self._state, clock, phase = phase, period = period)
		)

	def reset(self):
		pass

	def _step(self):
		changed = set() if self._vcd_writers else None
		# Performs the two phases of a delta cycle in a loop:
		converged = False
		while not converged:
			# 1. eval: run and suspend every non-waiting process once, queueing signal changes
			for process in self._processes:
				if process.runnable:
					process.runnable = False
					process.run()

			# 2. commit: apply every queued signal change, waking up any waiting processes
			converged = self._state.commit(changed)

		for vcd_writer in self._vcd_writers:
			for signal_state in changed:
				vcd_writer.update(self._timeline.now, signal_state.signal, signal_state.curr.value())

	def advance(self):
		self._step()
		self._timeline.advance()
		return any(not process.passive for process in self._processes)

	@property
	def now(self):
		return self._timeline.now

	@contextmanager
	def write_vcd(
		self, *, vcd_file: IO | str | None, gtkw_file: IO | str | None = None,
		traces: Iterable[Signal]
	) -> Generator[None, None, None]:
		vcd_writer = _VCDWriter(self._frag, vcd_file = vcd_file, gtkw_file = gtkw_file, traces = traces)
		try:
			self._vcd_writers.append(vcd_writer)
			yield
		finally:
			vcd_writer.close(self._timeline.now)
			self._vcd_writers.remove(vcd_writer)
