@tool
extends SceneTree

# =============================================================================
# validate_polycount.gd — CSG Triangle Budget Checker
# =============================================================================
# Audits a Godot 4 map scene for CSG polygon / triangle budget compliance.
# Budget limits (competitive FPS):
#   - Total faces:      < 50,000
#   - Per node:          < 5,000
#   - Total CSG nodes:  < 200
#
# Usage:
#   godot --headless --script client/scripts/validate_polycount.gd -- res://path/to/map.tscn
#
# Exit codes:
#   0  — All poly budgets met
#   1  — One or more budgets exceeded
#   2  — Scene file not found or cannot be loaded
# =============================================================================

# Budget constants (adjust for target platform)
const MAX_TOTAL_FACES := 50_000
const MAX_FACES_PER_NODE := 5_000
const MAX_CSG_NODES := 200
const MAX_MESH_INSTANCES := 150
const WARN_TOTAL_FACES := 40_000
const WARN_FACES_PER_NODE := 3_000

var _pass_count := 0
var _fail_count := 0
var _warning_count := 0
var _scene_path := ""
var _scene_instance: Node = null

# Accumulated stats
var _total_faces := 0
var _total_vertices := 0
var _total_csg_nodes := 0
var _total_mesh_instances := 0
var _node_face_counts: Dictionary = {}  # node_path -> face_count
var _highest_face_node := ""
var _highest_face_count := 0
var _csg_node_list: Array[String] = []
var _overbudget_nodes: Array[String] = []


func _init() -> void:
	_parse_arguments()

	if _scene_path.is_empty():
		printerr("ERROR: No scene path provided. Usage: godot --headless --script validate_polycount.gd -- res://path/to/map.tscn")
		quit(2)
		return

	print("")
	print("==========================================")
	print("  CSG POLYGON BUDGET VALIDATOR")
	print("  Scene: ", _scene_path)
	print("==========================================")
	print("")

	# Budget header
	print("  Budget limits:")
	print("    Max total faces:        ", _format_num(MAX_TOTAL_FACES))
	print("    Max faces per node:     ", _format_num(MAX_FACES_PER_NODE))
	print("    Max CSG nodes:          ", _format_num(MAX_CSG_NODES))
	print("    Max MeshInstance3D:     ", _format_num(MAX_MESH_INSTANCES))
	print("")

	var scene: PackedScene = _load_scene(_scene_path)
	if scene == null:
		printerr("ERROR: Scene not found: ", _scene_path)
		quit(2)
		return

	_scene_instance = _instantiate_scene(scene)
	if _scene_instance == null:
		printerr("ERROR: Failed to instantiate scene: ", _scene_path)
		quit(2)
		return

	root.add_child(_scene_instance)

	# Perform analysis
	_analyze_node(_scene_instance, "")

	# Run all checks
	_run_checks()

	# Print detailed report
	_print_detailed_report()

	# Summary
	_print_summary()

	if _fail_count > 0:
		quit(1)
	else:
		quit(0)


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

func _parse_arguments() -> void:
	"""Extract scene path from user arguments (after --)."""
	var args: PackedStringArray = OS.get_cmdline_user_args()
	if not args.is_empty():
		_scene_path = args[0].strip_edges()
		return

	# Fallback: scan full cmdline for --
	var full_args: PackedStringArray = OS.get_cmdline_args()
	var found_sep := false
	for i in range(full_args.size()):
		if full_args[i] == "--":
			found_sep = true
			continue
		if found_sep:
			_scene_path = full_args[i].strip_edges()
			break


# ---------------------------------------------------------------------------
# Scene helpers
# ---------------------------------------------------------------------------

func _load_scene(path: String) -> PackedScene:
	"""Load a .tscn/.scn file."""
	if not ResourceLoader.exists(path):
		if not path.ends_with(".tscn") and not path.ends_with(".scn"):
			path = path + ".tscn"
		if not ResourceLoader.exists(path):
			return null
	return ResourceLoader.load(path, "PackedScene", ResourceLoader.CACHE_MODE_REUSE)


func _instantiate_scene(scene: PackedScene) -> Node:
	"""Instantiate a packed scene."""
	return scene.instantiate()


# ---------------------------------------------------------------------------
# Check helpers
# ---------------------------------------------------------------------------

func _pass(msg: String) -> void:
	_pass_count += 1
	print("  ✓ ", msg)


func _fail(msg: String) -> void:
	_fail_count += 1
	print("  ✗ ", msg)


func _warn(msg: String) -> void:
	_warning_count += 1
	print("  ⚠ ", msg)


# ---------------------------------------------------------------------------
# Polygon analysis engine
# ---------------------------------------------------------------------------

func _analyze_node(node: Node, path: String) -> void:
	"""
	Recursively walk the scene tree and collect polygon statistics
	for all CSG and MeshInstance nodes.
	"""
	var current_path := path + "/" + node.name if path != "" else node.name

	# Handle CSG nodes
	var class_name := node.get_class()
	var is_csg := class_name in [
		"CSGCombiner3D", "CSGSphere3D", "CSGBox3D", "CSGCylinder3D",
		"CSGTorus3D", "CSGPolygon3D", "CSGMesh3D"
	]

	if is_csg:
		_total_csg_nodes += 1
		_csg_node_list.append(current_path)

	# Try to get mesh / face count
	var face_count := _estimate_faces(node, class_name)
	var vert_count := _estimate_vertices(node, class_name)

	if face_count > 0:
		_total_faces += face_count
		_total_vertices += vert_count
		_node_face_counts[current_path] = face_count

		if face_count > _highest_face_count:
			_highest_face_count = face_count
			_highest_face_node = current_path

		if face_count > MAX_FACES_PER_NODE:
			_overbudget_nodes.append(current_path)

	# Count MeshInstance3D
	if class_name == "MeshInstance3D" or node.is_class("MeshInstance3D"):
		_total_mesh_instances += 1

	# Recurse children
	for child in node.get_children():
		_analyze_node(child, current_path)


func _estimate_faces(node: Node, class_name: String) -> int:
	"""
	Attempt to estimate the number of triangle faces for a given node.
	For CSG nodes, Godot exposes get_meshes() which returns arrays.
	For MeshInstance3D, we inspect the mesh resource.
	"""
	var faces := 0

	# CSG nodes can provide their collision / bake mesh
	if class_name in ["CSGCombiner3D", "CSGSphere3D", "CSGBox3D", "CSGCylinder3D",
		"CSGTorus3D", "CSGPolygon3D", "CSGMesh3D"]:

		# Try to get the bake mesh (the triangulated mesh used for lightmapping)
		if node.has_method("get_bake_mesh"):
			var bake_mesh := node.get_bake_mesh()
			if bake_mesh != null:
				faces = _count_mesh_faces(bake_mesh)

		# If that fails, try the collision/debug mesh
		if faces == 0 and node.has_method("get_meshes"):
			var mesh_arrays := node.get_meshes()
			if mesh_arrays is Array and mesh_arrays.size() > 1:
				var mesh := mesh_arrays[1] as ArrayMesh
				if mesh != null:
					faces = _count_mesh_faces(mesh)

		# Fallback estimate: for simple primitives
		if faces == 0:
			match class_name:
				"CSGBox3D":
					faces = 12  # 6 faces * 2 triangles
				"CSGSphere3D":
					faces = 320  # typical sphere LOD
				"CSGCylinder3D":
					faces = 128  # typical cylinder
				"CSGTorus3D":
					faces = 512  # typical torus
				"CSGPolygon3D":
					faces = 64   # default polygon
				"CSGMesh3D":
					faces = 100  # fallback
				"CSGCombiner3D":
					faces = 0   # combiner itself has no geometry

	# MeshInstance3D: count faces from mesh resource
	if class_name == "MeshInstance3D" or node.is_class("MeshInstance3D"):
		if "mesh" in node and node.mesh != null:
			faces = _count_mesh_faces(node.mesh)

	return faces


func _estimate_vertices(node: Node, class_name: String) -> int:
	"""Estimate vertices from the node's mesh if available."""
	var verts := 0

	if class_name in ["CSGCombiner3D", "CSGSphere3D", "CSGBox3D", "CSGCylinder3D",
		"CSGTorus3D", "CSGPolygon3D", "CSGMesh3D"]:
		if node.has_method("get_bake_mesh"):
			var bake_mesh := node.get_bake_mesh()
			if bake_mesh != null:
				verts = _count_mesh_vertices(bake_mesh)

	if verts == 0 and (class_name == "MeshInstance3D" or node.is_class("MeshInstance3D")):
		if "mesh" in node and node.mesh != null:
			verts = _count_mesh_vertices(node.mesh)

	return verts


func _count_mesh_faces(mesh) -> int:
	"""Count triangle faces in any Mesh type."""
	if mesh == null:
		return 0

	var count := 0

	# ArrayMesh and similar
	if mesh.has_method("get_surface_count"):
		var surface_count := mesh.get_surface_count()
		for s in range(surface_count):
			var arrays := mesh.surface_get_arrays(s)
			if arrays is Array and arrays.size() > 0:
				var indices := arrays[Mesh.ARRAY_INDEX] as PackedInt32Array
				if indices != null and indices.size() > 0:
					count += indices.size() / 3
				else:
					# Non-indexed: count vertices / 3
					var verts := arrays[Mesh.ARRAY_VERTEX] as PackedVector3Array
					if verts != null:
						count += verts.size() / 3
	return count


func _count_mesh_vertices(mesh) -> int:
	"""Count vertices in any Mesh type."""
	if mesh == null:
		return 0

	var count := 0
	if mesh.has_method("get_surface_count"):
		var surface_count := mesh.get_surface_count()
		for s in range(surface_count):
			var arrays := mesh.surface_get_arrays(s)
			if arrays is Array and arrays.size() > 0:
				var verts := arrays[Mesh.ARRAY_VERTEX] as PackedVector3Array
				if verts != null:
					count += verts.size()
	return count


# ---------------------------------------------------------------------------
# Budget checks
# ---------------------------------------------------------------------------

func _run_checks() -> void:
	"""Run all poly budget validation checks."""
	print("--- Budget Checks ---")

	# Total faces check
	if _total_faces < WARN_TOTAL_FACES:
		_pass("Total faces: " + _format_num(_total_faces) + " (budget: < " + _format_num(MAX_TOTAL_FACES) + ")")
	elif _total_faces < MAX_TOTAL_FACES:
		_warn("Total faces approaching limit: " + _format_num(_total_faces) + " (budget: < " + _format_num(MAX_TOTAL_FACES) + ")")
	else:
		_fail("Total faces EXCEEDED: " + _format_num(_total_faces) + " (budget: < " + _format_num(MAX_TOTAL_FACES) + ")")

	# Per-node face check
	if _overbudget_nodes.size() == 0:
		_pass("No nodes exceed " + _format_num(MAX_FACES_PER_NODE) + " faces")
	else:
		for node_path in _overbudget_nodes:
			_fail("Node exceeds face budget: " + node_path + " (" + _format_num(_node_face_counts[node_path]) + " faces, max: " + _format_num(MAX_FACES_PER_NODE) + ")")

	# CSG node count check
	if _total_csg_nodes < MAX_CSG_NODES:
		_pass("CSG node count: " + str(_total_csg_nodes) + " (budget: < " + str(MAX_CSG_NODES) + ")")
	else:
		_fail("CSG node count EXCEEDED: " + str(_total_csg_nodes) + " (budget: < " + str(MAX_CSG_NODES) + ")")

	# Mesh instance count check
	if _total_mesh_instances < MAX_MESH_INSTANCES:
		_pass("MeshInstance3D count: " + str(_total_mesh_instances) + " (budget: < " + str(MAX_MESH_INSTANCES) + ")")
	else:
		_fail("MeshInstance3D count EXCEEDED: " + str(_total_mesh_instances) + " (budget: < " + str(MAX_MESH_INSTANCES) + ")")

	# Highest face node info
	print("")
	print("--- Hot Spots ---")
	if _highest_face_count > 0:
		if _highest_face_count > WARN_FACES_PER_NODE:
			_warn("Highest face node: " + _highest_face_node + " (" + _format_num(_highest_face_count) + " faces)")
		else:
			_pass("Highest face node: " + _highest_face_node + " (" + _format_num(_highest_face_count) + " faces)")
	else:
		_warn("No face data collected (scene may use non-CSG geometry)")


# ---------------------------------------------------------------------------
# Detailed report
# ---------------------------------------------------------------------------

func _print_detailed_report() -> void:
	"""Print a detailed per-node breakdown if there are few enough nodes."""
	print("")
	print("--- Detailed Breakdown ---")
	print("  Total triangles:   ", _format_num(_total_faces))
	print("  Total vertices:    ", _format_num(_total_vertices))
	print("  CSG nodes:         ", _total_csg_nodes)
	print("  MeshInstance3D:    ", _total_mesh_instances)
	print("  Nodes over budget: ", _overbudget_nodes.size())
	print("")

	# Top 10 most expensive nodes
	var sorted_nodes: Array[Dictionary] = []
	for node_path in _node_face_counts.keys():
		sorted_nodes.append({"path": node_path, "faces": _node_face_counts[node_path]})
	sorted_nodes.sort_custom(_sort_by_faces_desc)

	var top_n := mini(sorted_nodes.size(), 10)
	if top_n > 0:
		print("  Top " + str(top_n) + " most expensive nodes:")
		for i in range(top_n):
			var entry := sorted_nodes[i] as Dictionary
			var bar := _make_bar(entry.faces, _highest_face_count)
			print("    " + str(i + 1) + ". " + entry.path + "  —  " + _format_num(entry.faces) + " faces " + bar)
		print("")


func _sort_by_faces_desc(a: Dictionary, b: Dictionary) -> bool:
	return a.faces > b.faces


func _make_bar(value: int, max_value: int) -> String:
	"""Create a simple ASCII bar visualization."""
	if max_value == 0:
		return ""
	var ratio := float(value) / float(max_value)
	var bar_len := int(ratio * 20)
	if bar_len < 1 and value > 0:
		bar_len = 1
	var bar := ""
	for i in range(bar_len):
		bar += "█"
	return bar


# ---------------------------------------------------------------------------
# Formatting helpers
# ---------------------------------------------------------------------------

func _format_num(n: int) -> String:
	"""Format a number with thousand separators."""
	var s := str(n)
	var result := ""
	var count := 0
	for i in range(s.length() - 1, -1, -1):
		if count > 0 and count % 3 == 0:
			result = "," + result
		result = s[i] + result
		count += 1
	return result


func mini(a: int, b: int) -> int:
	return a if a < b else b


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

func _print_summary() -> void:
	"""Print final summary."""
	print("==========================================")
	print("  POLYGON BUDGET SUMMARY")
	print("==========================================")
	print("  Scene:            ", _scene_path)
	print("  Total faces:      ", _format_num(_total_faces))
	print("  Total vertices:   ", _format_num(_total_vertices))
	print("  CSG nodes:        ", _total_csg_nodes)
	print("  MeshInstance3D:   ", _total_mesh_instances)
	print("  Over-budget:      ", _overbudget_nodes.size())
	print("------------------------------------------")
	print("  Passed:           ", _pass_count)
	print("  Failed:           ", _fail_count)
	print("  Warnings:         ", _warning_count)
	print("------------------------------------------")
	if _fail_count > 0:
		print("  RESULT:           FAILED  ❌")
		printerr("Validation FAILED — ", _fail_count, " budget(s) exceeded.")
	else:
		print("  RESULT:           PASSED  ✅")
	print("==========================================")
	print("")