extends SceneTree

## Compares original-client JPG/PNG references with OpenWC checkpoint PNGs.
## Usage: godot --headless --path . --script res://src/tools/compare_render_checkpoints.gd --
##   --reference <directory> --candidate <directory> [--output <report.json>]
##   [--pixel-threshold 0.04] [--mean-threshold 0.015] [--changed-ratio-threshold 0.01]
## Use --self-test for a synthetic regression check that writes only under user://.

const REPORT_SCHEMA_VERSION := 1
const DEFAULT_PIXEL_THRESHOLD := 0.04
const DEFAULT_MEAN_THRESHOLD := 0.015
const DEFAULT_CHANGED_RATIO_THRESHOLD := 0.01


func _initialize() -> void:
	var arguments := _parse_arguments(OS.get_cmdline_user_args())
	if arguments.has("error"):
		push_error("RENDER_CHECKPOINT_DIFF: %s" % arguments.error)
		quit(2)
		return
	if arguments.get("self_test", false):
		quit(_run_self_test())
		return

	var reference_directory := String(arguments.get("reference", ""))
	var candidate_directory := String(arguments.get("candidate", ""))
	if reference_directory.is_empty() or candidate_directory.is_empty():
		push_error("RENDER_CHECKPOINT_DIFF: --reference and --candidate are required")
		quit(2)
		return

	var report := _compare_directories(reference_directory, candidate_directory, arguments)
	var output_path := String(arguments.get("output", ""))
	if not output_path.is_empty() and not _write_report(output_path, report):
		quit(2)
		return
	print("RENDER_CHECKPOINT_DIFF %s compared=%d failed=%d missing=%d" % [
		"PASS" if report.passed else "FAIL",
		report.compared_count,
		report.failed_count,
		report.missing_count,
	])
	quit(0 if report.passed else 1)


func _parse_arguments(raw_arguments: PackedStringArray) -> Dictionary:
	var parsed := {
		"pixel_threshold": DEFAULT_PIXEL_THRESHOLD,
		"mean_threshold": DEFAULT_MEAN_THRESHOLD,
		"changed_ratio_threshold": DEFAULT_CHANGED_RATIO_THRESHOLD,
	}
	var index := 0
	while index < raw_arguments.size():
		var argument := raw_arguments[index]
		if argument == "--self-test":
			parsed.self_test = true
			index += 1
			continue
		if argument not in ["--reference", "--candidate", "--output", "--pixel-threshold", "--mean-threshold", "--changed-ratio-threshold"]:
			return {"error": "unknown argument: %s" % argument}
		if index + 1 >= raw_arguments.size():
			return {"error": "missing value for %s" % argument}
		var key := argument.trim_prefix("--").replace("-", "_")
		var value := raw_arguments[index + 1]
		if key.ends_with("threshold"):
			if not value.is_valid_float():
				return {"error": "%s requires a number" % argument}
			parsed[key] = float(value)
		else:
			parsed[key] = value
		index += 2
	return parsed


func _compare_directories(reference_directory: String, candidate_directory: String, options: Dictionary) -> Dictionary:
	reference_directory = ProjectSettings.globalize_path(reference_directory)
	candidate_directory = ProjectSettings.globalize_path(candidate_directory)
	var reference_files := _reference_image_file_names(reference_directory)
	var results: Array[Dictionary] = []
	var failed_count := 0
	var missing_count := 0
	for file_name in reference_files:
		var reference_path := reference_directory.path_join(file_name)
		var checkpoint_name := file_name.get_basename()
		var candidate_file_names := _candidate_file_names(candidate_directory, checkpoint_name)
		if candidate_file_names.is_empty():
			results.append({"checkpoint": checkpoint_name, "reference": file_name, "status": "missing_candidate"})
			missing_count += 1
			continue
		for candidate_file_name in candidate_file_names:
			var candidate_path := candidate_directory.path_join(candidate_file_name)
			var comparison := _compare_images(reference_path, candidate_path, options)
			comparison.checkpoint = checkpoint_name
			comparison.reference = file_name
			comparison.candidate = candidate_file_name
			results.append(comparison)
			if comparison.status != "passed":
				failed_count += 1
	var no_reference_images := reference_files.is_empty()
	return {
		"schema_version": REPORT_SCHEMA_VERSION,
		"passed": not no_reference_images and failed_count == 0 and missing_count == 0,
		"reference_directory": reference_directory,
		"candidate_directory": candidate_directory,
		"pixel_threshold": options.pixel_threshold,
		"mean_threshold": options.mean_threshold,
		"changed_ratio_threshold": options.changed_ratio_threshold,
		"compared_count": results.size() - missing_count,
		"failed_count": failed_count,
		"missing_count": missing_count,
		"no_reference_images": no_reference_images,
		"results": results,
	}


func _compare_images(reference_path: String, candidate_path: String, options: Dictionary) -> Dictionary:
	var reference_image := Image.load_from_file(reference_path)
	var candidate_image := Image.load_from_file(candidate_path)
	if reference_image == null or candidate_image == null or reference_image.is_empty() or candidate_image.is_empty():
		return {"status": "load_error"}
	if reference_image.get_size() != candidate_image.get_size():
		return {
			"status": "size_mismatch",
			"reference_size": [reference_image.get_width(), reference_image.get_height()],
			"candidate_size": [candidate_image.get_width(), candidate_image.get_height()],
		}

	var error_sum := 0.0
	var changed_pixel_count := 0
	var pixel_count := reference_image.get_width() * reference_image.get_height()
	for y in reference_image.get_height():
		for x in reference_image.get_width():
			var perceptual_error := _perceptual_color_error(reference_image.get_pixel(x, y), candidate_image.get_pixel(x, y))
			error_sum += perceptual_error
			if perceptual_error > float(options.pixel_threshold):
				changed_pixel_count += 1
	var mean_perceptual_error := error_sum / float(pixel_count)
	var changed_pixel_ratio := float(changed_pixel_count) / float(pixel_count)
	var passed := mean_perceptual_error <= float(options.mean_threshold) and changed_pixel_ratio <= float(options.changed_ratio_threshold)
	return {
		"status": "passed" if passed else "threshold_exceeded",
		"width": reference_image.get_width(),
		"height": reference_image.get_height(),
		"mean_perceptual_error": mean_perceptual_error,
		"changed_pixel_ratio": changed_pixel_ratio,
	}


func _perceptual_color_error(reference_color: Color, candidate_color: Color) -> float:
	var reference_linear := reference_color.srgb_to_linear()
	var candidate_linear := candidate_color.srgb_to_linear()
	var red_delta := absf(reference_linear.r - candidate_linear.r)
	var green_delta := absf(reference_linear.g - candidate_linear.g)
	var blue_delta := absf(reference_linear.b - candidate_linear.b)
	var alpha_delta := absf(reference_linear.a - candidate_linear.a)
	return 0.2126 * red_delta + 0.7152 * green_delta + 0.0722 * blue_delta + 0.25 * alpha_delta


func _reference_image_file_names(directory_path: String) -> Array[String]:
	var directory := DirAccess.open(directory_path)
	if directory == null:
		return []
	var file_names: Array[String] = []
	for file_name in directory.get_files():
		var lower_name := file_name.to_lower()
		if not lower_name.begins_with("diagnostic_") and (lower_name.ends_with(".png") or lower_name.ends_with(".jpg") or lower_name.ends_with(".jpeg")):
			file_names.append(file_name)
	file_names.sort()
	return file_names


func _candidate_file_names(candidate_directory: String, checkpoint_name: String) -> Array[String]:
	var candidates: Array[String] = []
	for pass_name in ["cold_process", "warm_revisit"]:
		var file_name := "%s__%s.png" % [checkpoint_name, pass_name]
		if FileAccess.file_exists(candidate_directory.path_join(file_name)):
			candidates.append(file_name)
	if candidates.is_empty():
		var exact_file_name := "%s.png" % checkpoint_name
		if FileAccess.file_exists(candidate_directory.path_join(exact_file_name)):
			candidates.append(exact_file_name)
	return candidates


func _write_report(output_path: String, report: Dictionary) -> bool:
	var absolute_path := ProjectSettings.globalize_path(output_path)
	var error := DirAccess.make_dir_recursive_absolute(absolute_path.get_base_dir())
	if error != OK:
		push_error("RENDER_CHECKPOINT_DIFF: cannot create report directory: %s" % error_string(error))
		return false
	var report_file := FileAccess.open(absolute_path, FileAccess.WRITE)
	if report_file == null:
		push_error("RENDER_CHECKPOINT_DIFF: cannot write report: %s" % output_path)
		return false
	report_file.store_string(JSON.stringify(report, "  "))
	return true


func _run_self_test() -> int:
	var root := "user://render_checkpoint_diff_self_test"
	var reference_directory := root.path_join("reference")
	var identical_directory := root.path_join("identical")
	var changed_directory := root.path_join("changed")
	for directory_path in [reference_directory, identical_directory, changed_directory]:
		DirAccess.make_dir_recursive_absolute(ProjectSettings.globalize_path(directory_path))
	var reference_image := Image.create(2, 2, false, Image.FORMAT_RGBA8)
	reference_image.fill(Color(0.25, 0.5, 0.75, 1.0))
	var changed_image := reference_image.duplicate()
	changed_image.set_pixel(0, 0, Color.WHITE)
	if reference_image.save_jpg(reference_directory.path_join("synthetic.jpg"), 1.0) != OK:
		push_error("RENDER_CHECKPOINT_DIFF SELF_TEST: cannot save JPG reference")
		return 1
	for pass_name in ["cold_process", "warm_revisit"]:
		if reference_image.save_png(identical_directory.path_join("synthetic__%s.png" % pass_name)) != OK:
			return 1
		if changed_image.save_png(changed_directory.path_join("synthetic__%s.png" % pass_name)) != OK:
			return 1
	var options := {
		"pixel_threshold": DEFAULT_PIXEL_THRESHOLD,
		"mean_threshold": DEFAULT_MEAN_THRESHOLD,
		"changed_ratio_threshold": DEFAULT_CHANGED_RATIO_THRESHOLD,
	}
	var identical_report := _compare_directories(reference_directory, identical_directory, options)
	var changed_report := _compare_directories(reference_directory, changed_directory, options)
	if not identical_report.passed or changed_report.passed:
		push_error("RENDER_CHECKPOINT_DIFF SELF_TEST: expected identical pass and changed failure")
		return 1
	print("RENDER_CHECKPOINT_DIFF SELF_TEST PASS")
	return 0