extends Node

class_name OSCReceiver


@export var default_address:String = "*"
@export var default_port: int = 9001
@export var poll_interval: float = 0.05
@export var max_packets_per_poll: int = 10
@export var debug: bool = false
var _server: UDPServer
var _observers: Dictionary
var _timer: Timer
# UNIX_EPOCH = 1970-01-01 00:00:00
# NTP_EPOCH = 1900-01-01 00:00:00
# NTP_DELTA = UNIX_EPOCH - NTP_EPOCH in seconds
const NTP_DELTA = 2208988800
const ISIZE = 4294967296  # 2**32

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(osc_address: String, arg_types: String, callback: Callable) -> void:
	if not _observers.has([osc_address, arg_types]):
		_observers[[osc_address, arg_types]] = []

	_observers[[osc_address, arg_types]].append(callback)


func start_server(port:int = default_port, bind_address:String = default_address) -> void:
	_server = UDPServer.new()

	if _server.listen(port, bind_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() -> void:
	_timer.stop()
	remove_child(_timer)
	_timer.free()

	if _server:
		_server.stop()


func _debug(msg) -> void:
	if debug:
		print(msg)


func _poll() -> void:
	if not _server.is_listening():
		return

	_server.poll()

	var num_packets = 0

	while num_packets < max_packets_per_poll and _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()
		num_packets += 1

		_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)

		var callbacks = _observers.get([address, types], [])
		# Get callbacks registered with types wild-card.
		callbacks.append_array(_observers.get([address, "*"], []))

		if callbacks:
			result = _parse_osc_values(packet.slice(offset), types)

			if result[0] != OK:
				_debug("Invalid/Unsupported OSC message.")
			elif result[1].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 callbacks:
					callback.call(msg_info, result[1])

func _parse_osc_addr_and_types(packet: PackedByteArray) -> Array:
	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) -> Array:
	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 [ERR_PARSE_ERROR]
			"t":
				var sec = stream.get_u32()
				var frac = stream.get_u32()
				values.append(to_time(sec, frac))
			"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 [ERR_INVALID_DATA]

	return [OK, values]


## Return seconds and fractional part of NTP timestamp as 2-item array.
func to_frac(timestamp) -> Array:
	var sec = int(timestamp)
	return [sec, int(abs(timestamp - sec) * ISIZE)]


## Return NTP timestamp from integer seconds and fractional part.
func to_time(sec: int, frac: int) -> float:
	return sec + float(frac) / ISIZE