local Drag = {
	Dragging = false,
	Frame = nil,
	FramePosition = nil,
	FrameSize = nil, -- Added to store initial frame size
	StartPosition = nil,
	ChangedConnection = nil,
	Callback = nil,
}

local Helper = require(script.Parent.Helper)

local function DefaultCallback(_, Input: InputObject)
	local Delta = Input.Position - Drag.StartPosition
	local FramePosition: UDim2 = Drag.FramePosition

	if wax.shared.GetDPIScale then
		Delta = Delta / wax.shared.GetDPIScale()
	end

	Drag.Frame.Position = UDim2.new(
		FramePosition.X.Scale,
		FramePosition.X.Offset + Delta.X,
		FramePosition.Y.Scale,
		FramePosition.Y.Offset + Delta.Y
	)
end

function Drag.Setup(Frame: GuiObject, DragFrame: GuiObject, Callback: (Info: {}, Input: InputObject) -> () | nil)
	Callback = Callback or DefaultCallback

	DragFrame.InputBegan:Connect(function(Input: InputObject)
		if not Helper.IsClickInput(Input) then
			return
		end

		Drag.Dragging = true
		Drag.Frame = Frame
		Drag.FramePosition = Frame.Position
		Drag.StartPosition = Input.Position
		Drag.FrameSize = Frame.Size
		Drag.Callback = Callback

		Drag.ChangedConnection = Input.Changed:Connect(function()
			if Input.UserInputState ~= Enum.UserInputState.End then
				return
			end

			Drag.Dragging = false
			Drag.Frame = nil
			Drag.Callback = nil

			if Drag.ChangedConnection and Drag.ChangedConnection.Connected then
				Drag.ChangedConnection:Disconnect()
				Drag.ChangedConnection = nil
			end
		end)
	end)
end

wax.shared.Connect(wax.shared.UserInputService.InputChanged:Connect(function(Input: InputObject)
	if Drag.Dragging and Drag.Callback and Helper.IsMoveInput(Input) then
		Drag.Callback(Drag, Input)
	end
end))

return Drag


local UIHelper = {}

local CodeGen = require(script.Parent.Parent.CodeGen.Generator)

function UIHelper.IsMouseOverFrame(Frame: GuiObject, Position: Vector3): boolean
	local AbsPos, AbsSize = Frame.AbsolutePosition, Frame.AbsoluteSize
	return Position.X >= AbsPos.X
		and Position.X <= AbsPos.X + AbsSize.X
		and Position.Y >= AbsPos.Y
		and Position.Y <= AbsPos.Y + AbsSize.Y
end

function UIHelper.IsClickInput(Input: InputObject): boolean
	return Input.UserInputState == Enum.UserInputState.Begin
		and (Input.UserInputType == Enum.UserInputType.MouseButton1 or Input.UserInputType == Enum.UserInputType.Touch)
end

function UIHelper.IsMoveInput(Input: InputObject): boolean
	return Input.UserInputState == Enum.UserInputState.Change
		and (Input.UserInputType == Enum.UserInputType.MouseMovement or Input.UserInputType == Enum.UserInputType.Touch)
end

function UIHelper.QuickSerializeNumber(Number: number)
	if Number % 1 ~= 0 then
		return string.format("%.3f", Number)
	elseif Number == 1 / 0 then
		return "math.huge"
	elseif Number == -1 / 0 then
		return "-math.huge"
	end

	return Number
end

function UIHelper.QuickSerializeArgument(Argument)
	if typeof(Argument) == "string" then
		return string.format('"%s"', Argument)
	elseif typeof(Argument) == "number" then
		return UIHelper.QuickSerializeNumber(Argument)
	elseif typeof(Argument) == "Vector2" then
		return string.format(
			"%s, %s",
			UIHelper.QuickSerializeNumber(Argument.X),
			UIHelper.QuickSerializeNumber(Argument.Y)
		)
	elseif typeof(Argument) == "Vector3" then
		return string.format(
			"%s, %s, %s",
			UIHelper.QuickSerializeNumber(Argument.X),
			UIHelper.QuickSerializeNumber(Argument.Y),
			UIHelper.QuickSerializeNumber(Argument.Z)
		)
	elseif typeof(Argument) == "CFrame" then
		local Components = { Argument:GetComponents() }
		for Index, Value in pairs(Components) do
			Components[Index] = UIHelper.QuickSerializeNumber(Value)
		end
		return table.concat(Components, ", ")
	elseif typeof(Argument) == "table" then
		return "{...}"
	elseif typeof(Argument) == "Instance" then
		return CodeGen.GetFullPath(Argument, { DisableNilParentHandler = true })
	elseif typeof(Argument) == "userdata" then
		return "newproxy(" .. getmetatable(Argument) and "true" or "false" .. ")"
	end

	return tostring(Argument)
end

return UIHelper


local Resize = {}
Resize.__index = Resize

local Interface = require(script.Parent.Interface)
local Drag = require(script.Parent.Drag)

local HANDLE_SIZE = 6
local CORNER_HANDLE_SIZE = 20

function Resize.new(Options: {
	MainFrame: Frame,
	MinimumSize: Vector2? | UDim2?,
	MaximumSize: UDim2?,
	HandleSize: number?,
	CornerHandleSize: number?,
	Mirrored: boolean?,
	LockedPosition: boolean? | UDim2?,
})
	local MainFrame = Options.MainFrame
	local HandleSize = Options.HandleSize or HANDLE_SIZE
	local CornerHandleSize = Options.CornerHandleSize or CORNER_HANDLE_SIZE
	local Mirrored = Options.Mirrored or false
	local LockedPosition = Options.LockedPosition

	local MinimumSize
	if typeof(Options.MinimumSize) == "Vector2" then
		MinimumSize = UDim2.fromOffset(Options.MinimumSize.X, Options.MinimumSize.Y)
	elseif typeof(Options.MinimumSize) == "UDim2" then
		MinimumSize = Options.MinimumSize
	else
		MinimumSize = UDim2.fromOffset(100, 100)
	end

	local MaximumSize = Options.MaximumSize

	local self = setmetatable({
		MainFrame = MainFrame,
		ScreenGui = MainFrame:FindFirstAncestorOfClass("ScreenGui"),
		Parent = MainFrame.Parent,
	}, Resize)

	local function CalculateResizeProperties(
		initialFramePosition: UDim2,
		initialFrameSize: UDim2,
		mouseDelta: Vector2,
		resizeTypeX: string?,
		resizeTypeY: string?
	)
		if Mirrored then
			local parentAbsSize = self.Parent.AbsoluteSize

			local newSizeOffsetX = initialFrameSize.X.Offset
			local newSizeOffsetY = initialFrameSize.Y.Offset

			if resizeTypeX then
				local deltaX = 0
				if resizeTypeX == "Right" then
					deltaX = 2 * mouseDelta.X
				elseif resizeTypeX == "Left" then
					deltaX = -2 * mouseDelta.X
				end
				newSizeOffsetX = initialFrameSize.X.Offset + deltaX

				local minWidthAbs = MinimumSize.X.Scale * parentAbsSize.X + MinimumSize.X.Offset
				local maxWidthAbs = (MaximumSize and (MaximumSize.X.Scale * parentAbsSize.X + MaximumSize.X.Offset))
					or math.huge
				local scaleContributionX = initialFrameSize.X.Scale * parentAbsSize.X
				local minAllowedTotalOffsetX = minWidthAbs - scaleContributionX
				local maxAllowedTotalOffsetX = maxWidthAbs - scaleContributionX
				newSizeOffsetX = math.clamp(newSizeOffsetX, minAllowedTotalOffsetX, maxAllowedTotalOffsetX)
			end

			if resizeTypeY then
				local deltaY = 0
				if resizeTypeY == "Bottom" then
					deltaY = 2 * mouseDelta.Y
				elseif resizeTypeY == "Top" then
					deltaY = -2 * mouseDelta.Y
				end
				newSizeOffsetY = initialFrameSize.Y.Offset + deltaY

				local minHeightAbs = MinimumSize.Y.Scale * parentAbsSize.Y + MinimumSize.Y.Offset
				local maxHeightAbs = (MaximumSize and (MaximumSize.Y.Scale * parentAbsSize.Y + MaximumSize.Y.Offset))
					or math.huge
				local scaleContributionY = initialFrameSize.Y.Scale * parentAbsSize.Y
				local minAllowedTotalOffsetY = minHeightAbs - scaleContributionY
				local maxAllowedTotalOffsetY = maxHeightAbs - scaleContributionY
				newSizeOffsetY = math.clamp(newSizeOffsetY, minAllowedTotalOffsetY, maxAllowedTotalOffsetY)
			end

			local finalNewSize =
				UDim2.new(initialFrameSize.X.Scale, newSizeOffsetX, initialFrameSize.Y.Scale, newSizeOffsetY)
			local finalNewPosition = initialFramePosition
			if typeof(LockedPosition) == "UDim2" then
				finalNewPosition = LockedPosition
			end
			return finalNewSize, finalNewPosition
		else
			-- Non-mirrored logic
			local currentScreenGuiAbsSize = self.ScreenGui.AbsoluteSize
			local parentAbsSizeForMinMax = currentScreenGuiAbsSize -- As per original non-mirrored logic for min/max context

			-- These will store the final UDim offset values for position and the absolute pixel values for size calculation
			local finalPosOffsetX = initialFramePosition.X.Offset
			local finalPosOffsetY = initialFramePosition.Y.Offset

			-- Initial absolute pixel size of the frame
			local initialAbsWidthPx = initialFrameSize.X.Scale * self.Parent.AbsoluteSize.X + initialFrameSize.X.Offset
			local initialAbsHeightPx = initialFrameSize.Y.Scale * self.Parent.AbsoluteSize.Y + initialFrameSize.Y.Offset

			local newAbsWidthPx = initialAbsWidthPx
			local newAbsHeightPx = initialAbsHeightPx

			-- Min/max pixel dimensions
			local minWidthPx = MinimumSize.X.Scale * parentAbsSizeForMinMax.X + MinimumSize.X.Offset
			local minHeightPx = MinimumSize.Y.Scale * parentAbsSizeForMinMax.Y + MinimumSize.Y.Offset
			local maxWidthPx = MaximumSize and (MaximumSize.X.Scale * parentAbsSizeForMinMax.X + MaximumSize.X.Offset)
				or math.huge
			local maxHeightPx = MaximumSize and (MaximumSize.Y.Scale * parentAbsSizeForMinMax.Y + MaximumSize.Y.Offset)
				or math.huge

			-- Original edge calculation logic (assuming MainFrame.Position is center if AnchorPoint is 0.5,0.5 for these calcs)
			local initialAbsCenterX = currentScreenGuiAbsSize.X * initialFramePosition.X.Scale
				+ initialFramePosition.X.Offset
			local initialAbsSizeX_forEdgeCalc = initialFrameSize.X.Offset -- Original code used offset for this part of edge calculation
			local initialRightEdgeX = initialAbsCenterX + initialAbsSizeX_forEdgeCalc / 2
			local initialLeftEdgeX = initialAbsCenterX - initialAbsSizeX_forEdgeCalc / 2

			local initialAbsCenterY = currentScreenGuiAbsSize.Y * initialFramePosition.Y.Scale
				+ initialFramePosition.Y.Offset
			local initialAbsSizeY_forEdgeCalc = initialFrameSize.Y.Offset -- Original code used offset for this part of edge calculation
			local initialBottomEdgeY = initialAbsCenterY + initialAbsSizeY_forEdgeCalc / 2
			local initialTopEdgeY = initialAbsCenterY - initialAbsSizeY_forEdgeCalc / 2

			if resizeTypeX then
				if resizeTypeX == "Left" then
					local newLeftEdge = initialLeftEdgeX + mouseDelta.X
					newAbsWidthPx = math.clamp(initialRightEdgeX - newLeftEdge, minWidthPx, maxWidthPx)
					if newAbsWidthPx ~= (initialRightEdgeX - newLeftEdge) then -- Readjust edge if clamped
						newLeftEdge = initialRightEdgeX - newAbsWidthPx
					end
					if not LockedPosition then
						local newAbsCenterX = newLeftEdge + newAbsWidthPx / 2 -- Assuming center is halfway for position update
						finalPosOffsetX = newAbsCenterX - currentScreenGuiAbsSize.X * initialFramePosition.X.Scale
					end
				elseif resizeTypeX == "Right" then
					local newRightEdge = initialRightEdgeX + mouseDelta.X
					newAbsWidthPx = math.clamp(newRightEdge - initialLeftEdgeX, minWidthPx, maxWidthPx)
					if not LockedPosition then
						local newAbsCenterX = initialLeftEdgeX + newAbsWidthPx / 2 -- Assuming center is halfway
						finalPosOffsetX = newAbsCenterX - currentScreenGuiAbsSize.X * initialFramePosition.X.Scale
					end
				end
			end

			if resizeTypeY then
				if resizeTypeY == "Top" then
					local newTopEdge = initialTopEdgeY + mouseDelta.Y
					newAbsHeightPx = math.clamp(initialBottomEdgeY - newTopEdge, minHeightPx, maxHeightPx)
					if newAbsHeightPx ~= (initialBottomEdgeY - newTopEdge) then -- Readjust edge if clamped
						newTopEdge = initialBottomEdgeY - newAbsHeightPx
					end
					if not LockedPosition then
						local newAbsCenterY = newTopEdge + newAbsHeightPx / 2 -- Assuming center is halfway
						finalPosOffsetY = newAbsCenterY - currentScreenGuiAbsSize.Y * initialFramePosition.Y.Scale
					end
				elseif resizeTypeY == "Bottom" then
					local newBottomEdge = initialBottomEdgeY + mouseDelta.Y
					newAbsHeightPx = math.clamp(newBottomEdge - initialTopEdgeY, minHeightPx, maxHeightPx)
					if not LockedPosition then
						local newAbsCenterY = initialTopEdgeY + newAbsHeightPx / 2 -- Assuming center is halfway
						finalPosOffsetY = newAbsCenterY - currentScreenGuiAbsSize.Y * initialFramePosition.Y.Scale
					end
				end
			end

			-- Convert final absolute pixel dimensions back to UDim offsets for size
			local finalSizeOffsetX = newAbsWidthPx - (initialFrameSize.X.Scale * self.Parent.AbsoluteSize.X)
			local finalSizeOffsetY = newAbsHeightPx - (initialFrameSize.Y.Scale * self.Parent.AbsoluteSize.Y)

			local finalNewSize =
				UDim2.new(initialFrameSize.X.Scale, finalSizeOffsetX, initialFrameSize.Y.Scale, finalSizeOffsetY)
			local finalNewPosition = initialFramePosition -- Default if LockedPosition is true
			if typeof(LockedPosition) == "UDim2" then
				finalNewPosition = LockedPosition
			elseif not LockedPosition then -- Only update if not locked (boolean false)
				finalNewPosition = UDim2.new(
					initialFramePosition.X.Scale,
					finalPosOffsetX,
					initialFramePosition.Y.Scale,
					finalPosOffsetY
				)
			end
			return finalNewSize, finalNewPosition
		end
	end

	local function createDragHandler(resizeTypeX, resizeTypeY)
		return function(Info, Input: InputObject)
			local mouseDelta = Input.Position - Info.StartPosition
			if wax.shared.GetDPIScale then
				mouseDelta = mouseDelta / wax.shared.GetDPIScale()
			end

			local newSize, newPosition =
				CalculateResizeProperties(Info.FramePosition, Info.FrameSize, mouseDelta, resizeTypeX, resizeTypeY)

			MainFrame.Size = newSize
			MainFrame.Position = newPosition
		end
	end

	local LeftSide = Interface.New("Frame", {
		Size = UDim2.new(0, HandleSize, 1, -CornerHandleSize * 2),
		AnchorPoint = Vector2.new(0, 0.5),
		Position = UDim2.fromScale(0, 0.5),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6,
	})
	Drag.Setup(MainFrame, LeftSide, createDragHandler("Left", nil))

	local RightSide = Interface.New("Frame", {
		Size = UDim2.new(0, HandleSize, 1, -CornerHandleSize * 2),
		AnchorPoint = Vector2.new(1, 0.5),
		Position = UDim2.fromScale(1, 0.5),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6,
	})
	Drag.Setup(MainFrame, RightSide, createDragHandler("Right", nil))

	local TopSide = Interface.New("Frame", {
		Size = UDim2.new(1, -CornerHandleSize * 2, 0, HandleSize),
		AnchorPoint = Vector2.new(0.5, 0),
		Position = UDim2.fromScale(0.5, 0),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6,
	})
	Drag.Setup(MainFrame, TopSide, createDragHandler(nil, "Top"))

	local BottomSide = Interface.New("Frame", {
		Size = UDim2.new(1, -CornerHandleSize * 2, 0, HandleSize),
		AnchorPoint = Vector2.new(0.5, 1),
		Position = UDim2.fromScale(0.5, 1),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6,
	})
	Drag.Setup(MainFrame, BottomSide, createDragHandler(nil, "Bottom"))

	local TopLeftCorner = Interface.New("Frame", {
		Size = UDim2.fromOffset(CornerHandleSize, CornerHandleSize),
		AnchorPoint = Vector2.new(0, 0),
		Position = UDim2.fromScale(0, 0),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6 + 1,
	})
	Drag.Setup(MainFrame, TopLeftCorner, createDragHandler("Left", "Top"))

	local TopRightCorner = Interface.New("Frame", {
		Size = UDim2.fromOffset(CornerHandleSize, CornerHandleSize),
		AnchorPoint = Vector2.new(1, 0),
		Position = UDim2.fromScale(1, 0),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6 + 1,
	})
	Drag.Setup(MainFrame, TopRightCorner, createDragHandler("Right", "Top"))

	local BottomLeftCorner = Interface.New("Frame", {
		Size = UDim2.fromOffset(CornerHandleSize, CornerHandleSize),
		AnchorPoint = Vector2.new(0, 1),
		Position = UDim2.fromScale(0, 1),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6 + 1,
	})
	Drag.Setup(MainFrame, BottomLeftCorner, createDragHandler("Left", "Bottom"))

	local BottomRightCorner = Interface.New("Frame", {
		Size = UDim2.fromOffset(CornerHandleSize, CornerHandleSize),
		AnchorPoint = Vector2.new(1, 1),
		Position = UDim2.fromScale(1, 1),
		BackgroundTransparency = 1,
		Parent = MainFrame,
		ZIndex = 9e6 + 1,
	})
	Drag.Setup(MainFrame, BottomRightCorner, createDragHandler("Right", "Bottom"))

	return self
end

return Resize


local MarketplaceService = game:GetService("MarketplaceService")
local Players = game:GetService("Players")
local ReplicatedStorage = game:GetService("ReplicatedStorage")

local Lapis = require(ReplicatedStorage.Packages.Lapis)
local Sift = require(ReplicatedStorage.Packages.Sift)
local t = require(ReplicatedStorage.Packages.t)

local DEFAULT_DATA = { coins = 100, recentPurchases = {} }
local RECENT_PURCHASES_LIMIT = 100
local PRODUCTS = {
	[12345] = function(oldData)
		-- Product callbacks return a new version of the data.
		return Sift.Dictionary.merge(oldData, {
			coins = oldData.coins + 100,
		})
	end,
}

local collection = Lapis.createCollection("PlayerData", {
	defaultData = DEFAULT_DATA,
	validate = t.strictInterface({ coins = t.integer, recentPurchases = t.array(t.string) }),
})

local documents = {}

local function onPlayerAdded(player: Player)
	collection
		:load(`Player{player.UserId}`, { player.UserId })
		:andThen(function(document)
			if player.Parent == nil then
				document:close():catch(warn)
				return
			end

			documents[player] = document
		end)
		:catch(function(message)
			warn(`Player {player.Name}'s data failed to load: {message}`)
			player:Kick("Data failed to load.")
		end)
end

local function onPlayerRemoving(player: Player)
	local document = documents[player]

	if document ~= nil then
		documents[player] = nil
		document:close():catch(warn)
	end
end

local function processReceipt(receiptInfo): Enum.ProductPurchaseDecision
	local player = Players:GetPlayerByUserId(receiptInfo.PlayerId)
	if player == nil then
		return Enum.ProductPurchaseDecision.NotProcessedYet
	end

	while documents[player] == nil and player.Parent ~= nil do
		-- Wait until the document loads or the player leaves.
		task.wait()
	end

	local document = documents[player]
	if document == nil then
		return Enum.ProductPurchaseDecision.NotProcessedYet
	end

	local data = document:read()

	if table.find(data.recentPurchases, receiptInfo.PurchaseId) then
		-- The purchase has been added to the player's data, but it might not have saved yet.
		local saveOk = document:save():await()

		if saveOk then
			return Enum.ProductPurchaseDecision.PurchaseGranted
		else
			return Enum.ProductPurchaseDecision.NotProcessedYet
		end
	end

	-- The product callback must not yield. Otherwise, it can return outdated data and overwrite new changes.
	local productOk, dataWithProduct = pcall(PRODUCTS[receiptInfo.ProductId], data)
	if not productOk then
		return Enum.ProductPurchaseDecision.NotProcessedYet
	end

	local newRecentPurchases = Sift.Array.push(data.recentPurchases, receiptInfo.PurchaseId)
	if #newRecentPurchases > RECENT_PURCHASES_LIMIT then
		newRecentPurchases = Sift.Array.shift(newRecentPurchases, #newRecentPurchases - RECENT_PURCHASES_LIMIT)
	end

	document:write(Sift.Dictionary.merge(dataWithProduct, {
		recentPurchases = newRecentPurchases,
	}))

	local saveOk = document:save():await()
	if not saveOk then
		return Enum.ProductPurchaseDecision.NotProcessedYet
	end

	return Enum.ProductPurchaseDecision.PurchaseGranted
end

-- The ProcessReceipt callback must be set before the Players.PlayerAdded signal is fired, otherwise the player's
-- existing receipts won't be processed when they join.
MarketplaceService.ProcessReceipt = processReceipt

Players.PlayerAdded:Connect(onPlayerAdded)
Players.PlayerRemoving:Connect(onPlayerRemoving)

for _, player in Players:GetPlayers() do
	onPlayerAdded(player)
end


local Players = game:GetService("Players")
local ReplicatedStorage = game:GetService("ReplicatedStorage")

local Lapis = require(ReplicatedStorage.Packages.Lapis)
local t = require(ReplicatedStorage.Packages.t)

local DEFAULT_DATA = { coins = 100 }

local collection = Lapis.createCollection("PlayerData", {
	defaultData = DEFAULT_DATA,
	-- You can use t by osyrisrblx to type check your data at runtime.
	validate = t.strictInterface({ coins = t.integer }),
})

local documents = {}

local function onPlayerAdded(player: Player)
	-- The second argument associates the document with the player's UserId which is useful
	-- for GDPR compliance.
	collection
		:load(`Player{player.UserId}`, { player.UserId })
		:andThen(function(document)
			if player.Parent == nil then
				-- The player might have left before the document finished loading.
				-- The document needs to be closed because PlayerRemoving won't fire at this point.
				document:close():catch(warn)
				return
			end

			documents[player] = document
		end)
		:catch(function(message)
			warn(`Player {player.Name}'s data failed to load: {message}`)

			-- Optionally, you can kick the player when their data fails to load:
			player:Kick("Data failed to load.")
		end)
end

local function onPlayerRemoving(player: Player)
	local document = documents[player]

	-- The document won't be added to the dictionary if PlayerRemoving fires bofore it finishes loading.
	if document ~= nil then
		documents[player] = nil
		document:close():catch(warn)
	end
end

Players.PlayerAdded:Connect(onPlayerAdded)
Players.PlayerRemoving:Connect(onPlayerRemoving)

for _, player in Players:GetPlayers() do
	onPlayerAdded(player)
end


local RunService = game:GetService("RunService")

local Promise = require(script.Parent.Parent.Promise)

local UPDATE_INTERVAL = 5 * 60

local AutoSave = {}
AutoSave.__index = AutoSave

function AutoSave.new(data)
	return setmetatable({
		documents = {},
		data = data,
		gameClosed = false,
		ongoingLoads = 0,
		ongoingRemoveLocks = 0,
	}, AutoSave)
end

function AutoSave:addDocument(document)
	table.insert(self.documents, document)
end

function AutoSave:removeDocument(document)
	local index = table.find(self.documents, document)

	table.remove(self.documents, index)
end

function AutoSave:finishLoad(document)
	if self.gameClosed then
		document:close()
	end

	self.ongoingLoads -= 1
end

function AutoSave:onGameClose()
	self.gameClosed = true
	self.data.throttle.gameClosed = true

	while #self.documents > 0 do
		self.documents[#self.documents]:close()
	end

	Promise.try(function()
		while self.ongoingLoads > 0 do
			task.wait()
		end
	end)
		:andThen(function()
			return Promise.allSettled({
				self.data:waitForOngoingSaves(),
				Promise.try(function()
					while self.ongoingRemoveLocks > 0 do
						task.wait()
					end
				end),
			})
		end)
		:await()
end

function AutoSave:start()
	local nextUpdateAt = os.clock() + UPDATE_INTERVAL
	RunService.Heartbeat:Connect(function()
		if os.clock() >= nextUpdateAt then
			for _, document in self.documents do
				document:save():catch(warn)
			end

			nextUpdateAt += UPDATE_INTERVAL
		end
	end)

	game:BindToClose(function()
		self:onGameClose()
	end)
end

return AutoSave


local HttpService = game:GetService("HttpService")

local Document = require(script.Parent.Document)
local Error = require(script.Parent.Error)
local freezeDeep = require(script.Parent.freezeDeep)
local Migration = require(script.Parent.Migration)
local Promise = require(script.Parent.Parent.Promise)
local copyDeep = require(script.Parent.copyDeep)

local LOCK_EXPIRE = 30 * 60

--[=[
	Collections are analagous to [GlobalDataStore].

	@class Collection
]=]
local Collection = {}
Collection.__index = Collection

function Collection.new(name, options, data, autoSave, config)
	if typeof(options.defaultData) ~= "function" and options.validate ~= nil then
		assert(options.validate(options.defaultData))
	end

	local migrations = {}
	if options.migrations ~= nil then
		for _, migration in options.migrations do
			if typeof(migration) == "function" then
				table.insert(migrations, { migrate = migration })
			else
				table.insert(migrations, migration)
			end
		end
	end
	options.migrations = migrations

	options.freezeData = if options.freezeData ~= nil then options.freezeData else true

	freezeDeep(options)

	return setmetatable({
		dataStore = config:get("dataStoreService"):GetDataStore(name),
		options = options,
		data = data,
		autoSave = autoSave,
	}, Collection)
end

--[=[
	Loads the document with `key`, migrates it, and session locks it.

	If specified, the document's `DataStoreKeyInfo:GetUserIds()` will be set to `defaultUserIds` if the document has
	never been loaded.

	@param key string
	@param defaultUserIds {number}?
	@return Promise<Document>
]=]
function Collection:load(key, defaultUserIds)
	if self.autoSave.gameClosed then
		-- If game:BindToClose has been called, this infinitely yields so the document can't load.
		return Promise.new(function() end)
	end

	local lockId = HttpService:GenerateGUID(false)

	self.autoSave.ongoingLoads += 1

	return self.data
		:load(self.dataStore, key, function(value, keyInfo)
			if value == nil then
				local defaultData
				if typeof(self.options.defaultData) == "function" then
					local defaultOk, tailoredDefaultData = pcall(self.options.defaultData, key)
					if not defaultOk then
						return "fail",
							Error.new("DefaultDataThrew", `'defaultData' threw an error: {tailoredDefaultData}`)
					end

					if self.options.validate ~= nil then
						local validateOk, valid, message = pcall(self.options.validate, tailoredDefaultData)
						if not validateOk then
							return "fail", Error.new("ValidateThrew", `'validate' threw an error: {valid}`)
						elseif not valid then
							return "fail", Error.new("ValidateFailed", `Invalid data: {message}`)
						end
					end

					defaultData = copyDeep(tailoredDefaultData)
				else
					-- The data was validated when the collection was created.
					defaultData = if self.options.freezeData
						then self.options.defaultData
						else copyDeep(self.options.defaultData)
				end

				local data = {
					migrationVersion = #self.options.migrations,
					lastCompatibleVersion = Migration.getLastCompatibleVersion(self.options.migrations),
					lockId = lockId,
					data = defaultData,
				}

				return "succeed", data, defaultUserIds
			end

			if
				value.lockId ~= nil
				and (DateTime.now().UnixTimestampMillis - keyInfo.UpdatedTime) / 1000 < LOCK_EXPIRE
			then
				return "retry", Error.new("SessionLocked", "Could not acquire lock")
			end

			local savedVersion = value.migrationVersion

			local migrationOk, migrated, lastCompatibleVersion = Migration.migrate(self.options.migrations, value, key)
			if not migrationOk then
				return "fail", Error.new("MigrationError", migrated)
			end

			if self.options.validate ~= nil then
				local validateOk, valid, message = pcall(self.options.validate, migrated)
				if not validateOk then
					return "fail", Error.new("ValidateThrew", `'validate' threw an error: {valid}`)
				elseif not valid then
					return "fail", Error.new("ValidateFailed", `Invalid data: {message}`)
				end
			end

			local data = {
				migrationVersion = math.max(savedVersion, #self.options.migrations),
				lastCompatibleVersion = lastCompatibleVersion,
				lockId = lockId,
				data = migrated,
			}

			return "succeed", data, keyInfo:GetUserIds(), keyInfo:GetMetadata()
		end)
		:andThen(function(value, keyInfo)
			if value == "cancelled" then
				self.autoSave.ongoingLoads -= 1

				-- Infinitely yield because the load was cancelled by game:BindToClose.
				return Promise.new(function() end)
			end

			local data = value.data

			if self.options.freezeData then
				freezeDeep(data)
			end

			local document = Document.new(self, key, self.options.validate, lockId, data, keyInfo)

			self.autoSave:finishLoad(document)

			if self.autoSave.gameClosed then
				-- Infinitely yield because the document will automatically be closed.
				return Promise.new(function() end)
			end

			self.autoSave:addDocument(document)

			return document
		end)
		:catch(function(err)
			self.autoSave.ongoingLoads -= 1

			return Promise.reject(`DataStoreFailure({err.message})`)
		end)
end

--[=[
	Reads the data of the document with `key` regardless of whether it is session locked. This is useful for viewing a
	document without editing or session locking it. The data gets migrated but not saved.

	If the document has never been loaded, the promise will return `nil`.

	[DataStoreGetOptions.UseCache](https://create.roblox.com/docs/reference/engine/classes/DataStoreGetOptions#UseCache) is disabled.

	@param key string
	@return Promise<T?>
]=]
function Collection:read(key)
	return self.data:read(self.dataStore, key):andThen(function(value, keyInfo)
		if value == nil then
			return nil
		end

		local migrationOk, migrated = Migration.migrate(self.options.migrations, value, key)
		if not migrationOk then
			return Promise.reject(migrated)
		end

		if self.options.validate ~= nil then
			local validateOk, valid, message = pcall(self.options.validate, migrated)
			if not validateOk then
				return Promise.reject(`'validate' threw an error: {valid}`)
			elseif not valid then
				return Promise.reject(`Invalid data: {message}`)
			end
		end

		return value.data, keyInfo
	end)
end

--[=[
	Removes the data of the document with `key`.

	If a document is open while `remove` is called, the open document will fail to save/close and the data will still be removed.

	@param key string
	@return Promise<()>
]=]
function Collection:remove(key)
	return self.data:remove(self.dataStore, key)
end

return Collection


local DataStoreService = game:GetService("DataStoreService")

local Config = {}
Config.__index = Config

function Config.new()
	return setmetatable({
		config = {
			saveAttempts = 5,
			loadAttempts = 20,
			loadRetryDelay = 1,
			showRetryWarnings = true,
			dataStoreService = DataStoreService,
		},
	}, Config)
end

function Config:get(key)
	return self.config[key]
end

function Config:set(values)
	for key, value in values do
		if self.config[key] == nil then
			error(`Invalid config key "{tostring(key)}"`)
		end

		self.config[key] = value
	end
end

return Config


local RunService = game:GetService("RunService")

local Error = require(script.Parent.Parent.Error)
local Promise = require(script.Parent.Parent.Parent.Promise)

local GET_ASYNC_RETRY_ATTEMPTS = 5
local GET_ASYNC_RETRY_DELAY = 1

local REMOVE_ASYNC_RETRY_ATTEMPTS = 5
local REMOVE_ASYNC_RETRY_DELAY = 1

local getAsyncOptions = Instance.new("DataStoreGetOptions")
getAsyncOptions.UseCache = false

local function updateAsync(throttle, request)
	return Promise.new(function(resolve)
		local resultOutside, transformedOutside, keyInfo
		local ok, err = pcall(function()
			_, keyInfo = request.dataStore:UpdateAsync(request.key, function(...)
				if request.cancelOnGameClose and throttle.gameClosed then
					resultOutside = "cancelled"
					return nil
				end

				local result, transformed, userIds, metadata = request.transform(...)

				resultOutside = result
				transformedOutside = transformed

				if result == "succeed" then
					return transformed, userIds, metadata
				else
					return nil
				end
			end)
		end)

		if resultOutside == "cancelled" then
			resolve("cancelled")
		elseif not ok then
			resolve("retry", Error.new("RobloxApiError", err))
		else
			resolve(resultOutside, transformedOutside, keyInfo)
		end
	end)
end

local function getAsync(request)
	return Promise.new(function(resolve)
		local ok, value, keyInfo = pcall(function()
			return request.dataStore:GetAsync(request.key, getAsyncOptions)
		end)

		resolve(ok, value, keyInfo)
	end)
end

local function removeAsync(request)
	return Promise.new(function(resolve)
		local ok, message = pcall(function()
			request.dataStore:RemoveAsync(request.key)
		end)

		resolve(ok, message)
	end)
end

local Throttle = {}
Throttle.__index = Throttle

function Throttle.new(config)
	return setmetatable({
		config = config,
		updateAsyncQueue = {},
		getAsyncQueue = {},
		removeAsyncQueue = {},
		gameClosed = false,
	}, Throttle)
end

function Throttle:getUpdateAsyncBudget()
	return self.config:get("dataStoreService"):GetRequestBudgetForRequestType(Enum.DataStoreRequestType.UpdateAsync)
end

function Throttle:getGetAsyncBudget()
	return self.config:get("dataStoreService"):GetRequestBudgetForRequestType(Enum.DataStoreRequestType.GetAsync)
end

function Throttle:getRemoveAsyncBudget()
	return self.config
		:get("dataStoreService")
		:GetRequestBudgetForRequestType(Enum.DataStoreRequestType.SetIncrementAsync)
end

function Throttle:start()
	local function retryRequest(request, err)
		request.attempts -= 1

		if request.attempts == 0 then
			request.reject(err)
		else
			if self.config:get("showRetryWarnings") then
				warn(`DataStore operation failed. Retrying...\nError: {err.message}`)
			end

			task.wait(request.retryDelay)
		end
	end

	local function updateUpdateAsync()
		for index = #self.updateAsyncQueue, 1, -1 do
			local request = self.updateAsyncQueue[index]

			if request.attempts == 0 then
				table.remove(self.updateAsyncQueue, index)
			elseif request.promise == nil and request.cancelOnGameClose and self.gameClosed then
				request.resolve("cancelled")
				table.remove(self.updateAsyncQueue, index)
			end
		end

		for _, request in self.updateAsyncQueue do
			if self:getUpdateAsyncBudget() == 0 then
				break
			end

			if request.promise ~= nil then
				continue
			end

			local promise = updateAsync(self, request):andThen(function(result, value, keyInfo)
				if result == "cancelled" then
					request.attempts = 0
					request.resolve("cancelled")
				elseif result == "succeed" then
					request.attempts = 0
					request.resolve(value, keyInfo)
				elseif result == "fail" then
					request.attempts = 0
					request.reject(value)
				elseif result == "retry" then
					retryRequest(request, value)
				else
					error("unreachable")
				end

				request.promise = nil
			end)

			if promise:getStatus() == Promise.Status.Started then
				request.promise = promise
			end
		end
	end

	local function updateGetAsync()
		for index = #self.getAsyncQueue, 1, -1 do
			local request = self.getAsyncQueue[index]

			if request.attempts == 0 then
				table.remove(self.getAsyncQueue, index)
			end
		end

		for _, request in self.getAsyncQueue do
			if self:getGetAsyncBudget() == 0 then
				break
			end

			if request.promise ~= nil then
				continue
			end

			local promise = getAsync(request):andThen(function(ok, value, keyInfo)
				if ok then
					request.attempts = 0
					request.resolve(value, keyInfo)
				else
					retryRequest(request, value)
				end

				request.promise = nil
			end)

			if promise:getStatus() == Promise.Status.Started then
				request.promise = promise
			end
		end
	end

	local function updateRemoveAsync()
		for index = #self.removeAsyncQueue, 1, -1 do
			local request = self.removeAsyncQueue[index]

			if request.attempts == 0 then
				table.remove(self.removeAsyncQueue, index)
			end
		end

		for _, request in self.removeAsyncQueue do
			if self:getRemoveAsyncBudget() == 0 then
				break
			end

			if request.promise ~= nil then
				continue
			end

			local promise = removeAsync(request):andThen(function(ok, message)
				if ok then
					request.attempts = 0
					request.resolve()
				else
					retryRequest(request, message)
				end

				request.promise = nil
			end)

			if promise:getStatus() == Promise.Status.Started then
				request.promise = promise
			end
		end
	end

	RunService.PostSimulation:Connect(function()
		updateUpdateAsync()
		updateGetAsync()
		updateRemoveAsync()
	end)
end

function Throttle:updateAsync(dataStore, key, transform, cancelOnGameClose, retryAttempts, retryDelay)
	return Promise.new(function(resolve, reject)
		table.insert(self.updateAsyncQueue, {
			dataStore = dataStore,
			key = key,
			transform = transform,
			attempts = retryAttempts,
			retryDelay = retryDelay,
			cancelOnGameClose = cancelOnGameClose,
			resolve = resolve,
			reject = reject,
		})
	end)
end

function Throttle:getAsync(dataStore, key)
	return Promise.new(function(resolve, reject)
		table.insert(self.getAsyncQueue, {
			dataStore = dataStore,
			key = key,
			attempts = GET_ASYNC_RETRY_ATTEMPTS,
			retryDelay = GET_ASYNC_RETRY_DELAY,
			resolve = resolve,
			reject = reject,
		})
	end)
end

function Throttle:removeAsync(dataStore, key)
	return Promise.new(function(resolve, reject)
		table.insert(self.removeAsyncQueue, {
			dataStore = dataStore,
			key = key,
			attempts = REMOVE_ASYNC_RETRY_ATTEMPTS,
			retryDelay = REMOVE_ASYNC_RETRY_DELAY,
			resolve = resolve,
			reject = reject,
		})
	end)
end

return Throttle


local Error = require(script.Parent.Error)
local Promise = require(script.Parent.Parent.Promise)
local Throttle = require(script.Throttle)

local Data = {}
Data.__index = Data

function Data.new(config)
	local throttle = Throttle.new(config)

	throttle:start()

	return setmetatable({
		config = config,
		throttle = throttle,
		ongoingSaves = {},
	}, Data)
end

function Data:waitForOngoingSave(dataStore, key)
	if self.ongoingSaves[dataStore] == nil or self.ongoingSaves[dataStore][key] == nil then
		return Promise.resolve()
	end

	local ongoingSave = self.ongoingSaves[dataStore][key]

	return Promise.allSettled({
		ongoingSave.promise,
		if ongoingSave.pendingSave ~= nil then ongoingSave.pendingSave.promise else nil,
	})
end

function Data:waitForOngoingSaves()
	local promises = {}

	for _, ongoingSaves in self.ongoingSaves do
		for _, ongoingSave in ongoingSaves do
			if ongoingSave.pendingSave ~= nil then
				table.insert(promises, ongoingSave.pendingSave.promise)
			end

			table.insert(promises, ongoingSave.promise)
		end
	end

	return Promise.allSettled(promises)
end

function Data:read(dataStore, key)
	return self.throttle:getAsync(dataStore, key)
end

function Data:load(dataStore, key, transform)
	return self:waitForOngoingSave(dataStore, key):andThen(function()
		local attempts = self.config:get("loadAttempts")
		local retryDelay = self.config:get("loadRetryDelay")

		return self.throttle:updateAsync(dataStore, key, transform, true, attempts, retryDelay)
	end)
end

function Data:save(dataStore, key, transform)
	if self.ongoingSaves[dataStore] == nil then
		self.ongoingSaves[dataStore] = {}
	end

	local ongoingSave = self.ongoingSaves[dataStore][key]

	if ongoingSave == nil then
		local attempts = self.config:get("saveAttempts")
		local promise = self.throttle:updateAsync(dataStore, key, transform, false, attempts):finally(function()
			self.ongoingSaves[dataStore][key] = nil

			if next(self.ongoingSaves[dataStore]) == nil then
				self.ongoingSaves[dataStore] = nil
			end
		end)

		if promise:getStatus() == Promise.Status.Started then
			self.ongoingSaves[dataStore][key] = { promise = promise }
		end

		return promise
	elseif ongoingSave.pendingSave == nil then
		local pendingSave = { transform = transform }

		local function save()
			return self:save(dataStore, key, pendingSave.transform)
		end

		-- promise:finally(save) can't be used because if the ongoingSave promise rejects, so will the promise returned from finally.
		pendingSave.promise = ongoingSave.promise:andThen(save, save)

		ongoingSave.pendingSave = pendingSave

		return pendingSave.promise
	else
		ongoingSave.pendingSave.transform = transform

		return ongoingSave.pendingSave.promise
	end
end

function Data:removeLock(dataStore, key, lockIdToRemove)
	local function transform(value, keyInfo)
		if value == nil then
			return "fail", Error.new("DocumentRemoved", "The document was removed")
		end

		if value.lockId ~= lockIdToRemove then
			return "fail", Error.new("SessionLockStolen", "The session lock was stolen")
		end

		value.lockId = nil

		return "succeed", value, keyInfo:GetUserIds(), keyInfo:GetMetadata()
	end

	local attempts = self.config:get("saveAttempts")

	return self.throttle:updateAsync(dataStore, key, transform, false, attempts)
end

function Data:remove(dataStore, key)
	return self.throttle:removeAsync(dataStore, key)
end

return Data


local Error = require(script.Parent.Error)
local freezeDeep = require(script.Parent.freezeDeep)
local Promise = require(script.Parent.Parent.Promise)
local noYield = require(script.Parent.noYield)

local function runCallback(document, name, callback)
	if callback == nil then
		return Promise.resolve()
	end

	document.callingCallback = name

	return Promise.new(function(resolve, reject)
		local ok, message = pcall(noYield, callback)

		document.callingCallback = nil

		if not ok then
			reject(Error.new("BeforeSaveCloseCallbackThrew", `{name} callback threw error: {message}`))
		else
			resolve()
		end
	end)
end

--[=[
	@class Document
]=]
local Document = {}
Document.__index = Document

function Document.new(collection, key, validate, lockId, data, keyInfo)
	return setmetatable({
		collection = collection,
		key = key,
		validate = validate,
		lockId = lockId,
		data = data,
		userIds = keyInfo:GetUserIds(),
		lastKeyInfo = keyInfo,
		closed = false,
	}, Document)
end

--[=[
	Returns the document's data.

	@return any
]=]
function Document:read()
	return self.data
end

--[=[
	Updates the document's cached data. This method doesn't save the data to the DataStore; it only modifies the
	document's in-memory data.

	This method should be used when performing immutable updates to the document's data. For mutable updates, the data
	can be directly modified:
	```lua
	local data = document:read()
	data.coins += 100
	```

	:::warning
	Throws an error if the document was closed or if the data is invalid.
	:::

	@param data any
]=]
function Document:write(data)
	assert(not self.closed, "Cannot write to a closed document")

	if self.validate ~= nil then
		assert(self.validate(data))
	end

	if self.collection.options.freezeData then
		freezeDeep(data)
	end

	self.data = data
end

--[=[
	Adds a user id to the document's `DataStoreKeyInfo:GetUserIds()`. The change won't apply until the document is
	saved or closed.

	If the user id is already associated with the document the method won't do anything.

	@param userId number
]=]
function Document:addUserId(userId)
	assert(not self.closed, "Cannot add user id to a closed document")

	if table.find(self.userIds, userId) == nil then
		table.insert(self.userIds, userId)
	end
end

--[=[
	Removes a user id from the document's `DataStoreKeyInfo:GetUserIds()`. The change won't apply until the document is
	saved or closed.

	If the user id is not associated with the document the method won't do anything.

	@param userId number
]=]
function Document:removeUserId(userId)
	assert(not self.closed, "Cannot remove user id from a closed document")

	local index = table.find(self.userIds, userId)

	if index ~= nil then
		table.remove(self.userIds, index)
	end
end

--[=[
	Returns the last updated `DataStoreKeyInfo` returned from loading, saving, or closing the document.

	@return DataStoreKeyInfo
]=]
function Document:keyInfo()
	return self.lastKeyInfo
end

--[=[
	Saves the document's data. If the save is throttled and you call it multiple times, it will save only once with the
	latest data.

	Documents are saved automatically. This method is used mainly to handle developer product purchases
	(see the [example](../docs/DeveloperProduct)) or other situations requiring immediate saving.

	:::warning
	Throws an error if the document was closed.
	:::

	:::warning
	If the beforeSave callback yields or errors, the returned promise will reject and the data will not be saved.
	:::

	@return Promise<()>
]=]
function Document:save()
	assert(not self.closed, "Cannot save a closed document")
	assert(self.callingCallback == nil, `Cannot save in {self.callingCallback} callback`)

	return runCallback(self, "beforeSave", self.beforeSaveCallback)
		:andThen(function()
			return self.collection.data
				:save(self.collection.dataStore, self.key, function(value, keyInfo)
					if value == nil then
						return "fail", Error.new("DocumentRemoved", "The document was removed")
					end

					if value.lockId ~= self.lockId then
						return "fail", Error.new("SessionLockStolen", "The session lock was stolen")
					end

					if not self.collection.options.freezeData and self.validate ~= nil then
						local validateOk, valid, message = pcall(self.validate, self.data)
						if not validateOk then
							return "fail", Error.new("ValidateThrew", `'validate' threw an error: {valid}`)
						elseif not valid then
							return "fail", Error.new("ValidateFailed", `Invalid data: {message}`)
						end
					end

					value.data = self.data

					return "succeed", value, self.userIds, keyInfo:GetMetadata()
				end)
				:andThen(function(_, keyInfo)
					self.lastKeyInfo = keyInfo
				end)
		end)
		:catch(function(err)
			return Promise.reject(`DataStoreFailure({err.message})`)
		end)
end

--[=[
	Saves the document and removes the session lock. The document is unusable after calling. If a save is currently in
	progress it will close the document instead.

	If called again, it will return the promise from the original call.

	:::warning
	If the beforeSave or beforeClose callbacks yield or error, the returned promise will reject and the data will not be saved.
	:::

	@return Promise<()>
]=]
function Document:close()
	assert(self.callingCallback == nil, `Cannot close in {self.callingCallback} callback`)

	if self.closePromise == nil then
		self.closePromise = runCallback(self, "beforeSave", self.beforeSaveCallback)
			:andThenCall(runCallback, self, "beforeClose", self.beforeCloseCallback)
			:finally(function()
				self.closed = true

				self.collection.autoSave:removeDocument(self)
			end)
			:andThen(function()
				return self.collection.data:save(self.collection.dataStore, self.key, function(value, keyInfo)
					if value == nil then
						return "fail", Error.new("DocumentRemoved", "The document was removed")
					end

					if value.lockId ~= self.lockId then
						return "fail", Error.new("SessionLockStolen", "The session lock was stolen")
					end

					if not self.collection.options.freezeData and self.validate ~= nil then
						local validateOk, valid, message = pcall(self.validate, self.data)
						if not validateOk then
							return "fail", Error.new("ValidateThrew", `'validate' threw an error: {valid}`)
						elseif not valid then
							return "fail", Error.new("ValidateFailed", `Invalid data: {message}`)
						end
					end

					value.data = self.data
					value.lockId = nil

					return "succeed", value, self.userIds, keyInfo:GetMetadata()
				end)
			end)
			:andThen(function(_, keyInfo)
				self.lastKeyInfo = keyInfo
			end)
			:catch(function(err)
				if
					err.kind == "BeforeSaveCloseCallbackThrew"
					or err.kind == "ValidateThrew"
					or err.kind == "ValidateFailed"
				then
					self.collection.autoSave.ongoingRemoveLocks += 1

					self.collection.data
						:removeLock(self.collection.dataStore, self.key, self.lockId)
						:catch(function(removeLockErr)
							warn(`RemoveLockFailure({removeLockErr.message})`)
						end)
						:finally(function()
							self.collection.autoSave.ongoingRemoveLocks -= 1
						end)
				end

				return Promise.reject(`DataStoreFailure({err.message})`)
			end)
	end

	return self.closePromise
end

--[=[
	Sets a callback that is run inside `document:save` and `document:close` before it saves. The document can be read and written to in the
	callback.

	The callback will run before the beforeClose callback inside of `document:close`.

	:::warning
	Throws an error if it was called previously.
	:::

	@param callback () -> ()
]=]
function Document:beforeSave(callback)
	assert(self.beforeSaveCallback == nil, "Document:beforeSave can only be called once")

	self.beforeSaveCallback = callback
end

--[=[
	Sets a callback that is run inside `document:close` before it saves. The document can be read and written to in the
	callback.

	:::warning
	Throws an error if it was called previously.
	:::

	@param callback () -> ()
]=]
function Document:beforeClose(callback)
	assert(self.beforeCloseCallback == nil, "Document:beforeClose can only be called once")

	self.beforeCloseCallback = callback
end

return Document


local ReplicatedStorage = game:GetService("ReplicatedStorage")

local Promise = require(ReplicatedStorage.Packages.Promise)

local function defaultOptions()
	return {
		validate = function(data)
			return typeof(data.foo) == "string", "foo must be a string"
		end,
		defaultData = { foo = "bar" },
	}
end

return function(x)
	local assertEqual = x.assertEqual
	local shouldThrow = x.shouldThrow

	x.test("it should not merge close into save when save is running", function(context)
		local document = context.lapis.createCollection("collection", defaultOptions()):load("doc"):expect()

		-- It's not safe to merge saves when UpdateAsync is running.
		-- This will yield the UpdateAsync call until stopYield is called.
		context.dataStoreService.yield:startYield()

		local save = document:save()
		document:write({ foo = "new" })
		local close = document:close()

		context.dataStoreService.yield:stopYield()

		Promise.all({ save, close }):expect()

		local saved = context.read("collection", "doc")

		-- If data.foo == "bar", that means the close was merged with the save when it wasn't safe to.
		assert(saved.data.foo == "new", "")
	end)

	x.test("it should merge pending saves", function(context)
		local document = context.lapis.createCollection("collection", defaultOptions()):load("doc"):expect()

		context.dataStoreService.yield:startYield()

		local ongoingSave = document:save()

		local pendingSave = document:save()
		local pendingClose = document:close() -- This should override the pending save.

		context.dataStoreService.yield:stopYield()

		local values = Promise.all({ ongoingSave, pendingSave }):expect()

		-- This stops the close if it wasn't merged.
		context.dataStoreService.yield:startYield()

		-- Since the following code is resumed by the save promise, we need to wait for the close promise to resolve.
		task.wait()

		pendingClose:now("save and close didn't merge"):expect()

		-- save and close should never resolve with a value.
		-- It's checked in this test to make sure it works with save merging.
		assert(#values == 0, "")

		local saved = context.read("collection", "doc")

		assert(saved.lockId == nil, "")
	end)

	x.test("saves data", function(context)
		local document = context.lapis.createCollection("12345", defaultOptions()):load("doc"):expect()

		document:write({
			foo = "new value",
		})

		document:save():expect()

		local saved = context.read("12345", "doc")

		assert(typeof(saved) == "table", "")
		assert(typeof(saved.lockId) == "string", "")
		assert(saved.data.foo == "new value", "")
	end)

	x.test("writes the data", function(context)
		local document = context.lapis.createCollection("1", defaultOptions()):load("doc"):expect()

		document:write({
			foo = "baz",
		})

		assert(document:read().foo == "baz", "")
	end)

	x.test("write throws if data doesn't validate", function(context)
		local document = context.lapis.createCollection("2", defaultOptions()):load("doc"):expect()

		shouldThrow(function()
			document:write({
				foo = 5,
			})
		end, "foo must be a string")
	end)

	x.test("methods throw when called on a closed document", function(context)
		local document = context.lapis.createCollection("5", defaultOptions()):load("doc"):expect()

		local promise = document:close()

		shouldThrow(function()
			document:write({})
		end, "Cannot write to a closed document")

		shouldThrow(function()
			document:save()
		end, "Cannot save a closed document")

		shouldThrow(function()
			document:addUserId(1234)
		end, "Cannot add user id to a closed document")

		shouldThrow(function()
			document:removeUserId(1234)
		end, "Cannot remove user id from a closed document")

		promise:expect()
	end)

	x.test("close returns first promise when called again", function(context)
		local document = context.lapis.createCollection("col", defaultOptions()):load("doc"):expect()

		local promise = document:close()

		assertEqual(promise, document:close())
	end)

	x.test("loads with default data", function(context)
		local document = context.lapis.createCollection("o", defaultOptions()):load("a"):expect()

		assert(document:read().foo == "bar", "")
	end)

	x.test("loads with existing data", function(context)
		local collection = context.lapis.createCollection("xyz", defaultOptions())

		context.write("xyz", "xyz", {
			foo = "existing",
		})

		local document = collection:load("xyz"):expect()

		assert(document:read().foo == "existing", "")
	end)

	x.test("freezes document data", function(context)
		local collection = context.lapis.createCollection("collection", {
			defaultData = {},
		})

		context.write("collection", "document", { a = { b = 1 } })

		local document = collection:load("document"):expect()

		shouldThrow(function()
			document:read().a.b = 2
		end)

		document:write({ a = { b = 2 } })

		shouldThrow(function()
			document:read().a.b = 3
		end)
	end)

	x.test("doesn't save data when the lock was stolen", function(context)
		local collection = context.lapis.createCollection("hi", defaultOptions())

		local document = collection:load("hi"):expect()

		context.write("hi", "hi", {
			foo = "stolen",
		}, "stolenLockId")

		document:write({
			foo = "qux",
		})

		shouldThrow(function()
			document:save():expect()
		end, "The session lock was stolen")

		assert(context.read("hi", "hi").data.foo == "stolen", "")

		shouldThrow(function()
			document:close():expect()
		end, "The session lock was stolen")

		assert(context.read("hi", "hi").data.foo == "stolen", "")
	end)

	x.test("doesn't throw when the budget is exhausted", function(context)
		-- This makes sure the test doesn't pass by retyring after budget is added.
		context.lapis.setConfig({ loadAttempts = 1 })

		local document = context.lapis.createCollection("bye", defaultOptions()):load("bye"):expect()

		context.dataStoreService.budget.budgets[Enum.DataStoreRequestType.GetAsync] = 0
		context.dataStoreService.budget.budgets[Enum.DataStoreRequestType.SetIncrementAsync] = 0
		context.dataStoreService.budget.budgets[Enum.DataStoreRequestType.UpdateAsync] = 0

		local promise = document:save()

		-- This wait is necessary so that the request is run by Throttle.
		task.wait(0.1)

		context.dataStoreService.budget:update()

		promise:expect()
	end)

	x.test(":save doesn't resolve with any value", function(context)
		local document = context.lapis.createCollection("12345", defaultOptions()):load("doc"):expect()

		local a, b = document:save():expect()

		assert(a == nil, "")
		assert(b == nil, "")
	end)

	x.test(":close doesn't resolve with any value", function(context)
		local document = context.lapis.createCollection("12345", defaultOptions()):load("doc"):expect()

		local a, b = document:close():expect()

		assert(a == nil, "")
		assert(b == nil, "")
	end)

	x.nested("Document:beforeSave", function()
		x.test("throws when yielding", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function()
				task.wait()
			end)

			shouldThrow(function()
				document:save():expect()
			end, "beforeSave callback threw error: thread is not yieldable")
		end)

		x.test("throws when setting twice", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function() end)

			shouldThrow(function()
				document:beforeSave(function() end)
			end, "Document:beforeSave can only be called once")
		end)

		x.test("throws when calling close in callback", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function()
				document:close()
			end)

			shouldThrow(function()
				document:close():expect()
			end, "beforeSave callback threw error")
		end)

		x.test("throws when calling save in callback", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function()
				document:save()
			end)

			shouldThrow(function()
				document:close():expect()
			end, "beforeSave callback threw error")
		end)

		x.test("saves new data in document:save", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function()
				document:read() -- This checks that read doesn't error in the callback.
				document:write({ foo = "new" })
			end)

			document:save():expect()

			assertEqual(context.read("collection", "document").data.foo, "new")
		end)

		x.test("saves new data in document:close", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeSave(function()
				document:write({ foo = "new" })
			end)

			document:close():expect()

			assertEqual(context.read("collection", "document").data.foo, "new")
		end)
	end)

	x.nested("Document:beforeClose", function()
		x.test("throws when yielding", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeClose(function()
				task.wait()
			end)

			shouldThrow(function()
				document:close():expect()
			end, "beforeClose callback threw error: thread is not yieldable")
		end)

		x.test("throws when setting twice", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeClose(function() end)

			shouldThrow(function()
				document:beforeClose(function() end)
			end, "Document:beforeClose can only be called once")
		end)

		x.test("throws when calling close in callback", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeClose(function()
				document:close()
			end)

			shouldThrow(function()
				document:close():expect()
			end, "beforeClose callback threw error")
		end)

		x.test("throws when calling save in callback", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeClose(function()
				document:save()
			end)

			shouldThrow(function()
				document:close():expect()
			end, "beforeClose callback threw error")
		end)

		x.test("closes document even if beforeClose errors", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local promise = collection:load("document")
			local document = promise:expect()

			document:beforeClose(function()
				error("error")
			end)

			shouldThrow(function()
				document:close():expect()
			end)

			local secondPromise = collection:load("document")

			assert(secondPromise ~= promise, "collection:load should return a new promise")

			shouldThrow(function()
				document:write({ foo = "baz" })
			end, "Cannot write to a closed document")

			-- Ignore the could not acquire lock error.
			secondPromise:catch(function() end)
		end)

		x.test("saves new data", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			document:beforeClose(function()
				document:read() -- This checks that read doesn't error in the callback.

				document:write({ foo = "new" })
			end)

			document:close():expect()

			assertEqual(context.read("collection", "document").data.foo, "new")
		end)

		x.test("beforeSave runs before beforeClose", function(context)
			local document = context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()

			local order = ""

			document:beforeSave(function()
				order ..= "s"
			end)

			document:beforeClose(function()
				order ..= "c"
			end)

			document:close():expect()

			assertEqual(order, "sc")
		end)

		x.nested("keyInfo", function()
			x.test("gets load key info", function(context)
				local collection = context.lapis.createCollection("collection", defaultOptions())

				local before = context.getKeyInfo("collection", "document")
				local document = collection:load("document"):expect()
				local keyInfo = document:keyInfo()

				assert(before ~= keyInfo, "")
				assert(typeof(keyInfo) == "table", "")
				assert(keyInfo.Version == "0", "")
			end)

			x.test("updating user ids shouldn't affect key info", function(context)
				local collection = context.lapis.createCollection("collection", defaultOptions())

				local document = collection:load("document"):expect()
				local keyInfo = document:keyInfo()

				document:addUserId(123)

				assertEqual(#keyInfo:GetUserIds(), 0)
			end)

			x.test("key info is updated after :save", function(context)
				local document =
					context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()
				local keyInfo = document:keyInfo()

				document:save():expect()

				local newKeyInfo = document:keyInfo()

				assert(keyInfo ~= newKeyInfo, "")
				assert(newKeyInfo.Version == "1", "")
			end)

			x.test("key info is updated after :close", function(context)
				local document =
					context.lapis.createCollection("collection", defaultOptions()):load("document"):expect()
				local keyInfo = document:keyInfo()

				document:close():expect()

				local newKeyInfo = document:keyInfo()

				assert(keyInfo ~= newKeyInfo, "")
				assert(newKeyInfo.Version == "1", "")
			end)
		end)
	end)
end


type ErrorKind =
	"RobloxApiError"
	| "DefaultDataThrew"
	| "SessionLocked"
	| "MigrationError"
	| "BeforeSaveCloseCallbackThrew"
	| "DocumentRemoved"
	| "SessionLockStolen"
	| "ValidateThrew"
	| "ValidateFailed"

local Error = {}

function Error.new(kind: ErrorKind, message: string)
	return {
		kind = kind,
		message = message,
	}
end

return Error


local AutoSave = require(script.Parent.AutoSave)
local Collection = require(script.Parent.Collection)
local Config = require(script.Parent.Config)
local Data = require(script.Parent.Data)

local Internal = {}

function Internal.new(enableAutoSave)
	local config = Config.new()
	local data = Data.new(config)
	local autoSave = AutoSave.new(data)

	if enableAutoSave then
		autoSave:start()
	end

	local usedCollections = {}

	local internal = {}

	if not enableAutoSave then
		-- This exposes AutoSave to unit tests.
		internal.autoSave = autoSave
	end

	function internal.setConfig(values)
		config:set(values)
	end

	function internal.createCollection(name, options)
		if usedCollections[name] then
			error(`Collection "{name}" already exists`)
		end

		usedCollections[name] = true

		return Collection.new(name, options, data, autoSave, config)
	end

	return internal
end

return Internal


local Migration = {}

function Migration.getLastCompatibleVersion(migrations)
	local serverVersion = #migrations

	for version = serverVersion, 1, -1 do
		local migration = migrations[version]

		if migration.backwardsCompatible ~= true then
			return version
		end
	end

	return 0
end

function Migration.migrate(migrations, value, key)
	local serverVersion = #migrations
	local savedVersion = value.migrationVersion

	local data = value.data
	local lastCompatibleVersion = value.lastCompatibleVersion

	if serverVersion > savedVersion then
		for version = savedVersion + 1, #migrations do
			local ok, migrated = pcall(migrations[version].migrate, data, key)
			if not ok then
				return false, `Migration {version} threw an error: {migrated}`
			end

			if migrated == nil then
				return false, `Migration {version} returned 'nil'`
			end

			data = migrated
		end

		lastCompatibleVersion = Migration.getLastCompatibleVersion(migrations)
	elseif serverVersion < savedVersion then
		-- lastCompatibleVersion will be nil for documents that existed before backwards compatibilty was added and haven't been migrated to a new version since.
		if lastCompatibleVersion == nil or serverVersion < lastCompatibleVersion then
			return false,
				`Saved migration version {savedVersion} is not backwards compatible with version {serverVersion}`
		end
	end

	return true, data, lastCompatibleVersion
end

return Migration


export type Status = "Started" | "Resolved" | "Rejected" | "Cancelled"

export type Promise = {
	andThen: (
		self: Promise,
		successHandler: (...any) -> ...any,
		failureHandler: ((...any) -> ...any)?
	) -> Promise,
	andThenCall: <T...>(self: Promise, callback: (T...) -> ...any, T...) -> any,
	andThenReturn: (self: Promise, ...any) -> Promise,

	await: (self: Promise) -> (boolean, ...any),
	awaitStatus: (self: Promise) -> (Status, ...any),

	cancel: (self: Promise) -> (),
	catch: (self: Promise, failureHandler: (...any) -> ...any) -> Promise,
	expect: (self: Promise) -> ...any,

	finally: (self: Promise, finallyHandler: (status: Status) -> ...any) -> Promise,
	finallyCall: <T...>(self: Promise, callback: (T...) -> ...any, T...) -> Promise,
	finallyReturn: (self: Promise, ...any) -> Promise,

	getStatus: (self: Promise) -> Status,
	now: (self: Promise, rejectionValue: any?) -> Promise,
	tap: (self: Promise, tapHandler: (...any) -> ...any) -> Promise,
	timeout: (self: Promise, seconds: number, rejectionValue: any?) -> Promise,
}

export type TypedPromise<T...> = {
	andThen: (self: Promise, successHandler: (T...) -> ...any, failureHandler: ((...any) -> ...any)?) -> Promise,
	andThenCall: <T...>(self: Promise, callback: (T...) -> ...any, T...) -> Promise,
	andThenReturn: (self: Promise, ...any) -> Promise,

	await: (self: Promise) -> (boolean, T...),
	awaitStatus: (self: Promise) -> (Status, T...),

	cancel: (self: Promise) -> (),
	catch: (self: Promise, failureHandler: (...any) -> ...any) -> Promise,
	expect: (self: Promise) -> T...,

	finally: (self: Promise, finallyHandler: (status: Status) -> ...any) -> Promise,
	finallyCall: <T...>(self: Promise, callback: (T...) -> ...any, T...) -> Promise,
	finallyReturn: (self: Promise, ...any) -> Promise,

	getStatus: (self: Promise) -> Status,
	now: (self: Promise, rejectionValue: any?) -> Promise,
	tap: (self: Promise, tapHandler: (T...) -> ...any) -> Promise,
	timeout: (self: Promise, seconds: number, rejectionValue: any?) -> Promise,
}

return nil


local function copyDeep(value)
	if typeof(value) ~= "table" then
		return value
	end

	local new = table.clone(value)

	for k, v in value do
		if type(v) == "table" then
			new[k] = copyDeep(v)
		end
	end

	return new
end

return copyDeep


local function freezeDeep(value)
	if typeof(value) ~= "table" then
		return
	end

	if not table.isfrozen(value) then
		table.freeze(value)
	end

	for _, innerValue in value do
		freezeDeep(innerValue)
	end
end

return freezeDeep


local ReplicatedStorage = game:GetService("ReplicatedStorage")
local RunService = game:GetService("RunService")

local DataStoreServiceMock = require(ReplicatedStorage.DevPackages.DataStoreServiceMock)
local Internal = require(script.Parent.Internal)
local Promise = require(script.Parent.Parent.Promise)

local function defaultOptions()
	return {
		validate = function(data)
			return typeof(data.apples) == "number", "apples should be a number"
		end,
		defaultData = {
			apples = 20,
		},
	}
end

return function(x)
	local assertEqual = x.assertEqual
	local shouldThrow = x.shouldThrow

	x.beforeEach(function(context)
		local dataStoreService = DataStoreServiceMock.manual()

		context.dataStoreService = dataStoreService

		-- We want requests to overflow the throttle queue so that they result in errors.
		dataStoreService.budget:setMaxThrottleQueueSize(0)

		context.lapis = Internal.new(false)
		context.lapis.setConfig({ dataStoreService = dataStoreService, showRetryWarnings = false })

		context.write = function(name, key, data, lockId, userIds, metadata)
			local dataStore = dataStoreService.dataStores[name]["global"]

			dataStore:write(key, {
				migrationVersion = 0,
				lockId = lockId,
				data = data,
			}, userIds, metadata)
		end

		context.read = function(name, key)
			return dataStoreService.dataStores[name]["global"].data[key]
		end

		context.expectUnlocked = function(name, key)
			local data = dataStoreService.dataStores[name]["global"].data[key]

			if data ~= nil and data.lockId ~= nil then
				error("Document is locked")
			end
		end

		context.expectLocked = function(name, key)
			local data = dataStoreService.dataStores[name]["global"].data[key]

			if data == nil or data.lockId == nil then
				error("Document is not locked")
			end
		end

		context.expectUserIds = function(name, key, targetUserIds)
			local keyInfo = dataStoreService.dataStores[name]["global"].keyInfos[key]

			local currentUserIds = if keyInfo ~= nil then keyInfo:GetUserIds() else {}

			if #currentUserIds ~= #targetUserIds then
				error("Incorrect user ids length")
			end

			for index, value in targetUserIds do
				if currentUserIds[index] ~= value then
					error("Invalid user id")
				end
			end
		end

		context.getKeyInfo = function(name, key)
			return dataStoreService.dataStores[name]["global"].keyInfos[key]
		end
	end)

	x.test("throws when setting invalid config key", function(context)
		shouldThrow(function()
			context.lapis.setConfig({
				foo = true,
			})
		end, 'Invalid config key "foo"')
	end)

	x.test("throws when creating a duplicate collection", function(context)
		context.lapis.createCollection("foo", defaultOptions())

		shouldThrow(function()
			context.lapis.createCollection("foo", defaultOptions())
		end, 'Collection "foo" already exists')
	end)

	x.test("freezes default data", function(context)
		local defaultData = { a = { b = { c = 5 } } }

		context.lapis.createCollection("baz", {
			defaultData = defaultData,
		})

		shouldThrow(function()
			defaultData.a.b.c = 8
		end)
	end)

	x.test("validates default data as a table", function(context)
		shouldThrow(function()
			context.lapis.createCollection("bar", {
				validate = function()
					return false, "data is invalid"
				end,
			})
		end, "data is invalid")
	end)

	x.test("handles default data erroring", function(context)
		local collection = context.lapis.createCollection("collection", {
			defaultData = function()
				error("foo")
			end,
		})

		shouldThrow(function()
			collection:load("document"):expect()
		end, "'defaultData' threw an error", "foo")
	end)

	x.test("validates default data as a function", function(context)
		local collection = context.lapis.createCollection("collection", {
			defaultData = function()
				return {}
			end,
			validate = function()
				return false, "foo"
			end,
		})

		shouldThrow(function()
			collection:load("document"):expect()
		end, "Invalid data:", "foo")
	end)

	x.test("default data function should set default data", function(context)
		local collection = context.lapis.createCollection("collection", {
			defaultData = function()
				return "default"
			end,
		})

		local document = collection:load("document"):expect()

		assertEqual(document:read(), "default")
	end)

	x.test("should pass key to default data", function(context)
		local key
		local collection = context.lapis.createCollection("collection", {
			defaultData = function(passed)
				key = passed
				return {}
			end,
		})

		collection:load("document"):expect()
		assertEqual(key, "document")
	end)

	x.test("default data function should deep copy data", function(context)
		local returned = { {} }
		local collection = context.lapis.createCollection("collection", {
			defaultData = function()
				return returned
			end,
		})

		local document = collection:load("document"):expect()

		assert(document:read() ~= returned, "")
		assert(document:read()[1] ~= returned[1], "")
	end)

	x.test("default data function should freeze data", function(context)
		local collection = context.lapis.createCollection("collection", {
			defaultData = function()
				return {}
			end,
		})

		local document = collection:load("document"):expect()

		shouldThrow(function()
			document:read().foo = true
		end, "readonly")
	end)

	x.test("handles validate erroring", function(context)
		local created = false

		local collection = context.lapis.createCollection("collection", {
			validate = function()
				if created then
					error("foo")
				else
					return true
				end
			end,
		})

		created = true

		context.write("collection", "document", {})

		shouldThrow(function()
			collection:load("document"):expect()
		end, "'validate' threw an error", "foo")
	end)

	x.test("should not override data if validation fails", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())

		context.write("collection", "doc", { apples = "string" })

		local old = context.read("collection", "doc")

		shouldThrow(function()
			collection:load("doc"):expect()
		end, "apples should be a number")

		assertEqual(old, context.read("collection", "doc"))
	end)

	x.test("should session lock the document", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())
		local document = collection:load("doc"):expect()

		local otherLapis = Internal.new(false)
		otherLapis.setConfig({ dataStoreService = context.dataStoreService, loadAttempts = 1 })

		local otherCollection = otherLapis.createCollection("collection", defaultOptions())

		shouldThrow(function()
			otherCollection:load("doc"):expect()
		end, "Could not acquire lock")

		-- It should keep the session lock when saved.
		document:save():expect()

		shouldThrow(function()
			otherCollection:load("doc"):expect()
		end, "Could not acquire lock")

		-- It should remove the session lock when closed.
		document:close():expect()

		otherCollection:load("doc"):expect()
	end)

	x.test("load should retry when document is session locked", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())
		local document = collection:load("doc"):expect()

		local otherLapis = Internal.new(false)
		otherLapis.setConfig({
			dataStoreService = context.dataStoreService,
			loadAttempts = 4,
			loadRetryDelay = 0.5,
			showRetryWarnings = false,
		})

		local otherCollection = otherLapis.createCollection("collection", defaultOptions())
		local promise = otherCollection:load("doc")

		-- Wait for the document to attempt to load once.
		task.wait(0.1)

		-- Remove the sesssion lock.
		document:close():expect()

		promise:expect()
	end)

	x.test("second load should fail because of session lock", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())

		context.lapis.setConfig({ loadAttempts = 1 })

		local first = collection:load("document")
		local second = collection:load("document")

		first:expect()

		shouldThrow(function()
			second:expect()
		end, "Could not acquire lock")
	end)

	x.test("load returns a new promise when first load fails", function(context)
		context.lapis.setConfig({ loadAttempts = 1 })
		context.dataStoreService.errors:addSimulatedErrors(1)

		local collection = context.lapis.createCollection("ghi", defaultOptions())

		local promise1 = collection:load("ghi")

		shouldThrow(function()
			promise1:expect()
		end)

		local promise2 = collection:load("ghi")

		assert(promise1 ~= promise2, "load should return new promise")

		promise2:expect()
	end)

	x.nested("migrations", function()
		x.test("migrates the data", function(context)
			local collection = context.lapis.createCollection("migration", {
				validate = function(value)
					return value == "newData", "value does not equal newData"
				end,
				defaultData = "newData",
				migrations = {
					function()
						return "newData"
					end,
				},
			})

			context.write("migration", "migration", "data")

			collection:load("migration"):expect()

			local readData = collection:read("migration"):expect()

			assertEqual(readData, "newData")
		end)

		x.test("error is thrown if a migration returns nil", function(context)
			local collection = context.lapis.createCollection("collection", {
				defaultData = {},
				migrations = {
					function() end,
				},
			})

			context.write("collection", "document", {})

			shouldThrow(function()
				collection:load("document"):expect()
			end, "Migration 1 returned 'nil'")

			shouldThrow(function()
				collection:read("document"):expect()
			end, "Migration 1 returned 'nil'")
		end)

		x.test("passes key to migrations", function(context)
			local passedKey
			local collection = context.lapis.createCollection("collection", {
				defaultData = "newData",
				migrations = {
					function(_, key)
						passedKey = key
						return "newData"
					end,
				},
			})

			context.write("collection", "key", "data")

			collection:read("key"):expect()
			assertEqual(passedKey, "key")
			passedKey = nil

			collection:load("key"):expect()
			assertEqual(passedKey, "key")
		end)

		x.test("migrations should allow mutable updates", function(context)
			local collection = context.lapis.createCollection("collection", {
				validate = function(value)
					return typeof(value.coins) == "number"
				end,
				defaultData = { coins = 0 },
				migrations = {
					function(old)
						old.coins = 0

						return old
					end,
					function(old)
						old.coins = 100

						return old
					end,
				},
			})

			context.write("collection", "document", {})

			local document = collection:load("document"):expect()

			assertEqual(document:read().coins, 100)
		end)

		x.test("data should be frozen after a migration", function(context)
			local collection = context.lapis.createCollection("collection", {
				validate = function(value)
					return typeof(value.coins) == "number"
				end,
				defaultData = { coins = 0 },
				migrations = {
					function(old)
						old.coins = 0
						return old
					end,
				},
			})

			context.write("collection", "document", {})

			local document = collection:load("document"):expect()

			shouldThrow(function()
				document:read().coins = 100
			end, "readonly")
		end)

		x.test("migrations should work with tables and functions", function(context)
			local collection = context.lapis.createCollection("collection", {
				defaultData = "a",
				migrations = {
					{
						backwardsCompatible = false,
						migrate = function(old)
							return old
						end,
					},
					function(old)
						return old
					end,
				},
			})

			local dataStore = context.dataStoreService.dataStores.collection.global
			dataStore:write("document", {
				migrationVersion = 0,
				data = "a",
			})

			collection:load("document"):expect()
		end)

		x.nested("saved version ahead", function()
			x.test(
				"throws when migration version is ahead of latest version and is not backwards compatible",
				function(context)
					local collection = context.lapis.createCollection("collection", {
						defaultData = "a",
						migrations = {
							function(old)
								return old
							end,
						},
					})

					local dataStore = context.dataStoreService.dataStores.collection.global
					dataStore:write("document", {
						migrationVersion = 2,
						data = "b",
					})

					shouldThrow(function()
						collection:load("document"):expect()
					end, "Saved migration version 2 is not backwards compatible with version 1")

					shouldThrow(function()
						collection:read("document"):expect()
					end, "Saved migration version 2 is not backwards compatible with version 1")
				end
			)

			x.test("default data gets lastCompatibleVersion", function(context)
				local migrate = function(old)
					return old
				end

				local collection = context.lapis.createCollection("collection", {
					defaultData = "a",
					migrations = {
						{ migrate = migrate, backwardsCompatible = true },
					},
				})

				collection:load("document"):expect():close():expect()

				local otherLapis = Internal.new(false)
				otherLapis.setConfig({ dataStoreService = context.dataStoreService, loadAttempts = 1 })

				local otherCollection = otherLapis.createCollection("collection", {
					defaultData = "a",
				})

				-- This would error if lastCompatibleVersion = 0 wasn't saved.
				otherCollection:load("document"):expect()
			end)

			x.test("handles lastCompatibleVersion == nil", function(context)
				local collection = context.lapis.createCollection("collection", {
					defaultData = "a",
				})

				local dataStore = context.dataStoreService.dataStores.collection.global
				dataStore:write("document", {
					migrationVersion = 1,
					data = "b",
				})

				shouldThrow(function()
					collection:load("document"):expect()
				end, "Saved migration version 1 is not backwards compatible with version 0")

				shouldThrow(function()
					collection:read("document"):expect()
				end, "Saved migration version 1 is not backwards compatible with version 0")
			end)

			x.test("migration saves lastCompatibleVersion", function(context)
				local function migrate(old)
					return old
				end

				local collection = context.lapis.createCollection("collection", {
					defaultData = "a",
					migrations = {
						{ migrate = migrate, backwardsCompatible = false },
						{ migrate = migrate, backwardsCompatible = true },
						{ migrate = migrate, backwardsCompatible = true },
					},
				})

				local dataStore = context.dataStoreService.dataStores.collection.global
				dataStore:write("document", {
					migrationVersion = 0,
					data = "b",
				})

				collection:load("document"):expect():close():expect()

				local lapisWithV0 = Internal.new(false)
				lapisWithV0.setConfig({ dataStoreService = context.dataStoreService, loadAttempts = 1 })

				local collectionWithV0 = lapisWithV0.createCollection("collection", {
					defaultData = "a",
				})

				shouldThrow(function()
					collectionWithV0:load("document"):expect()
				end, "Saved migration version 3 is not backwards compatible with version 0")

				local lapisWithV1 = Internal.new(false)
				lapisWithV1.setConfig({ dataStoreService = context.dataStoreService, loadAttempts = 1 })

				local collectionWithV1 = lapisWithV1.createCollection("collection", {
					defaultData = "a",
					migrations = {
						{ migrate = migrate, backwardsCompatible = false },
						{ migrate = migrate, backwardsCompatible = true },
					},
				})

				-- This shouldn't error because v3 is backwards compatible with v1.
				collectionWithV1:load("document"):expect()
			end)

			x.test("keeps saved version", function(context)
				local collection = context.lapis.createCollection("collection", {
					defaultData = "a",
				})

				local dataStore = context.dataStoreService.dataStores.collection.global
				dataStore:write("document", {
					lastCompatibleVersion = 0,
					migrationVersion = 1,
					data = "b",
				})

				local document = collection:load("document"):expect()

				assertEqual(context.read("collection", "document").migrationVersion, 1)

				document:save("document"):expect()

				assertEqual(context.read("collection", "document").migrationVersion, 1)

				document:close("document"):expect()

				assertEqual(context.read("collection", "document").migrationVersion, 1)
			end)
		end)
	end)

	x.test("closing and immediately opening should return a new document", function(context)
		local collection = context.lapis.createCollection("ccc", defaultOptions())

		local document = collection:load("doc"):expect()

		local close = document:close()
		local open = collection:load("doc")

		close:expect()

		local newDocument = open:expect()

		assert(newDocument ~= document, "")
	end)

	x.test("closes all document on game:BindToClose", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())

		local one = collection:load("one"):expect()
		local two = collection:load("two"):expect()
		local three = collection:load("three"):expect()

		context.dataStoreService.yield:startYield()

		local thread = task.spawn(function()
			context.lapis.autoSave:onGameClose()
		end)

		assert(coroutine.status(thread) == "suspended", "onGameClose didn't wait for the documents to finish closing")

		-- Verify each document has been closed.
		for _, document in { one, two, three } do
			shouldThrow(function()
				document:save():expect()
			end, "Cannot save a closed document")
		end

		context.dataStoreService.yield:stopYield()

		-- Wait for documents to finish saving.
		task.wait(0.1)

		assert(coroutine.status(thread) == "dead", "")
	end)

	x.nested("user ids", function()
		x.test("it uses defaultUserIds on first load", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document", { 123 }):expect()
			context.expectUserIds("collection", "document", { 123 })
			document:close():expect()
			context.expectUserIds("collection", "document", { 123 })

			-- Since the document has already been created, the defaultUserIds should not override the saved ones.
			document = collection:load("document", { 321 }):expect()
			context.expectUserIds("collection", "document", { 123 })
			document:close():expect()
			context.expectUserIds("collection", "document", { 123 })
		end)

		x.test("adds new user ids", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document", {}):expect()

			document:addUserId(111)
			document:addUserId(111) -- It should not add this user id twice.
			document:addUserId(222)

			context.expectUserIds("collection", "document", {})

			document:save():expect()

			context.expectUserIds("collection", "document", { 111, 222 })

			document:close():expect()

			context.expectUserIds("collection", "document", { 111, 222 })
		end)

		x.test("removes new user ids", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document", { 333, 444, 555 }):expect()

			document:removeUserId(111) -- It should do nothing if the user id doesn't exist.
			document:removeUserId(444)

			context.expectUserIds("collection", "document", { 333, 444, 555 })

			document:save():expect()

			context.expectUserIds("collection", "document", { 333, 555 })

			document:close():expect()

			context.expectUserIds("collection", "document", { 333, 555 })
		end)
	end)

	x.nested("load during BindToClose", function()
		x.test("load infinitely yields after BindToClose", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			task.spawn(function()
				context.lapis.autoSave:onGameClose()
			end)

			shouldThrow(function()
				collection:load("document"):timeout(0.5):expect()
			end, "Timed out")
		end)

		x.test("load just before BindToClose", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			context.dataStoreService.yield:startYield()

			collection:load("document")

			local waited = false
			local finished = false
			local thread = task.spawn(function()
				RunService.PostSimulation:Wait()
				RunService.PostSimulation:Wait()
				waited = true
				context.lapis.autoSave:onGameClose()
				finished = true
			end)

			while not waited do
				task.wait()
			end

			context.dataStoreService.yield:stopYield()

			context.dataStoreService.yield:startYield()
			assert(
				coroutine.status(thread) == "suspended",
				"onGameClose didn't wait for the documents to finish closing"
			)
			context.dataStoreService.yield:stopYield()

			while not finished do
				task.wait()
			end

			context.expectUnlocked("collection", "document")

			assert(coroutine.status(thread) == "dead", "")
		end)

		x.test("BindToClose should finish if a document fails to load", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			context.write("collection", "document", "INVALID DATA")
			collection:load("document"):catch(function() end)

			-- Wait to close game so that the save request doesn't get cancelled.
			task.wait(0.1)

			Promise.try(function()
				context.lapis.autoSave:onGameClose()
			end)
				:timeout(1)
				:expect()
		end)
	end)

	x.nested("freezeData = false", function()
		x.test("default data should be deep copied", function(context)
			local defaultData = { foo = {} }

			local collection = context.lapis.createCollection("collection", {
				freezeData = false,
				defaultData = defaultData,
			})

			local document = collection:load("document"):expect()
			local data = document:read()

			assert(data ~= defaultData, "")
			assert(data.foo ~= defaultData.foo, "")
			assert(typeof(data.foo) == "table", "")
		end)

		x.test("data should not be frozen", function(context)
			local collection = context.lapis.createCollection("collection", {
				freezeData = false,
				defaultData = {},
			})
			local document = collection:load("document"):expect()

			-- This would error if the data was frozen.
			document:read().apples = 1

			-- Make sure write doesn't freeze the data.
			document:write(document:read())
			document:read().apples = 1
		end)

		x.test("should validate data in save and close", function(context)
			local valid = true
			local collection = context.lapis.createCollection("collection", {
				freezeData = false,
				validate = function()
					return valid, "data is invalid"
				end,
				defaultData = {},
			})

			local document = collection:load("document"):expect()

			valid = false

			shouldThrow(function()
				document:save():expect()
			end, "data is invalid")

			shouldThrow(function()
				document:close():expect()
			end, "data is invalid")
		end)

		x.test("should handle validate errors data in save and close", function(context)
			local throwError = false
			local collection = context.lapis.createCollection("collection", {
				freezeData = false,
				validate = function()
					if throwError then
						error("foo")
					end

					return true
				end,
				defaultData = {},
			})

			local document = collection:load("document"):expect()

			throwError = true

			shouldThrow(function()
				document:save():expect()
			end, "'validate' threw an error", "foo")

			shouldThrow(function()
				document:close():expect()
			end, "'validate' threw an error", "foo")
		end)
	end)

	x.nested("Collection:read", function()
		x.test("returns nil when there is no data", function(context)
			local collection = context.lapis.createCollection("collection", {
				defaultData = "data",
			})

			local data, keyInfo = collection:read("key"):expect()

			assertEqual(data, nil)
			assertEqual(keyInfo, nil)
		end)

		x.test("returns existing data", function(context)
			local collection = context.lapis.createCollection("collection", {
				defaultData = "data",
			})

			collection:load("key", { 321 }):expect()

			local data, keyInfo = collection:read("key"):expect()

			assertEqual(data, "data")
			assertEqual(keyInfo:GetUserIds()[1], 321)
		end)

		x.test("throws error when data is invalid", function(context)
			local collection = context.lapis.createCollection("collection", {
				defaultData = "data",
				validate = function(data)
					return data == "data", "data was invalid"
				end,
			})

			context.write("collection", "key", "INVALID DATA")

			shouldThrow(function()
				collection:read("key"):expect()
			end, "Invalid data")
		end)

		x.test("throws error when validate throws", function(context)
			local created = false
			local collection = context.lapis.createCollection("collection", {
				defaultData = "data",
				validate = function()
					if created then
						error("validate error")
					else
						return true
					end
				end,
			})
			created = true

			context.write("collection", "key", "data")

			shouldThrow(function()
				collection:read("key"):expect()
			end, "'validate' threw an error")
		end)
	end)

	x.nested("Collection:remove", function()
		x.test("should remove data", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			collection:load("document"):expect():close():expect()

			collection:remove("document"):expect()

			assertEqual(context.read("collection", "document"), nil)
		end)

		x.test("documents open during remove should fail to save/close", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document"):expect()

			collection:remove("document"):expect()

			shouldThrow(function()
				document:save():expect()
			end, "The document was removed")

			shouldThrow(function()
				document:close():expect()
			end, "The document was removed")
		end)
	end)

	x.nested("Document:close should still unlock after specific errors", function()
		x.test("shouldn't overwrite stolen lock", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())
			local document = collection:load("document"):expect()

			context.write("collection", "document", { apples = 20 }, "stolen lock")

			document:beforeSave(function()
				error("oh no")
			end)

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectLocked("collection", "document")
		end)

		x.test("doesn't work for session lock stolen error", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())
			local document = collection:load("document"):expect()

			context.write("collection", "document", { apples = 20 }, "another lock id")

			document:write({ apples = 100 })

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectLocked("collection", "document")
			assertEqual(context.read("collection", "document").data.apples, 20) -- Only the lock should have changed.
		end)

		x.test("beforeSave error", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document"):expect()

			document:write({ apples = 100 })

			document:beforeSave(function()
				error("oh no")
			end)

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectUnlocked("collection", "document")
			assertEqual(context.read("collection", "document").data.apples, 20) -- Only the lock should have changed.
		end)

		x.test("beforeClose error", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())

			local document = collection:load("document"):expect()

			document:write({ apples = 100 })

			document:beforeClose(function()
				error("oh no")
			end)

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectUnlocked("collection", "document")
			assertEqual(context.read("collection", "document").data.apples, 20) -- Only the lock should have changed.
		end)

		x.test("validate error", function(context)
			local collection = context.lapis.createCollection("collection", {
				validate = function(data)
					return typeof(data.apples) == "number", "apples should be a number"
				end,
				defaultData = { apples = 20 },
				freezeData = false,
			})

			local document = collection:load("document"):expect()

			document:read().apples = nil

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectUnlocked("collection", "document")
			assertEqual(context.read("collection", "document").data.apples, 20) -- Only the lock should have changed.
		end)

		x.test("validate threw error", function(context)
			local loaded = false
			local collection = context.lapis.createCollection("collection", {
				validate = function()
					if loaded then
						error("oh no")
					end
					return true
				end,
				defaultData = { apples = 20 },
				freezeData = false,
			})

			local document = collection:load("document"):expect()

			loaded = true

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			context.expectUnlocked("collection", "document")
			assertEqual(context.read("collection", "document").data.apples, 20) -- Only the lock should have changed.
		end)

		x.test("onGameClose should wait for the lock to remove", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())
			local document = collection:load("document"):expect()

			document:beforeSave(function()
				error("oh no")
			end)

			context.dataStoreService.yield:startYield()

			shouldThrow(function()
				document:close("document"):expect()
			end)

			local thread = task.spawn(function()
				context.lapis.autoSave:onGameClose()
			end)

			assert(coroutine.status(thread) == "suspended", "onGameClose didn't wait for locks to be removed")

			context.dataStoreService.yield:stopYield()

			-- Wait for locks to be removed.
			task.wait(0.1)

			assert(coroutine.status(thread) == "dead", "")
		end)

		x.test("should preserve userids/metadata", function(context)
			local collection = context.lapis.createCollection("collection", defaultOptions())
			local document = collection:load("document"):expect()

			document:beforeSave(function()
				error("oh no")
			end)

			context.write(
				"collection",
				"document",
				document:read(),
				context.read("collection", "document").lockId,
				{ 1234 },
				{ foo = "bar" }
			)

			shouldThrow(function()
				document:close("document"):expect()
			end)

			task.wait(0.1)

			local keyInfo = context.getKeyInfo("collection", "document")

			assertEqual(keyInfo:GetUserIds()[1], 1234)
			assertEqual(keyInfo:GetMetadata().foo, "bar")
		end)
	end)

	x.test("preserves metadata", function(context)
		local collection = context.lapis.createCollection("collection", defaultOptions())

		context.write("collection", "document", { apples = 30 }, nil, nil, { foo = "bar" })

		local function verifyMetadata()
			local keyInfo = context.getKeyInfo("collection", "document")
			assertEqual(keyInfo:GetMetadata().foo, "bar")
		end

		local document = collection:load("document"):expect()
		verifyMetadata()

		document:save():expect()
		verifyMetadata()

		document:close():expect()
		verifyMetadata()
	end)
end


local function noYield(callback)
	for _ in
		function()
			callback()
		end
	do
		break
	end
end

return noYield


local process = require("@lune/process")

process.exec("rojo", { "build", "--output", "crunchyroll.rbxm" }, { stdio = "forward" })


local fs = require("@lune/fs")
local process = require("@lune/process")
local serde = require("@lune/serde")

local mirror_configs = fs.readDir("./benches/mirror_list")

-- thank you nick from lute
local function get_git_version(): { major: number, minor: number, path: number }
	local git_ver_string = process.exec("git", { "--version" }).stdout
	local git_ver = string.split(git_ver_string, " ")[3]
	local components = string.split(git_ver, ".")

	return {
		major = tonumber(components[1]) or 0,
		minor = tonumber(components[2]) or 0,
		path = tonumber(components[3]) or 0,
	}
end

local git_version = get_git_version()

for _, file_name in mirror_configs do
	local path = "./benches/mirror_list/" .. file_name

	local config: unknown = serde.decode("toml", fs.readFile(path))
	assert(type(config) == "table", `[mirror] {path} must be a table`)

	-- Lovely type errors :3
	assert(type(config.revision) == "string", `[mirror] {path} config.revision must be a string`)
	assert(type(config.branch) == "string", `[mirror] {path} config.branch must be a string`)
	assert(type(config.name) == "string", `[mirror] {path} config.name must be a string`)
	assert(type(config.remote) == "string", `[mirror] {path} config.remote must be a string`)

	local mirror_path = "./benches/mirrors/" .. config.name
	local full_mirror_path = mirror_path .. "/src/"
	if fs.isDir(mirror_path) then
		process.exec(
			"git",
			{ "fetch", "--depth=1", "origin", config.revision },
			{ stdio = "forward", cwd = full_mirror_path }
		)

		process.exec(
			"git",
			{ "checkout", config.revision },
			{ stdio = "forward", cwd = full_mirror_path }
		)

		continue
	end

	if git_version.major >= 3 or (git_version.major == 2 and git_version.minor >= 49) then
		process.exec("git", {
			"clone",
			"--depth=1",
			"--revision",
			config.revision,
			config.remote,
			full_mirror_path,
		})
	else
		process.exec("git", {
			"clone",
			"--depth=1",
			"--branch",
			config.branch,
			config.remote,
			full_mirror_path,
		})
	end
end