extends Node3D
## Generic simulator: Blender designs, Godot (Jolt) does the physics.
##
##   godot --headless --path godot --fixed-fps 240 -- <config.json>
##
## The .glb exported by Blender is loaded at runtime and every object becomes
## a body according to its name prefix:
##   static_*  static, exact mesh (ConcavePolygonShape3D)
##   box_*     static, box the size of the mesh
##   gate_*    static, box that a "remove" event takes out of the world
##   ball_*    rigid, sphere
##   block_*   rigid, box
## Trajectories are recorded in Blender coordinates (Z up), so the render
## script uses them without converting anything.
##
## The prefixes also fix the order bodies are created in: Blender exports
## objects alphabetically, and in a chaotic scene (hundreds of balls on pegs)
## a different creation order gives different trajectories. Keep new prefixes
## in the same alphabetical order (ball_ < box_ < gate_ < static_) or cached
## simulations will no longer match.

var hz := 240                       # physics steps per second
var steps_per_frame := 8            # hz / 30 video fps

var cfg: Dictionary
var dynamic_bodies: Array[RigidBody3D] = []
var removable := {}                 # name -> CollisionObject3D
var tick := 0
var frames: Array = []
var events: Array = []
var scale_f := 1.0                  # the world is simulated scale_f times larger


func _ready() -> void:
	var args := OS.get_cmdline_user_args()
	if args.is_empty():
		push_error("missing config.json")
		get_tree().quit(1)
		return
	cfg = JSON.parse_string(FileAccess.get_file_as_string(args[0]))
	# Jolt has tolerances in meters (2 cm of allowed penetration): with small
	# pieces it is better to simulate at a larger scale. Gravity and positions
	# are scaled alike, so the dynamics are the same and the output goes back
	# to the original scale.
	scale_f = float(cfg.get("scale", 1.0))
	# more steps per second = stiffer contacts (tall stacks); the process runs
	# with --fixed-fps 240, so each frame does hz/240 steps
	hz = int(cfg.get("hz", 240))
	# recording fps: 30 for real time, more for slow motion
	steps_per_frame = hz / int(cfg.get("fps", 30))
	Engine.physics_ticks_per_second = hz
	Engine.max_physics_steps_per_frame = maxi(1, hz / 240)
	PhysicsServer3D.area_set_param(get_world_3d().space, PhysicsServer3D.AREA_PARAM_GRAVITY,
			float(cfg.get("gravity", 9.81)) * scale_f)
	events = cfg.get("events", [])
	events.sort_custom(func(a, b): return a["t"] < b["t"])

	var doc := GLTFDocument.new()
	var st := GLTFState.new()
	if doc.append_from_file(cfg["glb"], st) != OK:
		push_error("could not read " + cfg["glb"])
		get_tree().quit(1)
		return
	var root := doc.generate_scene(st)
	add_child(root)
	var meshes: Array[MeshInstance3D] = []
	_collect(root, meshes)
	for mi in meshes:
		_create_body(mi)
	root.queue_free()
	print("[sim] %d dynamic, %d meshes, g=%.2f" % [dynamic_bodies.size(), meshes.size(),
			float(cfg.get("gravity", 9.81))])
	_record()


func _collect(n: Node, out: Array[MeshInstance3D]) -> void:
	if n is MeshInstance3D:
		out.append(n)
	for c in n.get_children():
		_collect(c, out)


func _rule(body_name: String) -> Dictionary:
	var r: Dictionary = {}
	var longest := -1
	for prefix in cfg.get("rules", {}):
		if body_name.begins_with(prefix) and prefix.length() > longest:
			r = cfg["rules"][prefix]
			longest = prefix.length()
	return r


func _create_body(mi: MeshInstance3D) -> void:
	var body_name := String(mi.name)
	var aabb := mi.get_aabb()
	aabb = AABB(aabb.position * scale_f, aabb.size * scale_f)
	var r := _rule(body_name)
	var mat := PhysicsMaterial.new()
	mat.friction = float(r.get("friction", 0.5))
	mat.bounce = float(r.get("bounce", 0.0))
	var body: CollisionObject3D
	var shape := CollisionShape3D.new()

	if body_name.begins_with("static_"):
		body = StaticBody3D.new()
		var tri := ConcavePolygonShape3D.new()
		var faces := mi.mesh.get_faces()
		for n in faces.size():
			faces[n] *= scale_f
		tri.set_faces(faces)
		shape.shape = tri
	elif body_name.begins_with("box_") or body_name.begins_with("gate_"):
		body = StaticBody3D.new()
		var b := BoxShape3D.new()
		b.size = aabb.size
		shape.shape = b
		shape.position = aabb.get_center()
	elif body_name.begins_with("ball_") or body_name.begins_with("block_"):
		var rb := RigidBody3D.new()
		var vol: float
		if body_name.begins_with("ball_"):
			var s := SphereShape3D.new()
			s.radius = aabb.size.x * 0.5
			shape.shape = s
			vol = 4.0 / 3.0 * PI * pow(s.radius, 3)
		else:
			var b := BoxShape3D.new()
			b.size = aabb.size
			shape.shape = b
			vol = aabb.size.x * aabb.size.y * aabb.size.z
		shape.position = aabb.get_center()
		rb.mass = vol * float(r.get("density", 1000.0))
		rb.continuous_cd = bool(r.get("ccd", false))
		rb.can_sleep = bool(r.get("can_sleep", true))
		rb.linear_damp = float(r.get("damping", 0.0))
		rb.angular_damp = float(r.get("angular_damping", 0.0))
		if r.get("planar", false):
			# motion in the board plane (Blender XZ = Godot XY)
			rb.axis_lock_linear_z = true
			rb.axis_lock_angular_x = true
			rb.axis_lock_angular_y = true
		for prefix in cfg.get("frozen", []):
			if body_name.begins_with(prefix):
				rb.freeze = true
				rb.freeze_mode = RigidBody3D.FREEZE_MODE_KINEMATIC
		body = rb
		dynamic_bodies.append(rb)
	else:
		return

	body.name = body_name
	if body is StaticBody3D:
		body.physics_material_override = mat
	else:
		(body as RigidBody3D).physics_material_override = mat
	body.add_child(shape)
	add_child(body)
	var tr := mi.global_transform
	tr.origin *= scale_f
	body.global_transform = tr
	removable[body_name] = body


func _physics_process(_delta: float) -> void:
	tick += 1
	var t := tick / float(hz)
	while not events.is_empty() and events[0]["t"] <= t:
		_apply(events.pop_front())
	if tick % steps_per_frame == 0:
		_record()
	if t >= float(cfg["duration"]):
		_save()
		get_tree().quit()


func _apply(ev: Dictionary) -> void:
	var prefix: String = ev["prefix"]
	var n := 0
	for body_name in removable.keys():
		if not String(body_name).begins_with(prefix):
			continue
		var c: CollisionObject3D = removable[body_name]
		if not is_instance_valid(c):
			continue
		n += 1
		match ev["action"]:
			"remove":
				c.queue_free()
			"release":
				var rb := c as RigidBody3D
				rb.freeze = false
				rb.sleeping = false
				if ev.has("velocity"):
					var v: Array = ev["velocity"]    # in Blender coordinates
					rb.linear_velocity = Vector3(v[0], v[2], -v[1]) * scale_f
				if ev.has("spin"):
					var w: Array = ev["spin"]        # rad/s, Blender axes
					rb.angular_velocity = Vector3(w[0], w[2], -w[1])
	# a sleeping body never notices its support was removed: wake them all
	for rb in dynamic_bodies:
		if not rb.freeze:
			rb.sleeping = false
	print("[sim] t=%.2f %s %s (%d)" % [tick / float(hz), ev["action"], prefix, n])


func _record() -> void:
	var row := PackedFloat32Array()
	for rb in dynamic_bodies:
		var p := rb.global_position / scale_f
		var q := rb.global_transform.basis.get_rotation_quaternion()
		# Godot (Y up) -> Blender (Z up): (x, y, z) -> (x, -z, y)
		row.append_array([p.x, -p.z, p.y, q.w, q.x, -q.z, q.y])
	frames.append(row)


func _save() -> void:
	# binary: a JSON header in a separate file + raw float32
	var out: String = cfg["output"]
	var names := []
	for rb in dynamic_bodies:
		names.append(String(rb.name))
	var f := FileAccess.open(out + ".header.json", FileAccess.WRITE)
	f.store_string(JSON.stringify({"names": names, "frames": frames.size(),
			"fps": int(cfg.get("fps", 30)), "fields": ["x", "y", "z", "qw", "qx", "qy", "qz"]}))
	f.close()
	var b := FileAccess.open(out + ".bin", FileAccess.WRITE)
	for row in frames:
		b.store_buffer(row.to_byte_array())
	b.close()
	print("[sim] %d frames -> %s.bin" % [frames.size(), out])