g1-teleop/Main.gd

extends Node3D
## Main entry point for G1 Teleop Quest 3 app.
## Two-phase startup:
##   CONFIG phase: Dark VR environment with start screen UI panel
##   AR phase: Passthrough mixed reality with body tracking

enum Phase { CONFIG, AR }

@onready var body_tracker: Node = $BodyTracker
@onready var teleop_client: Node = $TeleopClient
@onready var xr_origin: XROrigin3D = $XROrigin3D
@onready var xr_camera: XRCamera3D = $XROrigin3D/XRCamera3D
@onready var webcam_quad: MeshInstance3D = $XROrigin3D/XRCamera3D/WebcamQuad
@onready var start_screen: Node3D = $XROrigin3D/StartScreen
@onready var left_controller: XRController3D = $XROrigin3D/LeftController
@onready var right_controller: XRController3D = $XROrigin3D/RightController
@onready var vr_pointer: Node3D = $VRUIPointer

var xr_interface: XRInterface
var xr_is_focused: bool = false
var current_phase: Phase = Phase.CONFIG
var _panel_positioned: bool = false


func _ready() -> void:
	# Hide webcam quad and start screen until positioned
	webcam_quad.visible = false
	start_screen.visible = false

	# Initialize OpenXR interface
	xr_interface = XRServer.find_interface("OpenXR")
	if xr_interface and xr_interface.is_initialized():
		print("[Main] OpenXR already initialized")
		_on_openxr_ready()
	elif xr_interface:
		xr_interface.connect("session_begun", _on_openxr_session_begun)
		xr_interface.connect("session_focussed", _on_openxr_focused)
		xr_interface.connect("session_stopping", _on_openxr_stopping)
		if not xr_interface.initialize():
			printerr("[Main] Failed to initialize OpenXR")
			get_tree().quit()
			return
		print("[Main] OpenXR initialized, waiting for session")
	else:
		printerr("[Main] OpenXR interface not found. Is the plugin enabled?")
		get_tree().quit()
		return

	# Enable XR on the viewport
	get_viewport().use_xr = true

	# Connect start screen signals
	start_screen.connect_requested.connect(_on_connect_requested)
	start_screen.launch_ar_requested.connect(_on_launch_ar_requested)

	# Connect teleop client connection state to start screen
	teleop_client.connection_state_changed.connect(_on_connection_state_changed)

	# Wire webcam frames (can happen anytime we're connected)
	teleop_client.webcam_frame_received.connect(webcam_quad._on_webcam_frame)

	# Setup VR pointer with references to controllers and XR origin
	vr_pointer.setup(xr_origin, xr_camera, left_controller, right_controller)

	# Setup body tracker visualization
	body_tracker.setup(xr_origin)

	print("[Main] Starting in CONFIG phase")


func _process(_delta: float) -> void:
	if not _panel_positioned and current_phase == Phase.CONFIG:
		# Wait until we have valid head tracking data to position the panel
		var hmd := XRServer.get_hmd_transform()
		if hmd.origin != Vector3.ZERO and hmd.origin.y > 0.3:
			_position_panel_in_front_of_user(hmd)
			_panel_positioned = true


func _position_panel_in_front_of_user(hmd: Transform3D) -> void:
	# Place the panel 1.2m in front of the user at their eye height
	var forward := -hmd.basis.z
	forward.y = 0  # Project onto horizontal plane
	if forward.length() < 0.01:
		forward = Vector3(0, 0, -1)
	forward = forward.normalized()

	var panel_pos := hmd.origin + forward * 1.2
	panel_pos.y = hmd.origin.y  # Same height as eyes

	start_screen.global_position = panel_pos
	# Face the panel toward the user
	# look_at() points -Z at target, but QuadMesh front face is +Z, so rotate 180
	start_screen.look_at(hmd.origin, Vector3.UP)
	start_screen.rotate_y(PI)
	start_screen.visible = true
	print("[Main] Panel positioned at %s (user at %s)" % [panel_pos, hmd.origin])


func _on_openxr_session_begun() -> void:
	print("[Main] OpenXR session begun")
	_on_openxr_ready()


func _on_openxr_ready() -> void:
	# In CONFIG phase, we stay in opaque/dark VR mode
	# Passthrough is only enabled when user clicks "Launch AR"
	if current_phase == Phase.AR:
		_enable_passthrough()


func _on_openxr_focused() -> void:
	xr_is_focused = true
	print("[Main] OpenXR session focused")


func _on_openxr_stopping() -> void:
	xr_is_focused = false
	print("[Main] OpenXR session stopping")


func _on_connect_requested(host: String, port: int) -> void:
	print("[Main] Connect requested: %s:%d" % [host, port])
	# If already connected, disconnect first
	if teleop_client.is_connected:
		teleop_client.disconnect_from_server()
		return

	teleop_client.server_host = host
	teleop_client.server_port = port
	teleop_client.connect_to_server()


func _on_connection_state_changed(connected: bool) -> void:
	start_screen.set_connected(connected)


func _on_launch_ar_requested() -> void:
	if current_phase == Phase.AR:
		return

	print("[Main] Launching AR mode")
	current_phase = Phase.AR

	# Enable passthrough
	_enable_passthrough()

	# Wire body tracker to teleop client
	body_tracker.tracking_data_ready.connect(teleop_client._on_tracking_data)

	# Show webcam quad, hide start screen
	webcam_quad.visible = true
	start_screen.hide_screen()


func _enable_passthrough() -> void:
	var openxr = xr_interface as OpenXRInterface
	if openxr:
		var modes = openxr.get_supported_environment_blend_modes()
		print("[Main] Supported blend modes: ", modes)
		if XRInterface.XR_ENV_BLEND_MODE_ALPHA_BLEND in modes:
			openxr.set_environment_blend_mode(XRInterface.XR_ENV_BLEND_MODE_ALPHA_BLEND)
			print("[Main] Passthrough enabled (alpha blend)")
			get_viewport().transparent_bg = true
			RenderingServer.set_default_clear_color(Color(0, 0, 0, 0))
		else:
			print("[Main] Alpha blend not supported, using opaque mode")