extends SceneTree

## Builds ADT terrain resources using one Texture2DArray/control-map material per tile.
##
## Usage:
##   godot --headless --path <project> --script res://src/tools/build_adt_control_splat_cache.gd -- \
##     --map Azeroth --jobs auto --texture-size 256 --lod 0 --force

const ADT_BUILDER_SCRIPT := preload("res://addons/mpq_extractor/loaders/adt_builder.gd")
const CONTROL_SPLAT_TILE_SCRIPT := preload("res://src/resources/control_splat_adt_tile.gd")

var _builder
var _loader
var _image_cache: Dictionary = {}


func _initialize() -> void:
	var args := OS.get_cmdline_user_args()
	var map_name := _get_arg_value(args, "--map", "Azeroth")
	var extracted_dir := _normalize_res_path(_get_arg_value(args, "--extracted", "res://data/extracted"))
	var output_dir := _normalize_res_path(_get_arg_value(args, "--output", "res://data/cache/terrain_control_splat_v3"))
	var lod := clampi(_get_arg_value(args, "--lod", "0").to_int(), 0, 3)
	var texture_size := clampi(_get_arg_value(args, "--texture-size", "256").to_int(), 16, 1024)
	var worker_index: Variant = _get_optional_int_arg(args, "--worker-index")
	var worker_count := maxi(1, _get_arg_value(args, "--worker-count", "1").to_int())
	var tile_x: Variant = _get_optional_int_arg(args, "--tile-x")
	var tile_y: Variant = _get_optional_int_arg(args, "--tile-y")
	var default_jobs := "1" if tile_x != null or tile_y != null else "auto"
	var jobs := _parse_jobs_arg(_get_arg_value(args, "--jobs", default_jobs))
	var force := args.has("--force")

	if jobs > 1 and worker_index == null:
		_run_parallel_controller(args, jobs)
		return

	_builder = ADT_BUILDER_SCRIPT.new()
	if not ClassDB.class_exists("ADTLoader"):
		push_error("ADTLoader not found. Rebuild GDExtension first.")
		quit(1)
		return

	_loader = ClassDB.instantiate("ADTLoader")
	if _loader == null:
		push_error("Failed to instantiate ADTLoader.")
		quit(1)
		return

	var map_dir := extracted_dir.path_join("World/Maps/%s" % map_name)
	var dir := DirAccess.open(map_dir)
	if dir == null:
		push_error("Cannot open map directory: %s" % map_dir)
		quit(1)
		return

	DirAccess.make_dir_recursive_absolute(ProjectSettings.globalize_path(output_dir.path_join(map_name)))

	var built := 0
	var skipped := 0
	var failed := 0
	var tile_index := 0
	var started_ms := Time.get_ticks_msec()
	var files := dir.get_files()
	var total := _count_worker_tiles(files, map_name, tile_x, tile_y, worker_index, worker_count)
	print("Control splat cache started. map=%s tiles=%d lod=%d texture_size=%d" % [
		map_name, total, lod, texture_size])

	for file_name in files:
		if not file_name.ends_with(".adt"):
			continue

		var stem := file_name.trim_suffix(".adt")
		var parts := stem.split("_")
		if parts.size() != 3 or parts[0] != map_name:
			continue
		if not parts[1].is_valid_int() or not parts[2].is_valid_int():
			continue

		var tx := parts[1].to_int()
		var ty := parts[2].to_int()
		if tile_x != null and tx != tile_x:
			continue
		if tile_y != null and ty != tile_y:
			continue
		if worker_index != null:
			if tile_index % worker_count != int(worker_index):
				tile_index += 1
				continue
		tile_index += 1

		var output_res_path := output_dir.path_join(map_name).path_join("%s_%d_%d.res" % [map_name, tx, ty])
		var output_abs_path := ProjectSettings.globalize_path(output_res_path)
		if not force and FileAccess.file_exists(output_abs_path):
			skipped += 1
			continue

		var source_res_path := map_dir.path_join(file_name)
		var source_abs_path := ProjectSettings.globalize_path(source_res_path)
		var data: Dictionary = _loader.call("load_adt", source_abs_path)
		if data.is_empty() or not data.has("chunks"):
			push_warning("Control splat failed to parse ADT: %s" % source_res_path)
			failed += 1
			continue

		var payload: Dictionary = _builder.build_control_splat_tile_render_payload(
			data,
			_image_cache,
			ProjectSettings.globalize_path(extracted_dir),
			lod,
			texture_size)
		var mesh: Mesh = payload.get("mesh", null)
		var texture_images: Array = payload.get("texture_images", [])
		if mesh == null or texture_images.is_empty():
			push_warning("Control splat produced incomplete tile: %s" % source_res_path)
			failed += 1
			continue
		var array_mesh := mesh as ArrayMesh
		if array_mesh != null and array_mesh.get_surface_count() > 0:
			array_mesh.surface_set_material(0, null)

		var tile: Resource = CONTROL_SPLAT_TILE_SCRIPT.new()
		tile.set("format_version", CONTROL_SPLAT_TILE_SCRIPT.FORMAT_VERSION)
		tile.set("map_name", map_name)
		tile.set("tile_x", tx)
		tile.set("tile_y", ty)
		tile.set("tile_origin", _builder.get_tile_origin_for_data(data))
		tile.set("lod", lod)
		tile.set("texture_size", texture_size)
		tile.set("terrain_mesh", mesh)
		tile.set("texture_images", texture_images)
		tile.set("alpha_atlas", payload.get("alpha_atlas", null))
		tile.set("layer_index_map", payload.get("layer_index_map", null))
		tile.set("texture_names", payload.get("texture_names", PackedStringArray()))

		var save_err := ResourceSaver.save(tile, output_res_path)
		if save_err != OK:
			push_warning("Failed to save control splat tile: %s (err=%d)" % [output_res_path, save_err])
			failed += 1
			continue

		built += 1
		if built % 25 == 0:
			print("Built %d control splat tiles..." % built)

	var elapsed_ms := Time.get_ticks_msec() - started_ms
	print("Control splat cache finished. built=%d skipped=%d failed=%d elapsed=%.2fs" % [
		built, skipped, failed, float(elapsed_ms) / 1000.0])
	quit(0 if failed == 0 else 2)


func _run_parallel_controller(args: PackedStringArray, jobs: int) -> void:
	var started_ms := Time.get_ticks_msec()
	var threads: Array[Thread] = []
	print("Starting parallel ADT control splat cache: jobs=%d" % jobs)
	for worker_index in jobs:
		var thread := Thread.new()
		var err := thread.start(_run_worker_process.bind(args, worker_index, jobs))
		if err != OK:
			push_error("Failed to start control splat worker thread %d (err=%d)" % [worker_index, err])
			quit(1)
			return
		threads.append(thread)

	var failed := 0
	for thread in threads:
		var exit_code: int = int(thread.wait_to_finish())
		if exit_code != 0:
			failed += 1

	var elapsed_ms := Time.get_ticks_msec() - started_ms
	print("Parallel control splat cache finished. jobs=%d failed_workers=%d elapsed=%.2fs" % [
		jobs, failed, float(elapsed_ms) / 1000.0])
	quit(0 if failed == 0 else 2)


func _run_worker_process(args: PackedStringArray, worker_index: int, worker_count: int) -> int:
	var worker_args := PackedStringArray([
		"--headless",
		"--path",
		ProjectSettings.globalize_path("res://"),
		"--script",
		"res://src/tools/build_adt_control_splat_cache.gd",
		"--",
	])

	var skip_next := false
	for i in args.size():
		if skip_next:
			skip_next = false
			continue
		var arg := args[i]
		if arg == "--jobs" or arg == "--worker-index" or arg == "--worker-count":
			skip_next = i + 1 < args.size()
			continue
		worker_args.append(arg)

	worker_args.append("--jobs")
	worker_args.append("1")
	worker_args.append("--worker-index")
	worker_args.append(str(worker_index))
	worker_args.append("--worker-count")
	worker_args.append(str(worker_count))

	print("Control splat worker %d/%d started" % [worker_index + 1, worker_count])
	var exit_code := OS.execute(OS.get_executable_path(), worker_args)
	print("Control splat worker %d/%d finished with exit=%d" % [worker_index + 1, worker_count, exit_code])
	return exit_code


func _count_worker_tiles(
		files: PackedStringArray,
		map_name: String,
		tile_x: Variant,
		tile_y: Variant,
		worker_index: Variant,
		worker_count: int) -> int:

	var tile_index := 0
	var count := 0
	for file_name in files:
		if not file_name.ends_with(".adt"):
			continue
		var stem := file_name.trim_suffix(".adt")
		var parts := stem.split("_")
		if parts.size() != 3 or parts[0] != map_name:
			continue
		if not parts[1].is_valid_int() or not parts[2].is_valid_int():
			continue
		var tx := parts[1].to_int()
		var ty := parts[2].to_int()
		if tile_x != null and tx != tile_x:
			continue
		if tile_y != null and ty != tile_y:
			continue
		if worker_index != null and tile_index % worker_count != int(worker_index):
			tile_index += 1
			continue
		tile_index += 1
		count += 1
	return count


func _parse_jobs_arg(value: String) -> int:
	if value.to_lower() == "auto":
		return maxi(1, OS.get_processor_count() - 1)
	return maxi(1, value.to_int())


func _get_arg_value(args: PackedStringArray, name: String, default_value: String) -> String:
	var idx := args.find(name)
	if idx >= 0 and idx + 1 < args.size():
		return args[idx + 1]
	return default_value


func _get_optional_int_arg(args: PackedStringArray, name: String) -> Variant:
	var idx := args.find(name)
	if idx >= 0 and idx + 1 < args.size() and args[idx + 1].is_valid_int():
		return args[idx + 1].to_int()
	return null


func _normalize_res_path(path: String) -> String:
	return path.rstrip("/")