class_name DeterministicRng
extends RefCounted

## A small deterministic random stream for simulation code.
##
## Each subsystem must create its own stream from the shared root seed and a
## stable stream ID. Saving both state and counter allows an exact continuation
## without coupling unrelated systems through one global random sequence.

const ALGORITHM_ID: StringName = &"xorshift32-sha256-v1"
const SNAPSHOT_VERSION := 1
const UINT32_MASK := 0xffffffff
const UINT32_RANGE := 0x100000000
const FALLBACK_NON_ZERO_STATE := 0x6d2b79f5

var _root_seed := 0
var _stream_id: StringName = &"default"
var _state := FALLBACK_NON_ZERO_STATE
var _counter := 0


static func from_root_seed(root_seed: int, stream_id: StringName) -> Variant:
	# Loading by resource path keeps this factory usable on a clean checkout,
	# before Godot has generated its global class-name cache.
	var result: Variant = load("res://scripts/core/deterministic_rng.gd").new()
	result.reset(root_seed, stream_id)
	return result


func reset(root_seed: int, stream_id: StringName) -> void:
	_root_seed = root_seed
	_stream_id = stream_id
	_state = _derive_non_zero_state(root_seed, stream_id)
	_counter = 0


func next_u32() -> int:
	# xorshift32 operates only on non-negative 32-bit values. Intermediates stay
	# below the signed 64-bit limit used by GDScript integers.
	var value := _state
	value ^= (value << 13) & UINT32_MASK
	value ^= value >> 17
	value ^= (value << 5) & UINT32_MASK
	_state = value & UINT32_MASK
	_counter += 1
	return _state


func randi_range(minimum: int, maximum: int) -> int:
	if minimum > maximum:
		push_error("DeterministicRng range minimum must not exceed maximum")
		return minimum

	# A crossing-zero subtraction can overflow only for ranges much larger than
	# this generator supports, so reject it before evaluating the subtraction.
	var width := 0
	if minimum < 0 and maximum > 0:
		if maximum > minimum + UINT32_MASK:
			push_error("DeterministicRng supports ranges of at most 2^32 values")
			return minimum
		width = maximum - minimum
	else:
		width = maximum - minimum
	if width > UINT32_MASK:
		push_error("DeterministicRng supports ranges of at most 2^32 values")
		return minimum

	var span := width + 1
	if span == 1:
		return minimum

	# Rejection sampling avoids modulo bias for spans which do not divide 2^32.
	var acceptance_limit := UINT32_RANGE - (UINT32_RANGE % span)
	while true:
		var sample := next_u32()
		if sample < acceptance_limit:
			return minimum + (sample % span)
	return minimum


func snapshot() -> Dictionary:
	return {
		"version": SNAPSHOT_VERSION,
		"algorithm": String(ALGORITHM_ID),
		"root_seed": _root_seed,
		"stream_id": String(_stream_id),
		"state": _state,
		"counter": _counter,
	}


func restore(saved_state: Dictionary) -> bool:
	if int(saved_state.get("version", -1)) != SNAPSHOT_VERSION:
		push_error("Unsupported DeterministicRng snapshot version")
		return false
	if StringName(saved_state.get("algorithm", "")) != ALGORITHM_ID:
		push_error("DeterministicRng snapshot algorithm does not match")
		return false

	var restored_state := int(saved_state.get("state", 0))
	var restored_counter := int(saved_state.get("counter", -1))
	if restored_state <= 0 or restored_state > UINT32_MASK:
		push_error("DeterministicRng snapshot contains an invalid state")
		return false
	if restored_counter < 0:
		push_error("DeterministicRng snapshot contains an invalid counter")
		return false

	_root_seed = int(saved_state.get("root_seed", 0))
	_stream_id = StringName(saved_state.get("stream_id", "default"))
	_state = restored_state
	_counter = restored_counter
	return true


func get_root_seed() -> int:
	return _root_seed


func get_stream_id() -> StringName:
	return _stream_id


func get_state() -> int:
	return _state


func get_counter() -> int:
	return _counter


static func _derive_non_zero_state(root_seed: int, stream_id: StringName) -> int:
	var stream_bytes := String(stream_id).to_utf8_buffer()
	var context := HashingContext.new()
	if context.start(HashingContext.HASH_SHA256) != OK:
		push_error("Could not initialize the deterministic RNG seed hash")
		return FALLBACK_NON_ZERO_STATE
	context.update("the-tree:deterministic-rng:v1\n".to_utf8_buffer())
	context.update(("%d\n%d\n" % [root_seed, stream_bytes.size()]).to_utf8_buffer())
	context.update(stream_bytes)
	var digest := context.finish()

	var derived_state := 0
	for byte_index in range(4):
		derived_state = (derived_state << 8) | int(digest[byte_index])
	if derived_state == 0:
		# The all-zero state is absorbing for xorshift32. A second digest word is
		# already independently distributed; retain a literal fallback as a final
		# guard for the astronomically unlikely double-zero case.
		for byte_index in range(4, 8):
			derived_state = (derived_state << 8) | int(digest[byte_index])
	if derived_state == 0:
		derived_state = FALLBACK_NON_ZERO_STATE
	return derived_state