extends SceneTree

const STAGE_SCENE := preload("res://scenes/stage/sandbox_stage.tscn")
const BONE_PATHS: Array[String] = [
	"Skeleton2D/Torso/Head",
	"Skeleton2D/Torso/LeftUpperArm",
	"Skeleton2D/Torso/LeftUpperArm/LeftLowerArm",
	"Skeleton2D/Torso/RightUpperArm",
	"Skeleton2D/Torso/RightUpperArm/RightLowerArm",
]
const MARKER_NAMES: Array[String] = ["Head", "Left_Hand", "Right_Hand", "Torso"]
const ROTATION_TOLERANCE := 0.0001
const POSITION_TOLERANCE := 0.01

var _checks := 0
var _failures := 0


func _initialize() -> void:
	call_deferred("_run")


func _run() -> void:
	var stage := STAGE_SCENE.instantiate() as Node2D
	root.add_child(stage)
	stage._snap_enabled = false
	var rig: StickmanRig = stage._spawner.spawn("stickman", Vector2.ZERO)
	stage._finalize_placed_node(rig)
	for frame in 3:
		await process_frame
	var pose_before := _capture_pose(rig)
	var markers_before := _capture_markers(rig)
	stage.set_mode(stage.StageMode.DIRECT)
	for frame in 3:
		await process_frame
	var pose_after := _capture_pose(rig)
	var markers_after := _capture_markers(rig)
	_check(stage.current_mode == stage.StageMode.DIRECT, "stage enters DIRECT")
	for path: String in BONE_PATHS:
		var before: Dictionary = pose_before[path]
		var after: Dictionary = pose_after[path]
		_check(absf(angle_difference(float(before["rotation"]), float(after["rotation"]))) <= ROTATION_TOLERANCE,
			"EDIT to DIRECT preserves %s rotation" % path.get_file())
		_check((before["position"] as Vector2).distance_to(after["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to DIRECT preserves %s position" % path.get_file())
	for marker_name: String in MARKER_NAMES:
		var before: Dictionary = markers_before[marker_name]
		var after: Dictionary = markers_after[marker_name]
		_check((before["position"] as Vector2).distance_to(after["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to DIRECT preserves %s IK marker" % marker_name)
	stage.set_mode(stage.StageMode.EDIT)
	for frame in 3:
		await process_frame
	var pose_round_trip := _capture_pose(rig)
	var markers_round_trip := _capture_markers(rig)
	_check(stage.current_mode == stage.StageMode.EDIT, "stage returns to EDIT")
	for path: String in BONE_PATHS:
		var before: Dictionary = pose_before[path]
		var round_trip: Dictionary = pose_round_trip[path]
		_check(absf(angle_difference(float(before["rotation"]), float(round_trip["rotation"]))) <= ROTATION_TOLERANCE,
			"EDIT to DIRECT to EDIT preserves %s rotation" % path.get_file())
		_check((before["position"] as Vector2).distance_to(round_trip["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to DIRECT to EDIT preserves %s position" % path.get_file())
	for marker_name: String in MARKER_NAMES:
		var before: Dictionary = markers_before[marker_name]
		var round_trip: Dictionary = markers_round_trip[marker_name]
		_check((before["position"] as Vector2).distance_to(round_trip["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to DIRECT to EDIT preserves %s IK marker" % marker_name)
	stage.set_mode(stage.StageMode.PLAY)
	for frame in 3:
		await physics_frame
	stage.set_mode(stage.StageMode.DIRECT)
	for frame in 4:
		await physics_frame
	var pose_after_play := _capture_pose(rig)
	var markers_after_play := _capture_markers(rig)
	_check(stage.current_mode == stage.StageMode.DIRECT, "stage enters DIRECT after PLAY")
	for path: String in BONE_PATHS:
		var before: Dictionary = pose_before[path]
		var after_play: Dictionary = pose_after_play[path]
		_check(absf(angle_difference(float(before["rotation"]), float(after_play["rotation"]))) <= ROTATION_TOLERANCE,
			"EDIT to PLAY to DIRECT restores %s rotation" % path.get_file())
		_check((before["position"] as Vector2).distance_to(after_play["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to PLAY to DIRECT restores %s position" % path.get_file())
	for marker_name: String in MARKER_NAMES:
		var before: Dictionary = markers_before[marker_name]
		var after_play: Dictionary = markers_after_play[marker_name]
		_check((before["position"] as Vector2).distance_to(after_play["position"] as Vector2) <= POSITION_TOLERANCE,
			"EDIT to PLAY to DIRECT restores %s IK marker" % marker_name)
	stage.queue_free()
	await process_frame
	if _failures == 0:
		print("RESULT: ALL PASSED (%d assertions)" % _checks)
		quit(0)
	else:
		print("RESULT: %d FAILURE(S) out of %d assertions" % [_failures, _checks])
		quit(1)


func _capture_pose(rig: StickmanRig) -> Dictionary:
	var result := {}
	for path: String in BONE_PATHS:
		var bone := rig.get_node(path) as Bone2D
		result[path] = {"rotation": bone.global_rotation, "position": bone.global_position}
	return result


func _capture_markers(rig: StickmanRig) -> Dictionary:
	var result := {}
	for marker_name: String in MARKER_NAMES:
		var marker := rig.get_node("IK_Targets/" + marker_name) as Marker2D
		result[marker_name] = {"position": marker.position, "rotation": marker.rotation}
	return result


func _check(condition: bool, label: String) -> void:
	_checks += 1
	if condition:
		print("PASS: %s" % label)
	else:
		_failures += 1
		push_error("FAIL: %s" % label)