--!strict
local Scheduler = {}
Scheduler.__index = Scheduler
export type Handle = {}
type Job = { Time: number, Order: number, Index: number, Handle: Handle, Callback: () -> () }
type State = { _heap: { Job }, _jobs: { [Handle]: Job? }, _order: number, _destroyed: boolean }
export type Scheduler = typeof(setmetatable({} :: State, Scheduler))

local function earlier(a: Job, b: Job): boolean
	return a.Time < b.Time or (a.Time == b.Time and a.Order < b.Order)
end

local function repair(heap: { Job }, index: number): ()
	local job, size = heap[index], #heap
	while index > 1 do
		local parent = math.floor(index / 2)
		if not earlier(job, heap[parent]) then break end
		heap[index] = heap[parent]
		heap[index].Index = index
		index = parent
	end
	while index * 2 <= size do
		local child = index * 2
		if child < size and earlier(heap[child + 1], heap[child]) then child += 1 end
		if not earlier(heap[child], job) then break end
		heap[index] = heap[child]
		heap[index].Index = index
		index = child
	end
	heap[index] = job
	job.Index = index
end

local function remove(heap: { Job }, job: Job): ()
	local index = job.Index
	if index == 0 then return end
	local last = table.remove(heap)
	job.Index = 0
	if last and index <= #heap then
		heap[index] = last
		repair(heap, index)
	end
end

function Scheduler.new(): Scheduler
	return setmetatable({ _heap = {}, _jobs = {}, _order = 0, _destroyed = false }, Scheduler)
end
function Scheduler.Schedule(self: Scheduler, delay: number, callback: () -> ()): Handle
	assert(not self._destroyed, "Scheduler is destroyed")
	assert(delay >= 0 and delay < math.huge, "Delay must be finite and nonnegative")
	assert(type(callback) == "function", "Callback must be a function")
	local heap, handle = self._heap, table.freeze({})
	self._order += 1
	local job: Job = {
		Time = os.clock() + delay, Order = self._order, Index = #heap + 1,
		Handle = handle, Callback = callback,
	}
	self._jobs[handle] = job
	heap[job.Index] = job
	repair(heap, job.Index)
	return handle
end

function Scheduler.Cancel(self: Scheduler, handle: Handle): boolean
	local job = self._jobs[handle]
	if not job then return false end
	remove(self._heap, job)
	self._jobs[handle] = nil
	return true
end

local function dispatch(self: Scheduler, handle: Handle): ()
	local job = self._jobs[handle]
	if not job then return end
	self._jobs[handle] = nil
	job.Callback()
end

function Scheduler.Update(self: Scheduler, limit: number?): number
	local budget = limit or 128
	assert(budget >= 0 and budget < math.huge and budget % 1 == 0, "Limit must be a finite nonnegative integer")
	local heap, now, count = self._heap, os.clock(), 0
	while count < budget do
		local job = heap[1]
		if not job or job.Time > now then break end
		remove(heap, job)
		count += 1
		task.defer(dispatch, self, job.Handle)
	end
	return count
end
function Scheduler.Clear(self: Scheduler): ()
	table.clear(self._jobs)
	table.clear(self._heap)
end
function Scheduler.Destroy(self: Scheduler): ()
	self._destroyed = true
	self._heap, self._jobs = {}, {}
end
return table.freeze(Scheduler)