g1-teleop/scripts/teleop_client.gd

extends Node
## WebSocket client that connects to teleop_server.py on the robot.
## Sends body tracking data as JSON, receives JPEG webcam frames as binary.
##
## Protocol:
##   Client → Server: JSON text messages with tracking data
##   Server → Client: Binary messages with JPEG webcam frames
##
## No SSL required — raw WebSocket over local network.

## Emitted when a JPEG webcam frame is received from the server.
## Camera IDs: 0=head, 1=left_wrist, 2=right_wrist
signal webcam_frame_received(jpeg_bytes: PackedByteArray)  ## Legacy (head only)
signal webcam_frame_head(jpeg_bytes: PackedByteArray)
signal webcam_frame_left(jpeg_bytes: PackedByteArray)
signal webcam_frame_right(jpeg_bytes: PackedByteArray)

## Emitted when connection state changes.
signal connection_state_changed(connected: bool)

## Server address — robot's IP and teleop_server.py port
@export var server_host: String = "10.0.0.64"
@export var server_port: int = 8765
@export var auto_connect: bool = true

## Reconnection settings
@export var reconnect_delay_sec: float = 2.0
@export var max_reconnect_attempts: int = 0  # 0 = unlimited

## Connection state
var ws := WebSocketPeer.new()
var is_connected: bool = false
var _reconnect_timer: float = 0.0
var _reconnect_attempts: int = 0
var _was_connected: bool = false
var _pending_data: Array = []

## Stats
var messages_sent: int = 0
var messages_received: int = 0
var bytes_sent: int = 0
var bytes_received: int = 0
var last_send_time: int = 0
var last_receive_time: int = 0


func _ready() -> void:
	if auto_connect:
		connect_to_server()


func _process(delta: float) -> void:
	ws.poll()

	var state := ws.get_ready_state()

	match state:
		WebSocketPeer.STATE_OPEN:
			if not is_connected:
				is_connected = true
				_was_connected = true
				_reconnect_attempts = 0
				print("[TeleopClient] Connected to ws://%s:%d" % [server_host, server_port])
				connection_state_changed.emit(true)

			# Process incoming messages
			while ws.get_available_packet_count() > 0:
				var packet := ws.get_packet()
				bytes_received += packet.size()
				messages_received += 1
				last_receive_time = Time.get_ticks_msec()

				# Check if this is a binary (JPEG) or text message
				if ws.was_string_packet():
					_handle_text_message(packet.get_string_from_utf8())
				else:
					# Binary = 1 byte camera_id + JPEG data
					if packet.size() > 1:
						var cam_id := packet[0]
						var jpeg_data := packet.slice(1)
						match cam_id:
							0:
								webcam_frame_head.emit(jpeg_data)
								webcam_frame_received.emit(jpeg_data)
							1:
								webcam_frame_left.emit(jpeg_data)
							2:
								webcam_frame_right.emit(jpeg_data)
							_:
								webcam_frame_received.emit(packet)

			# Send any pending tracking data
			for data in _pending_data:
				_send_json(data)
			_pending_data.clear()

		WebSocketPeer.STATE_CLOSING:
			pass  # Wait for close to complete

		WebSocketPeer.STATE_CLOSED:
			if is_connected or _was_connected:
				var code := ws.get_close_code()
				var reason := ws.get_close_reason()
				print("[TeleopClient] Disconnected (code=%d, reason=%s)" % [code, reason])
				is_connected = false
				connection_state_changed.emit(false)

			# Auto-reconnect
			if auto_connect and _was_connected:
				_reconnect_timer -= delta
				if _reconnect_timer <= 0:
					_reconnect_timer = reconnect_delay_sec
					_attempt_reconnect()

		WebSocketPeer.STATE_CONNECTING:
			pass  # Still connecting


func connect_to_server() -> void:
	var url := "ws://%s:%d" % [server_host, server_port]
	print("[TeleopClient] Connecting to %s..." % url)
	var err := ws.connect_to_url(url)
	if err != OK:
		printerr("[TeleopClient] Failed to initiate connection: ", err)


func disconnect_from_server() -> void:
	_was_connected = false
	ws.close()


func _attempt_reconnect() -> void:
	if max_reconnect_attempts > 0:
		_reconnect_attempts += 1
		if _reconnect_attempts > max_reconnect_attempts:
			print("[TeleopClient] Max reconnect attempts reached, giving up")
			_was_connected = false
			return
		print("[TeleopClient] Reconnect attempt %d/%d..." % [_reconnect_attempts, max_reconnect_attempts])
	else:
		_reconnect_attempts += 1
		if _reconnect_attempts % 5 == 1:  # Log every 5th attempt to avoid spam
			print("[TeleopClient] Reconnect attempt %d..." % _reconnect_attempts)

	connect_to_server()


## Called by body_tracker via signal when new tracking data is ready.
func _on_tracking_data(data: Dictionary) -> void:
	if is_connected:
		_send_json(data)
	# Don't buffer if not connected — tracking data is perishable


## Send a command message to the server (e.g. "recalibrate").
func send_command(command: String) -> void:
	if is_connected:
		_send_json({"type": command})


func _send_json(data: Dictionary) -> void:
	var json_str := JSON.stringify(data)
	var err := ws.send_text(json_str)
	if err == OK:
		messages_sent += 1
		bytes_sent += json_str.length()
		last_send_time = Time.get_ticks_msec()
	else:
		printerr("[TeleopClient] Failed to send: ", err)


func _handle_text_message(text: String) -> void:
	# Server may send JSON status/config messages
	var json := JSON.new()
	var err := json.parse(text)
	if err != OK:
		printerr("[TeleopClient] Invalid JSON from server: ", text.left(100))
		return

	var data: Dictionary = json.data
	var msg_type: String = data.get("type", "")

	match msg_type:
		"config":
			print("[TeleopClient] Server config: ", data)
		"status":
			print("[TeleopClient] Server status: ", data.get("message", ""))
		_:
			print("[TeleopClient] Unknown message type: ", msg_type)


## Get connection statistics as a dictionary
func get_stats() -> Dictionary:
	return {
		"connected": is_connected,
		"messages_sent": messages_sent,
		"messages_received": messages_received,
		"bytes_sent": bytes_sent,
		"bytes_received": bytes_received,
		"reconnect_attempts": _reconnect_attempts,
	}