extends Node
class_name VRPN

# tracking associated data
enum TrackingDataType { POS_QUAT, VELOCITY, ACCELERATION }

# magic cookie
const magic_cookie_start : String = "vrpn: ver."
# kinda redundant as we take the seq number as well
static var header_size : int = aligned_size(20)


var sensors : Dictionary[int,String] = {}
var messages : Dictionary[int,String] = {}

signal connected(s:StreamPeerTCP)
signal data(data:Array)
signal disconnected
signal error(msg:String)


@onready var _stream: StreamPeerTCP = StreamPeerTCP.new()

@export var tracker_receivers : Array[VRPN_Receiver] = []

@export var vrpn_server : String = "127.0.0.1"
@export var vrpn_port : int = 3883

func _ready() -> void:
	
	self.connect_to_host(vrpn_server,vrpn_port)
	
	if not connected.has_connections():
		connected.connect(self._on_connected)
	if not data.has_connections():
		data.connect(self._on_data)
	if not disconnected.has_connections():
		disconnected.connect(self._on_disconnected)
	if not error.has_connections():
		error.connect(self._on_error)

func _process(delta: float) -> void:
	# store old state
	var old_status = _stream.get_status()
	# poll
	_stream.poll()
	# receive all errors
	var new_status = _stream.get_status()
	# process if something changed
	if old_status != new_status:
		match new_status:
			_stream.STATUS_NONE:
				emit_signal("disconnected")
			_stream.STATUS_CONNECTING:
				pass
			_stream.STATUS_CONNECTED:
				emit_signal("connected",_stream)
			_stream.STATUS_ERROR:
				emit_signal("error","error with socket stream from {0}:{1}".format([_stream.get_connected_host(),_stream.get_connected_port()]))

	if new_status == _stream.STATUS_CONNECTED:
		var available_bytes: int = _stream.get_available_bytes()
		if available_bytes > 0:
			var res = _stream.get_partial_data(available_bytes)
			if res[0] != OK:
				emit_signal("error","error receiving data {0}".format([res[0]]))
			else:
				emit_signal("data", res[1])

func connect_to_host(host: String, port: int) -> void:
	if _stream.connect_to_host(host, port) != OK:
		emit_signal("error","error connecting to host '{0}:{1}'".format([host,port]))

func send(data: PackedByteArray) -> bool:
	if _stream.get_status() != _stream.STATUS_CONNECTED:
		emit_signal("error","stream not connected!")
		return false
	var res: int = _stream.put_data(data)
	if res != OK:
		emit_signal("error","error writing to stream {0}!".format([res]))
		return false
	return true
	

func _on_data(data : Array):
	var bytes = PackedByteArray(data)
	var as_cookie = bytes.get_string_from_ascii()
	
	# Cookie Hack!
	if as_cookie.begins_with(VRPN.magic_cookie_start): # 
		# kaboom we just send back the same cookie :) 
		self.send(bytes)
	else:
		VRPN.marshall_block(bytes,self)


func _on_connected(s : StreamPeerTCP):
	print("Connected to",s.get_connected_host()) # Replace with function body.
	
func _on_disconnected():
	push_warning("Disconnected") # Replace with function body.

func _on_error(msg:String):
	push_warning(msg) # Replace with function body.
	

static func marshall_block(data : PackedByteArray,session : VRPN) -> void:
	
	# need to fix that 
	var block_offset : int = 0
	
	while data.size() > block_offset:
		# reader for stream
		var header := StreamPeerBuffer.new()
		# get block addresses
		header.data_array = data.slice(block_offset,block_offset+header_size)
		# make sure we read as big endian
		header.big_endian = true 
		
		# read header
		var length := header.get_32() as int # length of message
		var time_sec := header.get_32() as int  # datetime sec
		var time_msec := header.get_32() as int # datetime micro sec
		var sender_id := header.get_32() as int  # sender id (tracker or interfaces)
		var message_type := header.get_32() as int # type of message (payload)
		var sequence_num := header.get_32() as int # inofficial sequence number (padding)

		# for debugging
		#print("length       '%d'" % length)
		#print("time_sec     '%d'" % time_sec)
		#print("time_msec    '%d'" % time_msec)
		#print("sender_id    '%d'" % sender_id)
		#print("message_type '%d'" % message_type)
		#print("sequence_num '%d'" % sequence_num)
		
		marshall_body(data.slice(block_offset+header_size,block_offset+length),message_type,sender_id,session)
		
		# next datablock
		block_offset += aligned_size(length)
		
static func decode_string(stream : StreamPeerBuffer) -> String:
	var len = stream.get_32()
	return stream.get_string(len)
	
		
static func marshall_body(data : PackedByteArray,message_type : int, sender_id: int, session : VRPN):
	var body := StreamPeerBuffer.new()
	body.data_array = data
	body.big_endian = true

	# only take message_type directly for negative (-1,-2)
	# messages that provide dynamic descriptors
	if message_type < 0:
		# message and sender descriptions
		match message_type:
			-1: # sensor names
				var name = decode_string(body)
				print("sensor name is '%s' with '%d" % [name,sender_id])
				session.sensors[sender_id] = name
			-2: # message names
				var name = decode_string(body)
				print("message name is '%s' for message_type '%d'" % [name,sender_id])
				session.messages[sender_id] = name
		return
	
	# now we use the string identifiers 
	# because they are supposedly dynamically assigned
	match session.messages[message_type]:
		'vrpn_Tracker Pos_Quat': # quat pos
			# get id
			var sensor_id = body.get_32()
			var padding = body.get_32() # padding 
			# position
			var pos = Vector3(body.get_double(),body.get_double(),body.get_double())
			# VRPN quat layout and Godot Quaternion c'tor identical with x,y,z,w
			var quat = Quaternion(body.get_double(),body.get_double(),body.get_double(),body.get_double()).normalized()
			# submit to listener
			for r in session.tracker_receivers:
				r._on_tracker({
					"type" : TrackingDataType.POS_QUAT,
					"tracker" : session.sensors[sender_id],
					"sensor" : sensor_id,
					"position" : pos,
					"rotation" : quat
				})
		'vrpn_Tracker Velocity':
			# get id
			var sensor_id = body.get_32()
			var padding = body.get_32() # padding 
			# position
			var vel = Vector3(body.get_double(),body.get_double(),body.get_double())
			# VRPN quat layout and Godot Quaternion c'tor identical with x,y,z,w
			var vel_rot = Quaternion(body.get_double(),body.get_double(),body.get_double(),body.get_double()).normalized()
			# submit to listener
			for r in session.tracker_receivers:
				r._on_tracker({
					"type" : TrackingDataType.VELOCITY,
					"tracker" : session.sensors[sender_id],
					"sensor" : sensor_id,
					"velocity_linear" : vel,
					"velocity_rotation" : vel_rot
				})
		'vrpn_Tracker Acceleration':
			# get id
			var sensor_id = body.get_32()
			var padding = body.get_32() # padding 
			# position
			var acc = Vector3(body.get_double(),body.get_double(),body.get_double())
			# VRPN quat layout and Godot Quaternion c'tor identical with x,y,z,w
			var acc_rot = Quaternion(body.get_double(),body.get_double(),body.get_double(),body.get_double()).normalized()

			var acc_dt = body.get_double() # padding 
			# submit to listener
			for r in session.tracker_receivers:
				r._on_tracker({
					"type" : TrackingDataType.ACCELERATION,
					"tracker" : session.sensors[sender_id],
					"sensor" : sensor_id,
					"acceleration_linear" : acc,
					"acceleration_rotation" : acc_rot,
					"acceleration_dt" : acc_dt
				})
		_:
			pass
			

static func aligned_size(actual_size : int, alignment : int = 8) -> int:
	return (actual_size + alignment - 1) & ~(alignment - 1)