extends Node

class_name OSCReceiver


@export var default_address:String = "0.0.0.0"
@export var default_port: int = 9001
@export var poll_interval: float = 0.05
@export var debug: bool = false
var _server: UDPServer
var _observers: Dictionary
var _timer: Timer


func _init():
	_observers = {}

	_timer = Timer.new()
	_timer.autostart = false
	_timer.one_shot = false
	_timer.timeout.connect(self._poll)
	add_child(_timer)


func _exit_tree():
	stop_server()


func register_callback(oscaddress, argtypes, callback):
	if not _observers.has([oscaddress, argtypes]):
		_observers[[oscaddress, argtypes]] = []

	_observers[[oscaddress, argtypes]].append(callback)


func start_server(address:String = default_address, port:int = default_port):
	_server = UDPServer.new()

	if _server.listen(port, address) != OK:
		_debug("OSCReceiver could not bind to port: %s" % port)
	else:
		_debug("OSCReceiver listening on port: %s" % port)

	_timer.start(poll_interval)


func stop_server():
	_timer.stop()
	remove_child(_timer)
	_timer.free()

	if _server:
		_server.stop()


func _debug(msg):
	if debug:
		print(msg)


func _poll():
	if not _server.is_listening():
		return

	_server.poll()

	if _server.is_connection_available():
		var peer: PacketPeerUDP = _server.take_connection()
		var packet = peer.get_packet()
		var sender_ip = peer.get_packet_ip()
		var sender_port = peer.get_packet_port()

		_debug("Accepted peer: %s:%s" % [sender_ip, sender_port])

		var result = _parse_osc_addr_and_types(packet)
		var address = result[0]
		var types = result[1]
		var offset = result[2]

		_debug("OSC address: %s" % address)
		_debug("OSC arg types: %s" % types)

		if _observers.has([address, types]):
			var values = _parse_osc_values(packet.slice(offset), types)

			if values == null:
				_debug("Invalid/Unsupported OSC message.")
			elif values.size() != types.length():
				_debug("Mismatch between expected / received number of OSC arguments.")
			else:
				var msg_info = {
					"ip": sender_ip,
					"port": sender_port,
					"address": address,
					"types": types
				}

				for callback in _observers[[address, types]]:
					callback.call(msg_info, values)

func _parse_osc_addr_and_types(packet: PackedByteArray):
	var asep = packet.find(0)
	var address = packet.slice(0, asep).get_string_from_ascii() 
	
	var toffset = asep + (4 - asep % 4)
	assert(char(packet.decode_u8(toffset)) == ",")
	var tsep = packet.find(0, toffset)
	var types = packet.slice(toffset + 1, tsep).get_string_from_ascii() 

	return [address, types, tsep + (4 - tsep % 4)]


func _parse_osc_values(packet: PackedByteArray, types: String):
	var result
	var values = []
	var stream = StreamPeerBuffer.new()
	stream.set_data_array(packet)
	stream.set_big_endian(true)

	for type_id in types:
		match type_id:
			"i":
				values.append(stream.get_32())
			"h":
				values.append(stream.get_62())
			"f":
				values.append(stream.get_float())
			"d":
				values.append(stream.get_double())
			"c":
				values.append(char(stream.get_32()))
			"s", "S":
				var value = PackedStringArray()
				var null_found = false

				while not null_found:
					for _dummy in range(4):
						var ch = stream.get_u8()
						if not null_found and ch != 0:
							value.append(char(ch))
						else:
							null_found = true

				values.append("".join(value))
			"b":
				var count = stream.get_u32()
				result = stream.get_data(count)
				
				if result[0] == OK:
					values.append(result[1])

					if count % 4:
						stream.seek(stream.get_position() + (4 - count % 4))
				else:
					_debug("Could not read OSC blob argument.")
					return
			"t":
				result = stream.get_data(8)

				if result[0] == OK:
					values.append(result[1])
				else:
					_debug("Could not read OSC timetag argument.")
					return
			"m", "r":
				values.append([
					stream.get_u8(),
					stream.get_u8(),
					stream.get_u8(),
					stream.get_u8(),
				])
			"T":
				values.append(true)
			"F":
				values.append(false)
			"I", "N":
				values.append(null)
			_:
				_debug("Argument type '%s' not supported." % type_id)
				return

	return values