local G={}local v={[Enum.DataStoreRequestType.GetAsync]={};[Enum.DataStoreRequestType.GetSortedAsync]={},[Enum.DataStoreRequestType.OnUpdate]={};[Enum.DataStoreRequestType.SetIncrementAsync]={};[Enum.DataStoreRequestType.SetIncrementSortedAsync]={}}local function d()for v,d in pairs(ConstantsMapping)do G[v]=d.START end G[Enum.DataStoreRequestType.UpdateAsync]=math.min(G[Enum.DataStoreRequestType.GetAsync],G[Enum.DataStoreRequestType.SetIncrementAsync])end local function l(v,d,l,C)if not Constants.BUDGETING_ENABLED then return end local N=d.RATE+C*d.RATE_PLR G[v]=math.min(G[v]+l*N,d.MAX_FACTOR*N)end local function C(v)if not Constants.BUDGETING_ENABLED then return end for v,d in pairs(v)do if G[d]then G[d]=math.max(0,G[d]-1)end end G[Enum.DataStoreRequestType.UpdateAsync]=math.min(G[Enum.DataStoreRequestType.GetAsync],G[Enum.DataStoreRequestType.SetIncrementAsync])end local function N(v)if not Constants.BUDGETING_ENABLED then return true end for v,d in pairs(v)do if G[d]and G[d]<1 then return false end end return true end local O=false if RunService:IsServer()then d();(coroutine.wrap(function()local d=tick()while Utils.accurateWait(Constants.BUDGET_UPDATE_INTERVAL)do local f=tick()local B=(f-d)/60 d=f local T=#Players:GetPlayers()if not O then for G,v in pairs(ConstantsMapping)do l(G,v,B,T)end G[Enum.DataStoreRequestType.UpdateAsync]=math.min(G[Enum.DataStoreRequestType.GetAsync],G[Enum.DataStoreRequestType.SetIncrementAsync])end for G,v in pairs(v)do for G=#v,1,-1 do local d=v[G]local l=d.Thread local O=d.Budget local f=d.Key local B=d.Lock local T=d.Cache if not(B and(B[f]or tick()-(T[f]or 0)<Constants.WRITE_COOLDOWN))and N(O)then table.remove(v,G)C(O)coroutine.resume(l)end end end end end))()game:BindToClose(function()for v,d in pairs(ConstantsMapping)do G[v]=math.max(G[v],Constants.BUDGET_ONCLOSE_BASE*(d.RATE/Constants.BUDGET_BASE))end G[Enum.DataStoreRequestType.UpdateAsync]=math.min(G[Enum.DataStoreRequestType.GetAsync],G[Enum.DataStoreRequestType.SetIncrementAsync])end)end function MockDataStoreManager.GetGlobalData()return Data.GlobalDataStore end function MockDataStoreManager.GetData(G,v)assert(type(G)=="\115\116\114\105\110\103")assert(type(v)=="\115\116\114\105\110\103")if not Data.DataStore[G]then Data.DataStore[G]={}end if not Data.DataStore[G][v]then Data.DataStore[G][v]={}end return Data.DataStore[G][v]end function MockDataStoreManager.GetOrderedData(G,v)assert(type(G)=="\115\116\114\105\110\103")assert(type(v)=="\115\116\114\105\110\103")if not Data.OrderedDataStore[G]then Data.OrderedDataStore[G]={}end if not Data.OrderedDataStore[G][v]then Data.OrderedDataStore[G][v]={}end return Data.OrderedDataStore[G][v]end function MockDataStoreManager.GetDataInterface(G)return Interfaces[G]end function MockDataStoreManager.SetDataInterface(G,v)assert(type(G)=="\116\097\098\108\101")assert(type(v)=="\116\097\098\108\101")Interfaces[G]=v end function MockDataStoreManager.GetBudget(v)if Constants.BUDGETING_ENABLED then return math.floor(G[v]or 0)else return math.huge end end function MockDataStoreManager.SetBudget(v,d)assert(type(d)=="\110\117\109\098\101\114")d=math.max(d,0)if v==Enum.DataStoreRequestType.UpdateAsync then G[Enum.DataStoreRequestType.SetIncrementAsync]=d G[Enum.DataStoreRequestType.GetAsync]=d end if G[v]then G[v]=d end end function MockDataStoreManager.ResetBudget()d()end function MockDataStoreManager.FreezeBudgetUpdates()O=true end function MockDataStoreManager.ThawBudgetUpdates()O=false end function MockDataStoreManager.YieldForWriteLockAndBudget(G,d,l,C,N)assert(type(G)=="\102\117\110\099\116\105\111\110")assert(type(d)=="\115\116\114\105\110\103")assert(type(l)=="\116\097\098\108\101")assert(type(C)=="\116\097\098\108\101")assert(#N>0)local O=N[1]if#v[O]>=Constants.THROTTLE_QUEUE_SIZE then return false end G()table.insert(v[O],1,{Key=d;Lock=l;Cache=C,Thread=coroutine.running();Budget=N})coroutine.yield()return true end function MockDataStoreManager.YieldForBudget(G,d)assert(type(G)=="\102\117\110\099\116\105\111\110")assert(#d>0)local l=d[1]if N(d)then C(d)elseif#v[l]>=Constants.THROTTLE_QUEUE_SIZE then return false else G()table.insert(v[l],1,{After=0;Thread=coroutine.running(),Budget=d})coroutine.yield()end return true end function MockDataStoreManager.ExportToJSON()local G={}if next(Data.GlobalDataStore)~=nil then G.GlobalDataStore=Data.GlobalDataStore end G.DataStore=Utils.prepareDataStoresForExport(Data.DataStore)G.OrderedDataStore=Utils.prepareDataStoresForExport(Data.OrderedDataStore)return HttpService:JSONEncode(G)end

local Budgets = {}

local budgetRequestQueues = {
	[Enum.DataStoreRequestType.GetAsync] = {};
	[Enum.DataStoreRequestType.GetSortedAsync] = {};
	[Enum.DataStoreRequestType.OnUpdate] = {};
	[Enum.DataStoreRequestType.SetIncrementAsync] = {};
	[Enum.DataStoreRequestType.SetIncrementSortedAsync] = {};
}

local function initBudget()
	for requestType, const in pairs(ConstantsMapping) do
		Budgets[requestType] = const.START
	end
	Budgets[Enum.DataStoreRequestType.UpdateAsync] = math.min(
		Budgets[Enum.DataStoreRequestType.GetAsync],
		Budgets[Enum.DataStoreRequestType.SetIncrementAsync]
	)
end

local function updateBudget(req, const, dt, n)
	if not Constants.BUDGETING_ENABLED then
		return
	end
	local rate = const.RATE + n * const.RATE_PLR
	Budgets[req] = math.min(
		Budgets[req] + dt * rate,
		const.MAX_FACTOR * rate
	)
end

local function stealBudget(budget)
	if not Constants.BUDGETING_ENABLED then
		return
	end
	for _, requestType in pairs(budget) do
		if Budgets[requestType] then
			Budgets[requestType] = math.max(0, Budgets[requestType] - 1)
		end
	end
	Budgets[Enum.DataStoreRequestType.UpdateAsync] = math.min(
		Budgets[Enum.DataStoreRequestType.GetAsync],
		Budgets[Enum.DataStoreRequestType.SetIncrementAsync]
	)
end

local function checkBudget(budget)
	if not Constants.BUDGETING_ENABLED then
		return true
	end
	for _, requestType in pairs(budget) do
		if Budgets[requestType] and Budgets[requestType] < 1 then
			return false
		end
	end
	return true
end

local isFrozen = false

if RunService:IsServer() then
	-- Only do budget/throttle updating on server (in case package required on client)

	initBudget()

	coroutine.wrap(function() -- Thread that increases budgets and de-throttles requests periodically
		local lastCheck = tick()
		while Utils.accurateWait(Constants.BUDGET_UPDATE_INTERVAL) do
			local now = tick()
			local dt = (now - lastCheck) / 60
			lastCheck = now
			local n = #Players:GetPlayers()

			if not isFrozen then
				for requestType, const in pairs(ConstantsMapping) do
					updateBudget(requestType, const, dt, n)
				end
				Budgets[Enum.DataStoreRequestType.UpdateAsync] = math.min(
					Budgets[Enum.DataStoreRequestType.GetAsync],
					Budgets[Enum.DataStoreRequestType.SetIncrementAsync]
				)
			end

			for _, budgetRequestQueue in pairs(budgetRequestQueues) do
				for i = #budgetRequestQueue, 1, -1 do
					local request = budgetRequestQueue[i]

					local thread = request.Thread
					local budget = request.Budget
					local key = request.Key
					local lock = request.Lock
					local cache = request.Cache

					if not (lock and (lock[key] or tick() - (cache[key] or 0) < Constants.WRITE_COOLDOWN)) and checkBudget(budget) then
						table.remove(budgetRequestQueue, i)
						stealBudget(budget)
						coroutine.resume(thread)
					end
				end
			end
		end
	end)()

	game:BindToClose(function()
		for requestType, const in pairs(ConstantsMapping) do
			Budgets[requestType] = math.max(
				Budgets[requestType],
				Constants.BUDGET_ONCLOSE_BASE * (const.RATE / Constants.BUDGET_BASE)
			)
		end
		Budgets[Enum.DataStoreRequestType.UpdateAsync] = math.min(
			Budgets[Enum.DataStoreRequestType.GetAsync],
			Budgets[Enum.DataStoreRequestType.SetIncrementAsync]
		)
	end)

end

function MockDataStoreManager.GetGlobalData()
	return Data.GlobalDataStore
end

function MockDataStoreManager.GetData(name, scope)
	assert(type(name) == "string")
	assert(type(scope) == "string")

	if not Data.DataStore[name] then
		Data.DataStore[name] = {}
	end
	if not Data.DataStore[name][scope] then
		Data.DataStore[name][scope] = {}
	end

	return Data.DataStore[name][scope]
end

function MockDataStoreManager.GetOrderedData(name, scope)
	assert(type(name) == "string")
	assert(type(scope) == "string")

	if not Data.OrderedDataStore[name] then
		Data.OrderedDataStore[name] = {}
	end
	if not Data.OrderedDataStore[name][scope] then
		Data.OrderedDataStore[name][scope] = {}
	end

	return Data.OrderedDataStore[name][scope]
end

function MockDataStoreManager.GetDataInterface(data)
	return Interfaces[data]
end

function MockDataStoreManager.SetDataInterface(data, interface)
	assert(type(data) == "table")
	assert(type(interface) == "table")

	Interfaces[data] = interface
end

function MockDataStoreManager.GetBudget(requestType)
	if Constants.BUDGETING_ENABLED then
		return math.floor(Budgets[requestType] or 0)
	else
		return math.huge
	end
end

function MockDataStoreManager.SetBudget(requestType, budget)
	assert(type(budget) == "number")
	budget = math.max(budget, 0)

	if requestType == Enum.DataStoreRequestType.UpdateAsync then
		Budgets[Enum.DataStoreRequestType.SetIncrementAsync] = budget
		Budgets[Enum.DataStoreRequestType.GetAsync] = budget
	end

	if Budgets[requestType] then
		Budgets[requestType] = budget
	end
end

function MockDataStoreManager.ResetBudget()
	initBudget()
end

function MockDataStoreManager.FreezeBudgetUpdates()
	isFrozen = true
end

function MockDataStoreManager.ThawBudgetUpdates()
	isFrozen = false
end

function MockDataStoreManager.YieldForWriteLockAndBudget(callback, key, writeLock, writeCache, budget)
	assert(type(callback) == "function")
	assert(type(key) == "string")
	assert(type(writeLock) == "table")
	assert(type(writeCache) == "table")
	assert(#budget > 0)

	local mainRequestType = budget[1]

	if #budgetRequestQueues[mainRequestType] >= Constants.THROTTLE_QUEUE_SIZE then
		return false -- no room in throttle queue
	end

	callback() -- would i.e. trigger a warning in output

	table.insert(budgetRequestQueues[mainRequestType], 1, {
		Key = key;
		Lock = writeLock;
		Cache = writeCache;
		Thread = coroutine.running();
		Budget = budget;
	})
	coroutine.yield()

	return true
end

function MockDataStoreManager.YieldForBudget(callback, budget)
	assert(type(callback) == "function")
	assert(#budget > 0)

	local mainRequestType = budget[1]

	if checkBudget(budget) then
		stealBudget(budget)
	elseif #budgetRequestQueues[mainRequestType] >= Constants.THROTTLE_QUEUE_SIZE then
		return false -- no room in throttle queue
	else
		callback() -- would i.e. trigger a warning in output

		table.insert(budgetRequestQueues[mainRequestType], 1, {
			After = 0; -- no write lock
			Thread = coroutine.running();
			Budget = budget;
		})
		coroutine.yield()
	end

	return true
end

function MockDataStoreManager.ExportToJSON()
	local export = {}

	if next(Data.GlobalDataStore) ~= nil then -- GlobalDataStore not empty
		export.GlobalDataStore = Data.GlobalDataStore
	end
	export.DataStore = Utils.prepareDataStoresForExport(Data.DataStore) -- can be nil
	export.OrderedDataStore = Utils.prepareDataStoresForExport(Data.OrderedDataStore) -- can be nil

	return HttpService:JSONEncode(export)
end
