extends SceneTree
## Gear-on-character capture: one frame per gear piece, worn by a real avatar.
##
## The gear library is authored in attach-point space, so the only way to judge a
## piece is to see it on the rig it ships with. Each shot stages `npc_casual`,
## attaches exactly one gear item through the runtime path (`_attach_gear`, which
## applies the same offsets the game uses) and frames the region that gear lives in.
##
## Usage (windowed):
##   godot --path . --script res://src/test/capture_npc_gear.gd -- \
##     --out=res://tmp/npc-gear
##
## Then compose the sheet:
##   python3 tools/npc_baseline.py sheet --dir tmp/npc-gear --out tmp/npc-gear/sheet.jpg

const Rules := preload("res://src/core/rules.gd")
const NpcAvatar := preload("res://src/game/npc_avatar.gd")
const NpcProfile := preload("res://src/core/npc_profile.gd")

## Gear order matches `NpcAvatar.GEAR_ATTACH`; the crop frames where the piece sits.
const SHOTS := [
	{"gear": "firefighter_helmet", "view": "head"},
	{"gear": "sunglasses", "view": "head"},
	{"gear": "lanyard", "view": "chest"},
	{"gear": "backpack", "view": "back"},
	{"gear": "shoulder_bag", "view": "back"},
	{"gear": "scba_tank", "view": "back"},
	{"gear": "camera", "view": "hand"},
	{"gear": "clipboard", "view": "hand"},
	{"gear": "phone", "view": "hand"},
	{"gear": "cup", "view": "hand"},
	{"gear": "umbrella", "view": "hand"},
	{"gear": "protest_sign", "view": "hand"},
	{"gear": "flashlight", "view": "hand"},
	{"gear": "wand", "view": "hand"},
	{"gear": "radio", "view": "hand"},
]
const VIEWS := {
	"head": {"yaw": 18.0, "distance": 0.95, "height": 1.42, "target_y": 1.72, "fov": 34.0},
	"chest": {"yaw": 18.0, "distance": 1.10, "height": 1.45, "target_y": 1.32, "fov": 34.0},
	"back": {"yaw": 158.0, "distance": 1.80, "height": 1.50, "target_y": 1.25, "fov": 36.0},
	"hand": {"yaw": 22.0, "distance": 1.25, "height": 1.35, "target_y": 0.95, "fov": 34.0},
}
const SETTLE_FRAMES := 3
const MAX_FRAMES := 3000

var out_dir := "user://screenshots/npc_gear"
var _stage: Node3D
var _camera: Camera3D
var _avatar: NpcAvatar
var _shot := 0
var _frames := 0


func _initialize() -> void:
	for arg in OS.get_cmdline_user_args():
		if str(arg).begins_with("--out="):
			out_dir = str(arg).substr(6)


func _process(_delta: float) -> bool:
	_frames += 1
	if _frames > MAX_FRAMES:
		push_error("capture_npc_gear timed out")
		quit(1)
		return true
	if _stage == null:
		_build_stage()
		_stage_shot()
		return false
	if _frames < SETTLE_FRAMES:
		return false
	_shoot()
	_shot += 1
	if _shot >= SHOTS.size():
		print("capture_npc_gear done -> %s" % out_dir)
		quit(0)
		return true
	_stage_shot()
	return false


func _build_stage() -> void:
	_stage = Node3D.new()
	_stage.name = "NpcGearStage"
	root.add_child(_stage)

	var env := Environment.new()
	env.background_mode = Environment.BG_COLOR
	env.background_color = Color(0.56, 0.59, 0.63)
	env.ambient_light_source = Environment.AMBIENT_SOURCE_COLOR
	env.ambient_light_color = Color(0.72, 0.75, 0.80)
	env.ambient_light_energy = 0.55
	env.tonemap_mode = Environment.TONE_MAPPER_ACES
	var world_env := WorldEnvironment.new()
	world_env.environment = env
	_stage.add_child(world_env)

	var sun := DirectionalLight3D.new()
	sun.rotation_degrees = Vector3(-38.0, -32.0, 0.0)
	sun.light_energy = 1.35
	sun.light_color = Color(1.0, 0.97, 0.92)
	sun.shadow_enabled = true
	_stage.add_child(sun)
	var fill := DirectionalLight3D.new()
	fill.rotation_degrees = Vector3(-12.0, 145.0, 0.0)
	fill.light_energy = 0.45
	fill.light_color = Color(0.75, 0.82, 1.0)
	_stage.add_child(fill)

	var ground := MeshInstance3D.new()
	var plane := PlaneMesh.new()
	plane.size = Vector2(24.0, 24.0)
	ground.mesh = plane
	var mat := StandardMaterial3D.new()
	mat.albedo_color = Color(0.30, 0.32, 0.35)
	mat.roughness = 0.92
	ground.material_override = mat
	_stage.add_child(ground)

	_camera = Camera3D.new()
	_camera.current = true
	_stage.add_child(_camera)


func _stage_shot() -> void:
	if _avatar != null and is_instance_valid(_avatar):
		_avatar.queue_free()
	_avatar = null
	var shot: Dictionary = SHOTS[_shot]
	var profile := NpcProfile.new()
	profile.id = "gear_" + str(shot["gear"])
	profile.display_name = str(shot["gear"]).to_upper()
	profile.npc_type = Rules.NpcType.TOURIST
	var avatar := NpcAvatar.new()
	avatar.name = "GearSubject"
	# The hoodie base has short sleeves and no gear of its own beyond the outfit's,
	# so the attach point of the tested piece is unobstructed.
	avatar.character_override = "npc_hoodie"
	avatar.setup(profile, "roam", 4242, Vector3.ZERO, Vector3(0, 0, 2),
		Vector3(0, 0.5, 1.2), str(profile.id))
	_stage.add_child(avatar)
	avatar.global_position = Vector3.ZERO
	avatar.rotation.y = 0.0
	avatar.set_physics_process(false)
	avatar.set_process(false)
	# Clear the gear the outfit attached for itself, or the shot shows two items in
	# the same hand and the tested piece is unreadable.
	for node in avatar.find_children("Gear_*", "Node3D", true, false):
		(node as Node3D).queue_free()
	# Attach the tested piece through the real runtime path, so the offsets under
	# test are exactly the ones the game applies.
	var gear := str(shot["gear"])
	if not avatar._attach_gear(gear):
		push_warning("gear %s did not attach" % gear)
	avatar._animate(1.0 / 60.0)
	_avatar = avatar
	_frame_view(VIEWS[str(shot["view"])])


func _frame_view(view: Dictionary) -> void:
	var yaw := deg_to_rad(float(view["yaw"]))
	var distance := float(view["distance"])
	var target := Vector3(0.0, float(view["target_y"]), 0.0)
	_camera.position = target + Vector3(sin(yaw) * distance,
		float(view["height"]) - target.y, cos(yaw) * distance)
	_camera.fov = float(view["fov"])
	_camera.look_at(target, Vector3.UP)


func _shoot() -> void:
	var dir := ProjectSettings.globalize_path(out_dir)
	DirAccess.make_dir_recursive_absolute(dir)
	var path := dir.path_join("%02d-%s.png" % [_shot, SHOTS[_shot]["gear"]])
	var img: Image = root.get_texture().get_image()
	var err := img.save_png(path)
	print("%s -> err=%d" % [path.get_file(), err])