class_name RegionalScaleBenchmark
extends RefCounted

# This workload measures primitive regional records and deterministic scheduler
# draining only. It does not load scenes or exercise the complete economy.
const SCHEMA_VERSION := 1
const WORKLOAD_ID := &"regional_records_scheduler_v1"
const WORLD_ID := &"regional_benchmark"
const FIXED_SEED := 424242
const LOCATION_COUNT := 5
const SETTLEMENT_COUNT := 5
const CARAVAN_COUNT := 20
const DETAILED_MEMBERS_PER_CARAVAN := 12
const NAMED_PERSON_COUNT := 2000
const COHORT_RESIDENT_COUNT := 50000
const MAX_ACTIVE_VISUALS := 40
const MAX_DUE_TICK := 256
const ROLE_IDS: Array[StringName] = [&"artisan", &"farmer", &"guard", &"merchant", &"scholar"]


func build_fixture(seed_value: int = FIXED_SEED) -> RegionalWorldState:
	var locations: Array[LocationStateRecord] = []
	var settlements: Array[SettlementStateRecord] = []
	var routes: Array[RouteStateRecord] = []
	var groups: Array[MobileGroupStateRecord] = []
	var persons: Array[PersonStateRecord] = []
	var cohorts: Array[PopulationCohortRecord] = []
	var polities: Array[PolityStateRecord] = []
	var relations: Array[DiplomaticRelationRecord] = []

	for index in LOCATION_COUNT:
		var location_id := _location_id(index)
		var location := LocationStateRecord.create(
			location_id,
			WORLD_ID,
			&"settlement",
			"Regional Settlement %02d" % index,
			_address(location_id, Vector3(float(index) * 100.0, 0.0, 0.0))
		)
		var settlement := SettlementStateRecord.create(
			_settlement_id(index),
			location_id,
			"Regional Settlement %02d" % index,
			_polity_id(index),
			2,
			_group_id(index),
			0,
			{&"food": 10000.0 + index * 100.0, &"wood": 5000.0 + index * 50.0}
		)
		var polity := PolityStateRecord.create(
			_polity_id(index), "Regional Polity %02d" % index, location_id, [_settlement_id(index)]
		)
		if location == null or settlement == null or polity == null:
			return null
		locations.append(location)
		settlements.append(settlement)
		polities.append(polity)

	for index in LOCATION_COUNT:
		var route := RouteStateRecord.create(
			_route_id(index),
			_location_id(index),
			_location_id((index + 1) % LOCATION_COUNT),
			20 + index,
			true
		)
		if route == null:
			return null
		routes.append(route)

	for group_index in CARAVAN_COUNT:
		var route_index := group_index % LOCATION_COUNT
		var member_ids: Array[StringName] = []
		for member_index in DETAILED_MEMBERS_PER_CARAVAN:
			member_ids.append(_person_id(group_index * DETAILED_MEMBERS_PER_CARAVAN + member_index))
		var departure_tick := group_index * 3
		var cargo := {
			&"food":
			float(
				KeyedRandom.range_int(seed_value, &"cargo_food", _group_id(group_index), 0, 50, 99)
			),
			&"herb":
			float(
				KeyedRandom.range_int(seed_value, &"cargo_herb", _group_id(group_index), 0, 10, 29)
			),
			&"wood":
			float(
				KeyedRandom.range_int(seed_value, &"cargo_wood", _group_id(group_index), 0, 20, 49)
			),
		}
		var group := MobileGroupStateRecord.create_on_route(
			_group_id(group_index),
			&"caravan",
			"Regional Caravan %02d" % group_index,
			_polity_id(route_index),
			_route_id(route_index),
			_location_id(route_index),
			_location_id((route_index + 1) % LOCATION_COUNT),
			departure_tick,
			departure_tick + 20 + route_index,
			member_ids,
			{},
			cargo,
			500.0
		)
		if group == null:
			return null
		groups.append(group)

	for person_index in NAMED_PERSON_COUNT:
		var person: PersonStateRecord
		if person_index < CARAVAN_COUNT * DETAILED_MEMBERS_PER_CARAVAN:
			var group_index := person_index / DETAILED_MEMBERS_PER_CARAVAN
			var home_index := group_index % LOCATION_COUNT
			person = PersonStateRecord.create_in_group(
				_person_id(person_index),
				"Regional Person %04d" % person_index,
				_role_id(seed_value, person_index),
				_polity_id(home_index),
				_location_id(home_index),
				_group_id(group_index)
			)
		else:
			var location_index := (
				(person_index - CARAVAN_COUNT * DETAILED_MEMBERS_PER_CARAVAN) % LOCATION_COUNT
			)
			var location_id := _location_id(location_index)
			person = PersonStateRecord.create_at_location(
				_person_id(person_index),
				"Regional Person %04d" % person_index,
				_role_id(seed_value, person_index),
				_polity_id(location_index),
				location_id,
				_address(
					location_id,
					Vector3(float(person_index % 64), 0.0, float((person_index / 64) % 64))
				)
			)
		if person == null:
			return null
		persons.append(person)

	var residents_per_location := COHORT_RESIDENT_COUNT / LOCATION_COUNT
	for index in LOCATION_COUNT:
		var cohort := PopulationCohortRecord.create_at_location(
			_cohort_id(index),
			&"aggregate_residents",
			_polity_id(index),
			residents_per_location,
			_location_id(index)
		)
		if cohort == null:
			return null
		cohorts.append(cohort)

	return RegionalWorldState.create(
		WORLD_ID, locations, settlements, routes, groups, persons, cohorts, polities, relations
	)


func run(execution_budget: int = 8, seed_value: int = FIXED_SEED) -> Dictionary:
	if execution_budget not in [1, 8, RegionalJobScheduler.UNLIMITED_BUDGET]:
		return {}
	var total_started_usec := Time.get_ticks_usec()
	var fixture_started_usec := Time.get_ticks_usec()
	var state := build_fixture(seed_value)
	var fixture_build_usec := maxi(Time.get_ticks_usec() - fixture_started_usec, 1)
	if state == null:
		return {}
	var fixture_audit := audit_fixture(state)
	if not fixture_audit["valid"]:
		return {}
	var state_dictionary := state.to_dictionary()
	var state_json := JSON.stringify(state_dictionary)
	var state_bytes := state_json.to_utf8_buffer().size()
	var fixture_checksum := state.checksum()

	var roundtrip_started_usec := Time.get_ticks_usec()
	var restored := RegionalWorldState.from_dictionary(state_dictionary)
	var roundtrip_usec := maxi(Time.get_ticks_usec() - roundtrip_started_usec, 1)
	if (
		restored == null
		or restored.to_dictionary() != state_dictionary
		or restored.checksum() != fixture_checksum
	):
		return {}

	var scheduler_started_usec := Time.get_ticks_usec()
	var scheduler := _build_scheduler(state, seed_value)
	var scheduler_build_usec := maxi(Time.get_ticks_usec() - scheduler_started_usec, 1)
	if scheduler == null:
		return {}
	var initial_jobs := scheduler.size()
	var scheduler_state_bytes := JSON.stringify(scheduler.to_dictionary()).to_utf8_buffer().size()
	var execution_order: Array[StringName] = []
	var processed_named_people := 0
	var processed_cohort_residents := 0
	var processed_cargo := 0.0
	var backlog_peak := scheduler.size()
	var drain_batches := 0
	var max_batch_size := 0
	var execution_started_usec := Time.get_ticks_usec()
	while not scheduler.is_empty():
		backlog_peak = maxi(backlog_peak, scheduler.size())
		var batch := scheduler.take_due(MAX_DUE_TICK, execution_budget)
		if batch.is_empty():
			return {}
		drain_batches += 1
		max_batch_size = maxi(max_batch_size, batch.size())
		for job in batch:
			execution_order.append(job.get_job_id())
			var payload := job.get_payload()
			processed_named_people += int(payload.get("named_people", 0))
			processed_cohort_residents += int(payload.get("cohort_residents", 0))
			processed_cargo += float(payload.get("cargo_total", 0.0))
	var execution_usec := maxi(Time.get_ticks_usec() - execution_started_usec, 1)
	var final_audit := audit_fixture(state)
	var headcount_conserved: bool = (
		final_audit["named_person_count"] == fixture_audit["named_person_count"]
		and final_audit["cohort_resident_count"] == fixture_audit["cohort_resident_count"]
		and processed_named_people == NAMED_PERSON_COUNT
		and processed_cohort_residents == COHORT_RESIDENT_COUNT
	)
	var cargo_conserved: bool = (
		is_equal_approx(float(final_audit["cargo_total"]), float(fixture_audit["cargo_total"]))
		and is_equal_approx(processed_cargo, float(fixture_audit["cargo_total"]))
	)
	var order_checksum := JSON.stringify(execution_order).sha256_text()
	var final_checksum := (
		JSON
		. stringify(
			{
				"fixture_checksum": fixture_checksum,
				"execution_order": execution_order,
				"processed_cargo": processed_cargo,
				"processed_cohort_residents": processed_cohort_residents,
				"processed_named_people": processed_named_people,
			}
		)
		. sha256_text()
	)
	return {
		"schema_version": SCHEMA_VERSION,
		"workload_id": String(WORKLOAD_ID),
		"benchmark_scope": "regional_data_and_scheduler_only",
		"rendered_throughput_measured": false,
		"full_economy_throughput_measured": false,
		"seed": seed_value,
		"location_count": LOCATION_COUNT,
		"settlement_count": SETTLEMENT_COUNT,
		"caravan_count": CARAVAN_COUNT,
		"detailed_members_per_caravan": DETAILED_MEMBERS_PER_CARAVAN,
		"detailed_caravan_member_count": fixture_audit["detailed_caravan_member_count"],
		"named_person_count": fixture_audit["named_person_count"],
		"cohort_resident_count": fixture_audit["cohort_resident_count"],
		"total_headcount": fixture_audit["total_headcount"],
		"max_active_visuals": MAX_ACTIVE_VISUALS,
		"fixture_valid": fixture_audit["valid"],
		"roundtrip_valid": true,
		"headcount_conserved": headcount_conserved,
		"cargo_conserved": cargo_conserved,
		"cargo_total": fixture_audit["cargo_total"],
		"processed_cargo_total": processed_cargo,
		"fixture_checksum": fixture_checksum,
		"execution_order_checksum": order_checksum,
		"final_checksum": final_checksum,
		"execution_order": execution_order,
		"initial_job_count": initial_jobs,
		"executed_job_count": execution_order.size(),
		"execution_budget": execution_budget,
		"backlog_initial": initial_jobs,
		"backlog_peak": backlog_peak,
		"backlog_final": scheduler.size(),
		"drain_batches": drain_batches,
		"max_batch_size": max_batch_size,
		"regional_state_bytes": state_bytes,
		"scheduler_state_bytes": scheduler_state_bytes,
		"fixture_build_usec": fixture_build_usec,
		"roundtrip_usec": roundtrip_usec,
		"scheduler_build_usec": scheduler_build_usec,
		"execution_usec": execution_usec,
		"elapsed_usec": maxi(Time.get_ticks_usec() - total_started_usec, 1),
	}


func audit_fixture(state: RegionalWorldState) -> Dictionary:
	if state == null:
		return {"valid": false}
	var valid := (
		state.get_world_id() == WORLD_ID
		and state.get_all_locations().size() == LOCATION_COUNT
		and state.get_all_settlements().size() == SETTLEMENT_COUNT
		and state.get_all_routes().size() == LOCATION_COUNT
		and state.get_all_mobile_groups().size() == CARAVAN_COUNT
		and state.get_all_persons().size() == NAMED_PERSON_COUNT
		and state.get_all_population_cohorts().size() == LOCATION_COUNT
		and state.get_all_polities().size() == LOCATION_COUNT
	)
	var detailed_members := 0
	var cargo_total := 0.0
	for group in state.get_all_mobile_groups():
		var member_ids := group.get_detailed_member_ids()
		detailed_members += member_ids.size()
		cargo_total += RegionalStateSupport.amount_ledger_total(group.get_cargo_ledger())
		valid = (
			valid
			and member_ids.size() == DETAILED_MEMBERS_PER_CARAVAN
			and (
				state.get_persons_in_group(group.get_group_id()).size()
				== DETAILED_MEMBERS_PER_CARAVAN
			)
			and group.is_on_route()
		)
		for person_id in member_ids:
			valid = valid and state.get_entity_type(person_id) == &"person"
	var cohort_residents := 0
	for cohort in state.get_all_population_cohorts():
		cohort_residents += cohort.get_count()
		valid = valid and cohort.is_at_location()
	for index in LOCATION_COUNT:
		valid = (
			valid
			and state.get_settlements_at_location(_location_id(index)).size() == 1
			and state.get_entity_type(_location_id(index)) == &"location"
			and state.get_entity_type(_settlement_id(index)) == &"settlement"
			and state.get_entity_type(_polity_id(index)) == &"polity"
		)
	valid = (
		valid
		and detailed_members == CARAVAN_COUNT * DETAILED_MEMBERS_PER_CARAVAN
		and cohort_residents == COHORT_RESIDENT_COUNT
	)
	return {
		"valid": valid,
		"named_person_count": state.get_all_persons().size(),
		"detailed_caravan_member_count": detailed_members,
		"cohort_resident_count": cohort_residents,
		"total_headcount": state.get_all_persons().size() + cohort_residents,
		"cargo_total": cargo_total,
	}


func _build_scheduler(state: RegionalWorldState, seed_value: int) -> RegionalJobScheduler:
	var scheduler := RegionalJobScheduler.new()
	for person in state.get_all_persons():
		if (
			scheduler.schedule_new(
				&"regional_person:%s" % person.get_person_id(),
				&"regional_person_review",
				KeyedRandom.range_int(
					seed_value, &"regional_person_job", person.get_person_id(), 0, 1, MAX_DUE_TICK
				),
				0,
				person.get_person_id(),
				{"named_people": 1}
			)
			== null
		):
			return null
	for cohort in state.get_all_population_cohorts():
		if (
			scheduler.schedule_new(
				&"regional_cohort:%s" % cohort.get_cohort_id(),
				&"regional_cohort_review",
				KeyedRandom.range_int(
					seed_value, &"regional_cohort_job", cohort.get_cohort_id(), 0, 1, MAX_DUE_TICK
				),
				1,
				cohort.get_cohort_id(),
				{"cohort_residents": cohort.get_count()}
			)
			== null
		):
			return null
	for group in state.get_all_mobile_groups():
		if (
			scheduler.schedule_new(
				&"regional_caravan:%s" % group.get_group_id(),
				&"regional_caravan_review",
				KeyedRandom.range_int(
					seed_value, &"regional_caravan_job", group.get_group_id(), 0, 1, MAX_DUE_TICK
				),
				2,
				group.get_group_id(),
				{"cargo_total": RegionalStateSupport.amount_ledger_total(group.get_cargo_ledger())}
			)
			== null
		):
			return null
	return scheduler


func _role_id(seed_value: int, person_index: int) -> StringName:
	var role_index := KeyedRandom.range_int(
		seed_value, &"person_role", _person_id(person_index), 0, 0, ROLE_IDS.size() - 1
	)
	return ROLE_IDS[role_index]


func _address(location_id: StringName, position: Vector3) -> SpatialAddress:
	return SpatialAddress.create(WORLD_ID, location_id, position)


func _location_id(index: int) -> StringName:
	return StringName("regional_location_%02d" % index)


func _settlement_id(index: int) -> StringName:
	return StringName("regional_settlement_%02d" % index)


func _route_id(index: int) -> StringName:
	return StringName("regional_route_%02d" % index)


func _group_id(index: int) -> StringName:
	return StringName("regional_caravan_%02d" % index)


func _person_id(index: int) -> StringName:
	return StringName("regional_person_%04d" % index)


func _cohort_id(index: int) -> StringName:
	return StringName("regional_residents_%02d" % index)


func _polity_id(index: int) -> StringName:
	return StringName("regional_polity_%02d" % index)