extends SceneTree

const BenchmarkScript := preload("res://simulation/benchmark/RegionalScaleBenchmark.gd")
const DEFAULT_OUTPUT_PATH := "user://regional_scale_latest.json"
const DEFAULT_HOST_LABEL := "unspecified"
const SAMPLE_COUNT := 3
const DEFAULT_BUDGET := 8


func _initialize() -> void:
	call_deferred("_run")


func _run() -> void:
	var benchmark := BenchmarkScript.new()
	var samples: Array[Dictionary] = []
	for _sample_index in SAMPLE_COUNT:
		var sample := benchmark.run(DEFAULT_BUDGET)
		if not _valid_sample(sample):
			push_error("Regional scale benchmark produced an invalid default-budget sample")
			quit(1)
			return
		samples.append(sample)
	var budget_one := benchmark.run(1)
	var unlimited := benchmark.run(RegionalJobScheduler.UNLIMITED_BUDGET)
	if not _valid_sample(budget_one) or not _valid_sample(unlimited):
		push_error("Regional scale benchmark produced an invalid budget-parity sample")
		quit(1)
		return
	if not _deterministic_samples(samples + [budget_one, unlimited]):
		push_error("Regional scale benchmark diverged across samples or execution budgets")
		quit(1)
		return

	var report := {
		"schema_version": BenchmarkScript.SCHEMA_VERSION,
		"captured_utc": Time.get_datetime_string_from_system(true) + "Z",
		"engine_version": String(Engine.get_version_info().get("string", "unknown")),
		"platform": OS.get_name(),
		"processor_count": OS.get_processor_count(),
		"host_label": _get_argument_value("--host-label=", DEFAULT_HOST_LABEL),
		"workload_id": String(BenchmarkScript.WORKLOAD_ID),
		"benchmark_scope": "regional_data_and_scheduler_only",
		"rendered_throughput_measured": false,
		"full_economy_throughput_measured": false,
		"sample_count": SAMPLE_COUNT,
		"default_execution_budget": DEFAULT_BUDGET,
		"fixture": _fixture_summary(samples[0]),
		"default_budget": _summarize_timings(samples),
		"budget_parity": {
			"budget_one_final_checksum": budget_one["final_checksum"],
			"budget_one_drain_batches": budget_one["drain_batches"],
			"unlimited_final_checksum": unlimited["final_checksum"],
			"unlimited_drain_batches": unlimited["drain_batches"],
			"identical": true,
		},
	}
	var output_path := _get_argument_value("--output=", DEFAULT_OUTPUT_PATH)
	var output := FileAccess.open(output_path, FileAccess.WRITE)
	if output == null:
		push_error("Could not write regional scale benchmark report to %s" % output_path)
		quit(1)
		return
	output.store_string(JSON.stringify(report, "\t") + "\n")
	output.close()
	print(
		(
			"[BENCH] regional | %d named + %d aggregate | %d caravans | "
			+ "%d jobs | median %.2f ms total | %.2f MiB state"
		)
		% [
			report["fixture"]["named_person_count"],
			report["fixture"]["cohort_resident_count"],
			report["fixture"]["caravan_count"],
			report["fixture"]["initial_job_count"],
			float(report["default_budget"]["elapsed_usec_median"]) / 1000.0,
			float(report["fixture"]["regional_state_bytes"]) / (1024.0 * 1024.0),
		]
	)
	print("[BENCH] Report: %s" % ProjectSettings.globalize_path(output_path))
	quit(0)


func _valid_sample(sample: Dictionary) -> bool:
	return (
		not sample.is_empty()
		and bool(sample.get("fixture_valid", false))
		and bool(sample.get("roundtrip_valid", false))
		and bool(sample.get("headcount_conserved", false))
		and bool(sample.get("cargo_conserved", false))
		and int(sample.get("backlog_final", -1)) == 0
	)


func _deterministic_samples(samples: Array) -> bool:
	if samples.is_empty():
		return false
	var expected: Dictionary = samples[0]
	for sample: Dictionary in samples:
		for key in [
			"fixture_checksum",
			"execution_order_checksum",
			"final_checksum",
			"initial_job_count",
			"executed_job_count",
			"regional_state_bytes",
			"scheduler_state_bytes",
		]:
			if sample.get(key) != expected.get(key):
				return false
	return true


func _fixture_summary(sample: Dictionary) -> Dictionary:
	var result := {}
	for key in [
		"seed",
		"location_count",
		"settlement_count",
		"caravan_count",
		"detailed_members_per_caravan",
		"detailed_caravan_member_count",
		"named_person_count",
		"cohort_resident_count",
		"total_headcount",
		"max_active_visuals",
		"cargo_total",
		"initial_job_count",
		"regional_state_bytes",
		"scheduler_state_bytes",
		"fixture_checksum",
		"execution_order_checksum",
		"final_checksum",
	]:
		result[key] = sample[key]
	return result


func _summarize_timings(samples: Array[Dictionary]) -> Dictionary:
	var result := {}
	for key in [
		"fixture_build_usec",
		"roundtrip_usec",
		"scheduler_build_usec",
		"execution_usec",
		"elapsed_usec",
	]:
		var values: Array[int] = []
		for sample in samples:
			values.append(int(sample[key]))
		values.sort()
		result[key + "_samples"] = values
		result[key + "_median"] = values[values.size() / 2]
	result["drain_batches"] = samples[0]["drain_batches"]
	result["backlog_initial"] = samples[0]["backlog_initial"]
	result["backlog_peak"] = samples[0]["backlog_peak"]
	result["backlog_final"] = samples[0]["backlog_final"]
	return result


func _get_argument_value(prefix: String, default_value: String) -> String:
	for argument in OS.get_cmdline_user_args():
		if argument.begins_with(prefix):
			var value := argument.trim_prefix(prefix)
			if not value.is_empty():
				return value
	return default_value