extends Node3D
## Reads XR_FB_body_tracking joints each frame via Godot's XRBodyTracker.
## Computes chest-relative wrist positions and emits tracking data.
## Visualizes all body joints as colored spheres.
##
## Meta body tracking provides 70 joints. We use:
##   CHEST (5)  - torso orientation (solves body rotation problem)
##   HEAD (7)   - head pose
##   LEFT_HAND_WRIST (12)  - left wrist
##   RIGHT_HAND_WRIST (17) - right wrist
##   Hand joints (18-69)   - finger positions
##
## Output poses are in Godot's coordinate system (Y-up, -Z forward).
## The server handles conversion to robot conventions.

## Emitted every frame with tracking data dict ready to send.
signal tracking_data_ready(data: Dictionary)

## Joint indices from XR_FB_body_tracking
## Reference: Meta OpenXR body tracking extension
const JOINT_ROOT := 0
const JOINT_HIPS := 1
const JOINT_SPINE_LOWER := 2
const JOINT_SPINE_MIDDLE := 3
const JOINT_SPINE_UPPER := 4
const JOINT_CHEST := 5
const JOINT_NECK := 6
const JOINT_HEAD := 7

const JOINT_LEFT_SHOULDER := 8
const JOINT_LEFT_SCAPULA := 9
const JOINT_LEFT_ARM_UPPER := 10
const JOINT_LEFT_ARM_LOWER := 11
const JOINT_LEFT_HAND_WRIST := 12

const JOINT_RIGHT_SHOULDER := 13
const JOINT_RIGHT_SCAPULA := 14
const JOINT_RIGHT_ARM_UPPER := 15
const JOINT_RIGHT_ARM_LOWER := 16
const JOINT_RIGHT_HAND_WRIST := 17

# Hand joints start at index 18 for left hand, 43 for right hand
# 25 joints per hand (same layout as XR_EXT_hand_tracking)
const JOINT_LEFT_HAND_START := 18
const JOINT_RIGHT_HAND_START := 43
const HAND_JOINT_COUNT := 25

## Total body joint count
const BODY_JOINT_COUNT := 70

## Joint visualization colors
const COLOR_BODY := Color(1.0, 0.7, 0.2, 1.0)       # Orange - spine/torso
const COLOR_HEAD := Color(1.0, 1.0, 1.0, 1.0)        # White
const COLOR_LEFT_ARM := Color(0.3, 0.5, 1.0, 1.0)    # Blue
const COLOR_RIGHT_ARM := Color(0.3, 1.0, 0.5, 1.0)   # Green
const COLOR_LEFT_HAND := Color(0.5, 0.7, 1.0, 1.0)   # Light blue
const COLOR_RIGHT_HAND := Color(0.5, 1.0, 0.7, 1.0)  # Light green

## Tracking state
var body_tracker_name: StringName = &"/user/body_tracker"
var is_tracking: bool = false
var frames_since_last_send: int = 0

## Send rate control: target ~30 Hz (every other frame at 72 Hz)
@export var send_every_n_frames: int = 2

## Debug logging
@export var debug_log: bool = false
var _log_counter: int = 0

## Joint visualization
var _xr_origin: XROrigin3D
var _joint_spheres: Array = []  # Array of MeshInstance3D, indexed by joint


func setup(xr_origin: XROrigin3D) -> void:
	_xr_origin = xr_origin
	_create_joint_spheres()


func _ready() -> void:
	print("[BodyTracker] Initialized, waiting for body tracking data...")


func _create_joint_spheres() -> void:
	var body_mesh := SphereMesh.new()
	body_mesh.radius = 0.025
	body_mesh.height = 0.05

	# Only create spheres for body joints (0-17), skip hand joints (18-69)
	# Hand joints are already visualized by vr_ui_pointer.gd
	_joint_spheres.resize(BODY_JOINT_COUNT)
	for i in range(BODY_JOINT_COUNT):
		if i >= JOINT_LEFT_HAND_START:
			_joint_spheres[i] = null
			continue
		var s := MeshInstance3D.new()
		s.mesh = body_mesh
		var mat := StandardMaterial3D.new()
		mat.shading_mode = BaseMaterial3D.SHADING_MODE_UNSHADED
		mat.albedo_color = _get_joint_color(i)
		s.material_override = mat
		s.visible = false
		add_child(s)
		_joint_spheres[i] = s


func _get_joint_color(joint_idx: int) -> Color:
	if joint_idx == JOINT_HEAD:
		return COLOR_HEAD
	elif joint_idx <= JOINT_NECK:
		return COLOR_BODY
	elif joint_idx >= JOINT_LEFT_SHOULDER and joint_idx <= JOINT_LEFT_HAND_WRIST:
		return COLOR_LEFT_ARM
	elif joint_idx >= JOINT_RIGHT_SHOULDER and joint_idx <= JOINT_RIGHT_HAND_WRIST:
		return COLOR_RIGHT_ARM
	elif joint_idx >= JOINT_LEFT_HAND_START and joint_idx < JOINT_RIGHT_HAND_START:
		return COLOR_LEFT_HAND
	elif joint_idx >= JOINT_RIGHT_HAND_START:
		return COLOR_RIGHT_HAND
	return COLOR_BODY


func _process(_delta: float) -> void:
	frames_since_last_send += 1
	if frames_since_last_send < send_every_n_frames:
		return
	frames_since_last_send = 0

	var tracker := XRServer.get_tracker(body_tracker_name) as XRBodyTracker
	if tracker == null:
		if is_tracking:
			print("[BodyTracker] Lost body tracking")
			is_tracking = false
			_hide_all_spheres()
		return

	if not tracker.get_has_tracking_data():
		if is_tracking:
			print("[BodyTracker] Body tracking data unavailable")
			is_tracking = false
			_hide_all_spheres()
		return

	if not is_tracking:
		print("[BodyTracker] Body tracking active!")
		is_tracking = true

	# Update joint sphere positions
	_update_joint_spheres(tracker)

	# Read key joint poses
	var chest_xform := tracker.get_joint_transform(JOINT_CHEST)
	var head_xform := tracker.get_joint_transform(JOINT_HEAD)
	var left_wrist_xform := tracker.get_joint_transform(JOINT_LEFT_HAND_WRIST)
	var right_wrist_xform := tracker.get_joint_transform(JOINT_RIGHT_HAND_WRIST)

	# Compute chest-relative wrist positions
	# This is the key advantage: eliminates head-position subtraction artifacts
	var chest_inv := chest_xform.affine_inverse()
	var left_wrist_rel := chest_inv * left_wrist_xform
	var right_wrist_rel := chest_inv * right_wrist_xform

	# Build tracking data packet
	var data := {}
	data["type"] = "tracking"
	data["timestamp"] = Time.get_ticks_msec()

	# Head pose (world space) - 7 floats: pos(3) + quat(4)
	data["head"] = _xform_to_pose7(head_xform)

	# Chest pose (world space) - 7 floats: pos(3) + quat(4)
	data["chest"] = _xform_to_pose7(chest_xform)

	# Chest-relative wrist poses as 4x4 matrices (16 floats each, column-major)
	# These replace the world-space wrist poses; no head subtraction needed
	data["left_wrist"] = _xform_to_mat16(left_wrist_rel)
	data["right_wrist"] = _xform_to_mat16(right_wrist_rel)

	# Hand joint positions (25 joints per hand, relative to wrist)
	var left_hand_pos := _get_hand_positions(tracker, JOINT_LEFT_HAND_START, left_wrist_xform)
	var right_hand_pos := _get_hand_positions(tracker, JOINT_RIGHT_HAND_START, right_wrist_xform)
	data["left_hand_pos"] = left_hand_pos
	data["right_hand_pos"] = right_hand_pos

	# Hand joint rotations (25 joints per hand, relative to wrist)
	var left_hand_rot := _get_hand_rotations(tracker, JOINT_LEFT_HAND_START, left_wrist_xform)
	var right_hand_rot := _get_hand_rotations(tracker, JOINT_RIGHT_HAND_START, right_wrist_xform)
	data["left_hand_rot"] = left_hand_rot
	data["right_hand_rot"] = right_hand_rot

	# Body joint transforms for retargeting (world-space pose7)
	# Order: hips, chest, L_upper_arm, L_lower_arm, L_wrist, R_upper_arm, R_lower_arm, R_wrist
	# Each joint: [x, y, z, qx, qy, qz, qw] = 7 floats, total 56 floats
	var hips_xform := tracker.get_joint_transform(JOINT_HIPS)
	var left_upper_arm_xform := tracker.get_joint_transform(JOINT_LEFT_ARM_UPPER)
	var left_lower_arm_xform := tracker.get_joint_transform(JOINT_LEFT_ARM_LOWER)
	var right_upper_arm_xform := tracker.get_joint_transform(JOINT_RIGHT_ARM_UPPER)
	var right_lower_arm_xform := tracker.get_joint_transform(JOINT_RIGHT_ARM_LOWER)

	# Use XRHandTracker wrist transforms instead of body tracker wrist —
	# hand tracking has much better wrist rotation detection
	var left_wrist_body := left_wrist_xform
	var right_wrist_body := right_wrist_xform
	var left_ht = XRServer.get_tracker(&"/user/hand_tracker/left") as XRHandTracker
	if left_ht and left_ht.has_tracking_data:
		var ht_wrist := left_ht.get_hand_joint_transform(XRHandTracker.HAND_JOINT_WRIST)
		if ht_wrist.origin != Vector3.ZERO:
			left_wrist_body = ht_wrist
	var right_ht = XRServer.get_tracker(&"/user/hand_tracker/right") as XRHandTracker
	if right_ht and right_ht.has_tracking_data:
		var ht_wrist := right_ht.get_hand_joint_transform(XRHandTracker.HAND_JOINT_WRIST)
		if ht_wrist.origin != Vector3.ZERO:
			right_wrist_body = ht_wrist

	var body_joints := []
	body_joints.append_array(_xform_to_pose7(hips_xform))
	body_joints.append_array(_xform_to_pose7(chest_xform))
	body_joints.append_array(_xform_to_pose7(left_upper_arm_xform))
	body_joints.append_array(_xform_to_pose7(left_lower_arm_xform))
	body_joints.append_array(_xform_to_pose7(left_wrist_body))
	body_joints.append_array(_xform_to_pose7(right_upper_arm_xform))
	body_joints.append_array(_xform_to_pose7(right_lower_arm_xform))
	body_joints.append_array(_xform_to_pose7(right_wrist_body))
	data["body_joints"] = body_joints

	# Debug logging every ~2 seconds
	if debug_log:
		_log_counter += 1
		if _log_counter >= 60:
			_log_counter = 0
			print("[BodyTracker] chest=", chest_xform.origin,
				  " head=", head_xform.origin,
				  " L_wrist_rel=", left_wrist_rel.origin,
				  " R_wrist_rel=", right_wrist_rel.origin)

	tracking_data_ready.emit(data)


## Convert Transform3D to 7-float pose: [x, y, z, qx, qy, qz, qw]
func _xform_to_pose7(xform: Transform3D) -> Array:
	var pos := xform.origin
	var quat := xform.basis.get_rotation_quaternion()
	return [pos.x, pos.y, pos.z, quat.x, quat.y, quat.z, quat.w]


## Convert Transform3D to 16-float column-major 4x4 matrix
## Layout: [m00, m10, m20, 0, m01, m11, m21, 0, m02, m12, m22, 0, tx, ty, tz, 1]
## This matches NumPy's Fortran ('F') order used by the existing teleop code.
func _xform_to_mat16(xform: Transform3D) -> Array:
	var b := xform.basis
	var o := xform.origin
	# Godot Basis: b[col][row] — but we store column-major
	# Column 0: basis x-axis
	# Column 1: basis y-axis
	# Column 2: basis z-axis
	# Column 3: origin (translation)
	return [
		b.x.x, b.x.y, b.x.z, 0.0,  # Column 0
		b.y.x, b.y.y, b.y.z, 0.0,  # Column 1
		b.z.x, b.z.y, b.z.z, 0.0,  # Column 2
		o.x,   o.y,   o.z,   1.0   # Column 3
	]


## Get 25 hand joint positions relative to wrist, as flat array (75 floats)
## Each joint: [x, y, z] relative to wrist
func _get_hand_positions(tracker: XRBodyTracker, start_idx: int, wrist_xform: Transform3D) -> Array:
	var wrist_inv := wrist_xform.affine_inverse()
	var positions := []
	positions.resize(HAND_JOINT_COUNT * 3)
	for i in range(HAND_JOINT_COUNT):
		var joint_xform := tracker.get_joint_transform(start_idx + i)
		var rel := wrist_inv * joint_xform
		positions[i * 3 + 0] = rel.origin.x
		positions[i * 3 + 1] = rel.origin.y
		positions[i * 3 + 2] = rel.origin.z
	return positions


func _update_joint_spheres(tracker: XRBodyTracker) -> void:
	if _joint_spheres.is_empty() or _xr_origin == null:
		return
	var origin_xform := _xr_origin.global_transform
	for i in range(BODY_JOINT_COUNT):
		if _joint_spheres[i] == null:
			continue
		var xform := tracker.get_joint_transform(i)
		if xform.origin == Vector3.ZERO:
			_joint_spheres[i].visible = false
			continue
		_joint_spheres[i].global_position = origin_xform * xform.origin
		_joint_spheres[i].visible = true


func _hide_all_spheres() -> void:
	for s in _joint_spheres:
		if s != null:
			s.visible = false


## Get 25 hand joint rotations relative to wrist, as flat array (225 floats)
## Each joint: 9 floats (3x3 rotation matrix, column-major)
func _get_hand_rotations(tracker: XRBodyTracker, start_idx: int, wrist_xform: Transform3D) -> Array:
	var wrist_rot_inv := wrist_xform.basis.inverse()
	var rotations := []
	rotations.resize(HAND_JOINT_COUNT * 9)
	for i in range(HAND_JOINT_COUNT):
		var joint_xform := tracker.get_joint_transform(start_idx + i)
		var rel_basis := wrist_rot_inv * joint_xform.basis
		# Store as column-major 3x3
		rotations[i * 9 + 0] = rel_basis.x.x
		rotations[i * 9 + 1] = rel_basis.x.y
		rotations[i * 9 + 2] = rel_basis.x.z
		rotations[i * 9 + 3] = rel_basis.y.x
		rotations[i * 9 + 4] = rel_basis.y.y
		rotations[i * 9 + 5] = rel_basis.y.z
		rotations[i * 9 + 6] = rel_basis.z.x
		rotations[i * 9 + 7] = rel_basis.z.y
		rotations[i * 9 + 8] = rel_basis.z.z
	return rotations