class_name RegionalJobScheduler
extends RefCounted

const SCHEMA_VERSION := 1
const UNLIMITED_BUDGET := -1

var _jobs_by_id: Dictionary = {}
var _job_heap: Array[ScheduledJobRecord] = []
var _heap_index_by_id: Dictionary = {}
var _dedupe_job_ids: Dictionary = {}
var _sequence_job_ids: Dictionary = {}
var _next_stable_sequence := 0
var _last_drain_tick := -1
var _executed_count := 0
var _last_executed_job: ScheduledJobRecord


static func from_dictionary(record_data: Dictionary) -> RegionalJobScheduler:
	var scheduler := RegionalJobScheduler.new()
	return scheduler if scheduler.restore_from_dictionary(record_data) else null


func size() -> int:
	return _jobs_by_id.size()


func is_empty() -> bool:
	return _jobs_by_id.is_empty()


func has_job(job_id: StringName) -> bool:
	return _jobs_by_id.has(job_id)


func get_by_id(job_id: StringName) -> ScheduledJobRecord:
	return _jobs_by_id.get(job_id) as ScheduledJobRecord


func schedule(job: ScheduledJobRecord) -> bool:
	if (
		job == null
		or not job.is_valid()
		or has_job(job.get_job_id())
		or _sequence_job_ids.has(job.get_stable_sequence())
		or (_last_executed_job != null and not ScheduledJobRecord.compare(_last_executed_job, job))
	):
		return false
	var dedupe_key := job.get_dedupe_key()
	if not dedupe_key.is_empty() and _dedupe_job_ids.has(dedupe_key):
		return false
	var stored := ScheduledJobRecord.from_dictionary(job.to_dictionary())
	if stored == null:
		return false
	_jobs_by_id[stored.get_job_id()] = stored
	_sequence_job_ids[stored.get_stable_sequence()] = stored.get_job_id()
	if not dedupe_key.is_empty():
		_dedupe_job_ids[dedupe_key] = stored.get_job_id()
	_heap_push(stored)
	_next_stable_sequence = maxi(_next_stable_sequence, stored.get_stable_sequence() + 1)
	return true


func schedule_new(
	job_id: StringName,
	job_type: StringName,
	due_tick: int,
	phase: int,
	entity_id: StringName,
	payload: Dictionary = {},
	dedupe_key: StringName = &"",
	repeat_interval: int = ScheduledJobRecord.NO_REPEAT
) -> ScheduledJobRecord:
	var job := ScheduledJobRecord.create(
		job_id,
		job_type,
		due_tick,
		phase,
		entity_id,
		_next_stable_sequence,
		payload,
		dedupe_key,
		repeat_interval
	)
	return job if schedule(job) else null


func cancel(job_id: StringName) -> ScheduledJobRecord:
	var job := get_by_id(job_id)
	if job == null:
		return null
	var remove_index := int(_heap_index_by_id[job_id])
	var last_index := _job_heap.size() - 1
	if remove_index != last_index:
		_heap_swap(remove_index, last_index)
	_job_heap.pop_back()
	_heap_index_by_id.erase(job_id)
	_jobs_by_id.erase(job_id)
	_sequence_job_ids.erase(job.get_stable_sequence())
	var dedupe_key := job.get_dedupe_key()
	if not dedupe_key.is_empty():
		_dedupe_job_ids.erase(dedupe_key)
	if remove_index < _job_heap.size():
		var parent := (remove_index - 1) / 2
		if (
			remove_index > 0
			and ScheduledJobRecord.compare(_job_heap[remove_index], _job_heap[parent])
		):
			_heap_sift_up(remove_index)
		else:
			_heap_sift_down(remove_index)
	return job


func peek_next() -> ScheduledJobRecord:
	return _job_heap[0] if not _job_heap.is_empty() else null


func get_all_sorted() -> Array[ScheduledJobRecord]:
	var jobs := _job_heap.duplicate()
	jobs.sort_custom(ScheduledJobRecord.compare)
	return jobs


func get_due_count(current_tick: int) -> int:
	var count := 0
	for job in _jobs_by_id.values():
		if (job as ScheduledJobRecord).get_due_tick() <= current_tick:
			count += 1
	return count


func take_due(
	current_tick: int, execution_budget: int = UNLIMITED_BUDGET
) -> Array[ScheduledJobRecord]:
	var due: Array[ScheduledJobRecord] = []
	if current_tick < 0 or current_tick < _last_drain_tick or execution_budget < UNLIMITED_BUDGET:
		return due
	_last_drain_tick = current_tick
	if execution_budget == 0:
		return due
	var remaining := execution_budget
	while not _job_heap.is_empty():
		var next := _job_heap[0]
		if next.get_due_tick() > current_tick:
			break
		if remaining == 0:
			break
		var job := cancel(next.get_job_id())
		due.append(job)
		_executed_count += 1
		_last_executed_job = job
		if job.is_repeating():
			var next_occurrence := job.create_next_occurrence(_next_stable_sequence)
			if not schedule(next_occurrence):
				push_error("Could not schedule next occurrence for '%s'" % job.get_job_id())
		if remaining > 0:
			remaining -= 1
	return due


func execute_due(
	current_tick: int, execution_budget: int, executor: Callable
) -> Array[ScheduledJobRecord]:
	var executed: Array[ScheduledJobRecord] = []
	if not executor.is_valid() or execution_budget < UNLIMITED_BUDGET:
		return executed
	var remaining := execution_budget
	while remaining != 0:
		var batch := take_due(current_tick, 1)
		if batch.is_empty():
			break
		var job := batch[0]
		executed.append(job)
		executor.call(job)
		if remaining > 0:
			remaining -= 1
	return executed


func restore(records: Array[ScheduledJobRecord], cursor: Dictionary = {}) -> bool:
	var candidate := RegionalJobScheduler.new()
	for record in records:
		if not candidate.schedule(record):
			return false
	var next_sequence := int(cursor.get("next_stable_sequence", candidate._next_stable_sequence))
	var last_drain_tick := int(cursor.get("last_drain_tick", -1))
	var executed_count := int(cursor.get("executed_count", 0))
	var saved_last_order: Variant = cursor.get("last_executed_order")
	var last_executed_job := _parse_cursor_order(saved_last_order)
	if (
		next_sequence < candidate._next_stable_sequence
		or last_drain_tick < -1
		or executed_count < 0
		or (executed_count == 0 and last_executed_job != null)
		or (executed_count > 0 and last_executed_job == null)
		or (saved_last_order != null and last_executed_job == null)
		or (
			last_executed_job != null
			and (
				last_executed_job.get_due_tick() > last_drain_tick
				or next_sequence <= last_executed_job.get_stable_sequence()
			)
		)
	):
		return false
	if (
		last_executed_job != null
		and not candidate._job_heap.is_empty()
		and not ScheduledJobRecord.compare(last_executed_job, candidate._job_heap[0])
	):
		return false
	_adopt(candidate)
	_next_stable_sequence = next_sequence
	_last_drain_tick = last_drain_tick
	_executed_count = executed_count
	_last_executed_job = last_executed_job
	return true


func restore_from_dictionary(record_data: Dictionary) -> bool:
	if int(record_data.get("schema_version", -1)) != SCHEMA_VERSION:
		return false
	if not record_data.has_all(["jobs", "cursor"]) or not record_data["jobs"] is Array:
		return false
	if not record_data["cursor"] is Dictionary:
		return false
	if not (
		record_data["cursor"]
		. has_all(
			[
				"next_stable_sequence",
				"last_drain_tick",
				"executed_count",
				"last_executed_order",
			]
		)
	):
		return false
	var records: Array[ScheduledJobRecord] = []
	for saved_job in record_data["jobs"]:
		if not saved_job is Dictionary:
			return false
		var job := ScheduledJobRecord.from_dictionary(saved_job)
		if job == null:
			return false
		records.append(job)
	return restore(records, record_data["cursor"])


func get_cursor() -> Dictionary:
	return {
		"next_stable_sequence": _next_stable_sequence,
		"last_drain_tick": _last_drain_tick,
		"executed_count": _executed_count,
		"last_executed_order": _serialize_cursor_order(_last_executed_job),
	}


func to_dictionary() -> Dictionary:
	var jobs: Array[Dictionary] = []
	for job in get_all_sorted():
		jobs.append(job.to_dictionary())
	return {
		"schema_version": SCHEMA_VERSION,
		"jobs": jobs,
		"cursor": get_cursor(),
	}


func _heap_push(job: ScheduledJobRecord) -> void:
	var index := _job_heap.size()
	_job_heap.append(job)
	_heap_index_by_id[job.get_job_id()] = index
	_heap_sift_up(index)


func _heap_sift_up(start_index: int) -> void:
	var index := start_index
	while index > 0:
		var parent := (index - 1) / 2
		if not ScheduledJobRecord.compare(_job_heap[index], _job_heap[parent]):
			break
		_heap_swap(index, parent)
		index = parent


func _heap_sift_down(start_index: int) -> void:
	var index := start_index
	while true:
		var left := index * 2 + 1
		if left >= _job_heap.size():
			return
		var right := left + 1
		var smallest := left
		if (
			right < _job_heap.size()
			and ScheduledJobRecord.compare(_job_heap[right], _job_heap[left])
		):
			smallest = right
		if not ScheduledJobRecord.compare(_job_heap[smallest], _job_heap[index]):
			return
		_heap_swap(index, smallest)
		index = smallest


func _heap_swap(first_index: int, second_index: int) -> void:
	var first := _job_heap[first_index]
	_job_heap[first_index] = _job_heap[second_index]
	_job_heap[second_index] = first
	_heap_index_by_id[_job_heap[first_index].get_job_id()] = first_index
	_heap_index_by_id[_job_heap[second_index].get_job_id()] = second_index


func _adopt(source: RegionalJobScheduler) -> void:
	_jobs_by_id = source._jobs_by_id
	_job_heap = source._job_heap
	_heap_index_by_id = source._heap_index_by_id
	_dedupe_job_ids = source._dedupe_job_ids
	_sequence_job_ids = source._sequence_job_ids
	_next_stable_sequence = source._next_stable_sequence


static func _serialize_cursor_order(job: ScheduledJobRecord) -> Variant:
	if job == null:
		return null
	return {
		"due_tick": job.get_due_tick(),
		"phase": job.get_phase(),
		"entity_id": String(job.get_entity_id()),
		"stable_sequence": job.get_stable_sequence(),
	}


static func _parse_cursor_order(value: Variant) -> ScheduledJobRecord:
	if value == null:
		return null
	if (
		not value is Dictionary
		or not value.has_all(["due_tick", "phase", "entity_id", "stable_sequence"])
	):
		return null
	return ScheduledJobRecord.create(
		&"__cursor__",
		&"__cursor__",
		int(value["due_tick"]),
		int(value["phase"]),
		StringName(value["entity_id"]),
		int(value["stable_sequence"])
	)