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/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

# Spinning G1 models flanking the start screen
var _g1_model_left: Node3D
var _g1_model_right: Node3D
var _config_light: DirectionalLight3D
const G1_SPIN_SPEED := 0.5  # radians per second

# Gaze balls: [exit_ar, slot_1, slot_2, quit_app] — left to right
var _gaze_balls: Array = []       # Array of MeshInstance3D
var _gaze_mats: Array = []        # Array of StandardMaterial3D
var _gaze_timers: Array = [0.0, 0.0, 0.0, 0.0]
var _gaze_laser: MeshInstance3D
const GAZE_BALL_COUNT := 4
const GAZE_ACTIVATE_TIME := 5.0
const GAZE_ANGLE_THRESHOLD := 8.0  # degrees
const GAZE_BALL_BASE_SCALE := 1.0
const GAZE_BALL_MAX_SCALE := 1.8
const GAZE_BALL_COLORS: Array = [
	Color(1.0, 0.2, 0.2, 0.35),   # 0: red — exit AR
	Color(0.8, 0.6, 0.2, 0.35),   # 1: yellow — reserved
	Color(0.2, 0.8, 0.4, 0.35),   # 2: green — reserved
	Color(0.2, 0.3, 1.0, 0.35),   # 3: blue — quit app
]
const GAZE_BALL_BASE_COLORS: Array = [
	Color(1.0, 0.2, 0.2),
	Color(0.8, 0.6, 0.2),
	Color(0.2, 0.8, 0.4),
	Color(0.2, 0.3, 1.0),
]


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")
		xr_interface.connect("pose_recentered", _on_pose_recentered)
		_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)
		xr_interface.connect("pose_recentered", _on_pose_recentered)
		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

	# CONFIG phase: keep background opaque so passthrough is hidden
	# (blend mode is alpha_blend from project settings, but we render opaque black)
	get_viewport().transparent_bg = false
	RenderingServer.set_default_clear_color(Color(0, 0, 0, 1))

	# 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

	# Spin G1 models in CONFIG phase
	if current_phase == Phase.CONFIG:
		if _g1_model_left and _g1_model_left.visible:
			_g1_model_left.rotate_y(G1_SPIN_SPEED * delta)
		if _g1_model_right and _g1_model_right.visible:
			_g1_model_right.rotate_y(G1_SPIN_SPEED * delta)

	if current_phase == Phase.AR:
		_update_gaze_balls(delta)


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

	# Position spinning G1 models on each side of the panel
	var right := forward.cross(Vector3.UP).normalized()
	_setup_g1_models(panel_pos, right, hmd.origin)


func _setup_g1_models(panel_pos: Vector3, right: Vector3, user_pos: Vector3) -> void:
	var g1_scene = load("res://models/g1_full.obj")
	if g1_scene == null:
		printerr("[Main] Failed to load G1 model")
		return

	if _g1_model_left == null:
		_g1_model_left = _create_g1_instance(g1_scene)
		add_child(_g1_model_left)
	if _g1_model_right == null:
		_g1_model_right = _create_g1_instance(g1_scene)
		add_child(_g1_model_right)

	# Place 0.7m to each side of the panel, slightly below eye height
	_g1_model_left.global_position = panel_pos - right * 0.7 + Vector3.DOWN * 0.3
	_g1_model_right.global_position = panel_pos + right * 0.7 + Vector3.DOWN * 0.3

	# Reset rotation to just the upright correction, then face user
	_g1_model_left.rotation = Vector3(-PI / 2.0, 0, 0)
	_g1_model_right.rotation = Vector3(-PI / 2.0, 0, 0)

	_g1_model_left.visible = true
	_g1_model_right.visible = true

	# Add a directional light for shading the models
	if _config_light == null:
		_config_light = DirectionalLight3D.new()
		_config_light.light_energy = 1.0
		_config_light.rotation_degrees = Vector3(-45, 30, 0)
		add_child(_config_light)
	_config_light.visible = true
	print("[Main] G1 models positioned")


func _create_g1_instance(mesh_resource: Mesh) -> MeshInstance3D:
	var inst := MeshInstance3D.new()
	inst.mesh = mesh_resource
	# Scale down to ~0.3m tall (adjust if needed)
	inst.scale = Vector3(0.0015, 0.0015, 0.0015)
	# Rotate 90 degrees around X to stand the model upright (OBJ has Z-up)
	inst.rotate_x(-PI / 2.0)
	inst.visible = false
	# Add a shaded material so the model is visible with lighting
	var mat := StandardMaterial3D.new()
	mat.albedo_color = Color(0.4, 0.6, 0.9, 1.0)  # Light blue
	inst.material_override = mat
	return inst


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_pose_recentered() -> void:
	print("[Main] Pose recentered — repositioning UI")
	var hmd := XRServer.get_hmd_transform()
	if current_phase == Phase.CONFIG:
		_position_panel_in_front_of_user(hmd)
	elif current_phase == Phase.AR:
		_position_quad_in_front_of_user(webcam_quad, hmd, 1.2, -0.3)
		_create_gaze_balls(hmd)


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 (only connect once)
	if not body_tracker.tracking_data_ready.is_connected(teleop_client._on_tracking_data):
		body_tracker.tracking_data_ready.connect(teleop_client._on_tracking_data)

	# Position webcam quad in front of user (stationary)
	var hmd := XRServer.get_hmd_transform()
	_position_quad_in_front_of_user(webcam_quad, hmd, 1.2, -0.3)
	webcam_quad.visible = true

	# Hide start screen and G1 models
	start_screen.hide_screen()
	if _g1_model_left:
		_g1_model_left.visible = false
	if _g1_model_right:
		_g1_model_right.visible = false
	if _config_light:
		_config_light.visible = false

	# Create and position gaze exit balls
	_create_gaze_balls(hmd)


func _enable_passthrough() -> void:
	# Enable Meta Quest passthrough via XR_FB_PASSTHROUGH extension.
	# Requires openxr/extensions/meta/passthrough=true in project settings
	# and meta_xr_features/passthrough=2 in export presets.

	# Make viewport transparent so passthrough shows through
	get_viewport().transparent_bg = true
	RenderingServer.set_default_clear_color(Color(0, 0, 0, 0))

	# Set blend mode to alpha blend — the Meta vendor plugin intercepts this
	# and activates FB passthrough even though it's not in supported modes list.
	var openxr = xr_interface as OpenXRInterface
	if openxr:
		openxr.environment_blend_mode = XRInterface.XR_ENV_BLEND_MODE_ALPHA_BLEND
		print("[Main] Passthrough enabled (alpha blend + transparent bg)")


func _disable_passthrough() -> void:
	get_viewport().transparent_bg = false
	RenderingServer.set_default_clear_color(Color(0, 0, 0, 1))

	var openxr = xr_interface as OpenXRInterface
	if openxr:
		openxr.environment_blend_mode = XRInterface.XR_ENV_BLEND_MODE_OPAQUE
		print("[Main] Passthrough disabled")


func _position_quad_in_front_of_user(quad: Node3D, hmd: Transform3D, distance: float, y_offset: float) -> void:
	var forward := -hmd.basis.z
	forward.y = 0
	if forward.length() < 0.01:
		forward = Vector3(0, 0, -1)
	forward = forward.normalized()

	var pos := hmd.origin + forward * distance
	pos.y = hmd.origin.y + y_offset

	quad.global_position = pos
	quad.look_at(hmd.origin, Vector3.UP)
	quad.rotate_y(PI)


func _create_gaze_balls(hmd: Transform3D) -> void:
	var forward := -hmd.basis.z
	forward.y = 0
	if forward.length() < 0.01:
		forward = Vector3(0, 0, -1)
	forward = forward.normalized()
	var right := forward.cross(Vector3.UP).normalized()

	var base_pos := hmd.origin + forward * 1.0 + Vector3.UP * 0.35

	# Create balls if first time, spread evenly left to right
	if _gaze_balls.is_empty():
		for i in range(GAZE_BALL_COUNT):
			var ball := _make_gaze_sphere(GAZE_BALL_COLORS[i])
			_gaze_balls.append(ball)
			_gaze_mats.append(ball.material_override)
			add_child(ball)

	# Position: spread from -0.45 to +0.45 across the row
	var spread := 0.3  # spacing between balls
	var half_width := spread * (GAZE_BALL_COUNT - 1) / 2.0
	for i in range(GAZE_BALL_COUNT):
		var offset := -half_width + spread * i
		_gaze_balls[i].global_position = base_pos + right * offset
		_gaze_balls[i].visible = true
		_gaze_balls[i].scale = Vector3.ONE
		_gaze_timers[i] = 0.0

	# Gaze laser beam
	if _gaze_laser == null:
		_gaze_laser = MeshInstance3D.new()
		var cyl := CylinderMesh.new()
		cyl.top_radius = 0.002
		cyl.bottom_radius = 0.002
		cyl.height = 1.0
		_gaze_laser.mesh = cyl
		var lmat := StandardMaterial3D.new()
		lmat.shading_mode = BaseMaterial3D.SHADING_MODE_UNSHADED
		lmat.albedo_color = Color(1.0, 1.0, 1.0, 0.4)
		lmat.transparency = BaseMaterial3D.TRANSPARENCY_ALPHA
		lmat.no_depth_test = true
		_gaze_laser.material_override = lmat
		add_child(_gaze_laser)
	_gaze_laser.visible = false


func _make_gaze_sphere(color: Color) -> MeshInstance3D:
	var mesh := MeshInstance3D.new()
	var sphere := SphereMesh.new()
	sphere.radius = 0.035
	sphere.height = 0.07
	mesh.mesh = sphere
	var mat := StandardMaterial3D.new()
	mat.shading_mode = BaseMaterial3D.SHADING_MODE_UNSHADED
	mat.albedo_color = color
	mat.transparency = BaseMaterial3D.TRANSPARENCY_ALPHA
	mat.no_depth_test = true
	mesh.material_override = mat
	return mesh


func _update_gaze_balls(delta: float) -> void:
	if _gaze_balls.is_empty() or not _gaze_balls[0].visible:
		return

	var hmd := XRServer.get_hmd_transform()
	if hmd.origin == Vector3.ZERO:
		return

	var gaze_dir := (-hmd.basis.z).normalized()
	var threshold := deg_to_rad(GAZE_ANGLE_THRESHOLD)
	var gazing_at_any := false

	for i in range(GAZE_BALL_COUNT):
		var ball: MeshInstance3D = _gaze_balls[i]
		var to_ball: Vector3 = (ball.global_position - hmd.origin).normalized()
		var looking := gaze_dir.angle_to(to_ball) < threshold
		if looking:
			_gaze_timers[i] += delta
			if not gazing_at_any:
				gazing_at_any = true
				_show_gaze_laser(hmd.origin, ball.global_position)
		else:
			_gaze_timers[i] = 0.0
		_update_gaze_visual(_gaze_mats[i] as StandardMaterial3D, GAZE_BALL_BASE_COLORS[i] as Color, _gaze_timers[i], ball)

		if _gaze_timers[i] >= GAZE_ACTIVATE_TIME:
			_on_gaze_activated(i)
			return

	if not gazing_at_any and _gaze_laser:
		_gaze_laser.visible = false


func _on_gaze_activated(index: int) -> void:
	match index:
		0:  # Exit AR
			print("[Main] Gaze exit: returning to CONFIG")
			_exit_ar_mode()
		3:  # Quit app
			print("[Main] Gaze exit: quitting app")
			get_tree().quit()
		_:  # Reserved slots — no action yet
			print("[Main] Gaze ball %d activated (no action assigned)" % index)
			_gaze_timers[index] = 0.0


func _update_gaze_visual(mat: StandardMaterial3D, base_color: Color, timer: float, ball: MeshInstance3D) -> void:
	var progress := clampf(timer / GAZE_ACTIVATE_TIME, 0.0, 1.0)
	# Opacity goes from 0.35 to 1.0 as gaze progresses
	var alpha := lerpf(0.35, 1.0, progress)
	# Color brightens toward white at the end
	var color := base_color.lerp(Color(1.0, 1.0, 1.0), progress * 0.5)
	color.a = alpha
	mat.albedo_color = color

	# Scale ball up when being gazed at
	var s := lerpf(GAZE_BALL_BASE_SCALE, GAZE_BALL_MAX_SCALE, progress)
	ball.scale = Vector3(s, s, s)


func _show_gaze_laser(from: Vector3, to: Vector3) -> void:
	if _gaze_laser == null:
		return
	var dir := to - from
	var length := dir.length()
	var cyl := _gaze_laser.mesh as CylinderMesh
	if cyl:
		cyl.height = length
	# Position at midpoint, orient Y-axis along the ray
	_gaze_laser.global_position = from + dir * 0.5
	var up := dir.normalized()
	var arbitrary := Vector3.UP if abs(up.dot(Vector3.UP)) < 0.99 else Vector3.RIGHT
	var right := up.cross(arbitrary).normalized()
	var forward := right.cross(up).normalized()
	_gaze_laser.global_transform.basis = Basis(right, up, forward)
	_gaze_laser.visible = true


func _exit_ar_mode() -> void:
	current_phase = Phase.CONFIG

	# Disable passthrough
	_disable_passthrough()

	# Hide gaze balls, laser, and webcam
	for ball in _gaze_balls:
		ball.visible = false
		ball.scale = Vector3.ONE
	for i in range(_gaze_timers.size()):
		_gaze_timers[i] = 0.0
	if _gaze_laser:
		_gaze_laser.visible = false
	webcam_quad.visible = false

	# Show start screen again, reposition in front of user
	var hmd := XRServer.get_hmd_transform()
	_position_panel_in_front_of_user(hmd)
	print("[Main] Returned to CONFIG phase")