local Types = require(script.Parent.Types)

local widgets = {} :: Types.WidgetUtility

return function(Iris: Types.Internal)
    widgets.GuiService = game:GetService("GuiService")
    widgets.RunService = game:GetService("RunService")
    widgets.UserInputService = game:GetService("UserInputService")
    widgets.ContextActionService = game:GetService("ContextActionService")
    widgets.TextService = game:GetService("TextService")

    widgets.ICONS = {
        BLANK_SQUARE = "rbxassetid://83265623867126",
        RIGHT_POINTING_TRIANGLE = "rbxassetid://105541346271951",
        DOWN_POINTING_TRIANGLE = "rbxassetid://95465797476827",
        MULTIPLICATION_SIGN = "rbxassetid://133890060015237", -- best approximation for a close X which roblox supports, needs to be scaled about 2x
        BOTTOM_RIGHT_CORNER = "rbxassetid://125737344915000", -- used in window resize icon in bottom right
        CHECKMARK = "rbxassetid://109638815494221",
        BORDER = "rbxassetid://133803690460269",
        ALPHA_BACKGROUND_TEXTURE = "rbxassetid://114090016039876", -- used for color4 alpha
        UNKNOWN_TEXTURE = "rbxassetid://95045813476061",
    }

    widgets.IS_STUDIO = widgets.RunService:IsStudio()
    function widgets.getTime()
        -- time() always returns 0 in the context of plugins
        if widgets.IS_STUDIO then
            return os.clock()
        else
            return time()
        end
    end

    -- acts as an offset where the absolute position of the base frame is not zero, such as IgnoreGuiInset or for stories
    widgets.GuiOffset = if Iris._config.IgnoreGuiInset then -widgets.GuiService:GetGuiInset() else Vector2.zero
    -- the registered mouse position always ignores the topbar, so needs a separate variable offset
    widgets.MouseOffset = if Iris._config.IgnoreGuiInset then Vector2.zero else widgets.GuiService:GetGuiInset()

    -- the topbar inset changes updates a frame later.
    local connection: RBXScriptConnection
    connection = widgets.GuiService:GetPropertyChangedSignal("TopbarInset"):Once(function()
        widgets.MouseOffset = if Iris._config.IgnoreGuiInset then Vector2.zero else widgets.GuiService:GetGuiInset()
        widgets.GuiOffset = if Iris._config.IgnoreGuiInset then -widgets.GuiService:GetGuiInset() else Vector2.zero
        connection:Disconnect()
    end)
    -- in case the topbar doesn't change, we cancel the event.
    task.delay(5, function()
        connection:Disconnect()
    end)

    function widgets.getMouseLocation()
        return widgets.UserInputService:GetMouseLocation() - widgets.MouseOffset
    end

    function widgets.isPosInsideRect(pos: Vector2, rectMin: Vector2, rectMax: Vector2)
        return pos.X >= rectMin.X and pos.X <= rectMax.X and pos.Y >= rectMin.Y and pos.Y <= rectMax.Y
    end

    function widgets.findBestWindowPosForPopup(refPos: Vector2, size: Vector2, outerMin: Vector2, outerMax: Vector2)
        local CURSOR_OFFSET_DIST = 20

        if refPos.X + size.X + CURSOR_OFFSET_DIST > outerMax.X then
            if refPos.Y + size.Y + CURSOR_OFFSET_DIST > outerMax.Y then
                -- placed to the top
                refPos += Vector2.new(0, -(CURSOR_OFFSET_DIST + size.Y))
            else
                -- placed to the bottom
                refPos += Vector2.new(0, CURSOR_OFFSET_DIST)
            end
        else
            -- placed to the right
            refPos += Vector2.new(CURSOR_OFFSET_DIST)
        end

        return Vector2.new(math.max(math.min(refPos.X + size.X, outerMax.X) - size.X, outerMin.X), math.max(math.min(refPos.Y + size.Y, outerMax.Y) - size.Y, outerMin.Y))
    end

    function widgets.getScreenSizeForWindow(thisWidget: Types.Widget) -- possible parents are GuiBase2d, CoreGui, PlayerGui
        if thisWidget.Instance:IsA("GuiBase2d") then
            return thisWidget.Instance.AbsoluteSize
        else
            local rootParent = thisWidget.Instance.Parent
            if rootParent:IsA("GuiBase2d") then
                return rootParent.AbsoluteSize
            else
                if rootParent.Parent:IsA("GuiBase2d") then
                    return rootParent.AbsoluteSize
                else
                    return workspace.CurrentCamera.ViewportSize
                end
            end
        end
    end

    function widgets.extend(superClass: Types.WidgetClass, subClass: Types.WidgetClass): Types.WidgetClass
        local newClass = table.clone(superClass)
        for index, value in subClass do
            newClass[index] = value
        end
        return newClass
    end

    function widgets.UIPadding(Parent: GuiObject, PxPadding: Vector2)
        local UIPaddingInstance = Instance.new("UIPadding")
        UIPaddingInstance.PaddingLeft = UDim.new(0, PxPadding.X)
        UIPaddingInstance.PaddingRight = UDim.new(0, PxPadding.X)
        UIPaddingInstance.PaddingTop = UDim.new(0, PxPadding.Y)
        UIPaddingInstance.PaddingBottom = UDim.new(0, PxPadding.Y)
        UIPaddingInstance.Parent = Parent
        return UIPaddingInstance
    end

    function widgets.UIListLayout(Parent: GuiObject, FillDirection: Enum.FillDirection, Padding: UDim)
        local UIListLayoutInstance = Instance.new("UIListLayout")
        UIListLayoutInstance.SortOrder = Enum.SortOrder.LayoutOrder
        UIListLayoutInstance.Padding = Padding
        UIListLayoutInstance.FillDirection = FillDirection
        UIListLayoutInstance.Parent = Parent
        return UIListLayoutInstance
    end

    function widgets.UIStroke(Parent: GuiObject, Thickness: number, Color: Color3, Transparency: number)
        local UIStrokeInstance = Instance.new("UIStroke")
        UIStrokeInstance.Thickness = Thickness
        UIStrokeInstance.Color = Color
        UIStrokeInstance.Transparency = Transparency
        UIStrokeInstance.ApplyStrokeMode = Enum.ApplyStrokeMode.Border
        UIStrokeInstance.LineJoinMode = Enum.LineJoinMode.Round
        UIStrokeInstance.Parent = Parent
        return UIStrokeInstance
    end

    function widgets.UICorner(Parent: GuiObject, PxRounding: number?)
        local UICornerInstance = Instance.new("UICorner")
        UICornerInstance.CornerRadius = UDim.new(PxRounding and 0 or 1, PxRounding or 0)
        UICornerInstance.Parent = Parent
        return UICornerInstance
    end

    function widgets.UISizeConstraint(Parent: GuiObject, MinSize: Vector2?, MaxSize: Vector2?)
        local UISizeConstraintInstance = Instance.new("UISizeConstraint")
        UISizeConstraintInstance.MinSize = MinSize or UISizeConstraintInstance.MinSize -- made these optional
        UISizeConstraintInstance.MaxSize = MaxSize or UISizeConstraintInstance.MaxSize
        UISizeConstraintInstance.Parent = Parent
        return UISizeConstraintInstance
    end

    -- below uses Iris

    function widgets.applyTextStyle(thisInstance: TextLabel & TextButton & TextBox)
        thisInstance.FontFace = Iris._config.TextFont
        thisInstance.TextSize = Iris._config.TextSize
        thisInstance.TextColor3 = Iris._config.TextColor
        thisInstance.TextTransparency = Iris._config.TextTransparency
        thisInstance.TextXAlignment = Enum.TextXAlignment.Left
        thisInstance.TextYAlignment = Enum.TextYAlignment.Center
        thisInstance.RichText = Iris._config.RichText
        thisInstance.TextWrapped = Iris._config.TextWrapped

        thisInstance.AutoLocalize = false
    end

    function widgets.applyInteractionHighlights(Property: string, Button: GuiButton, Highlightee: GuiObject, Colors: { [string]: any })
        local exitedButton = false
        widgets.applyMouseEnter(Button, function()
            Highlightee[Property .. "Color3"] = Colors.HoveredColor
            Highlightee[Property .. "Transparency"] = Colors.HoveredTransparency

            exitedButton = false
        end)

        widgets.applyMouseLeave(Button, function()
            Highlightee[Property .. "Color3"] = Colors.Color
            Highlightee[Property .. "Transparency"] = Colors.Transparency

            exitedButton = true
        end)

        widgets.applyInputBegan(Button, function(input: InputObject)
            if not (input.UserInputType == Enum.UserInputType.MouseButton1 or input.UserInputType == Enum.UserInputType.Gamepad1) then
                return
            end
            Highlightee[Property .. "Color3"] = Colors.ActiveColor
            Highlightee[Property .. "Transparency"] = Colors.ActiveTransparency
        end)

        widgets.applyInputEnded(Button, function(input: InputObject)
            if not (input.UserInputType == Enum.UserInputType.MouseButton1 or input.UserInputType == Enum.UserInputType.Gamepad1) or exitedButton then
                return
            end
            if input.UserInputType == Enum.UserInputType.MouseButton1 then
                Highlightee[Property .. "Color3"] = Colors.HoveredColor
                Highlightee[Property .. "Transparency"] = Colors.HoveredTransparency
            end
            if input.UserInputType == Enum.UserInputType.Gamepad1 then
                Highlightee[Property .. "Color3"] = Colors.Color
                Highlightee[Property .. "Transparency"] = Colors.Transparency
            end
        end)

        Button.SelectionImageObject = Iris.SelectionImageObject
    end

    function widgets.applyInteractionHighlightsWithMultiHighlightee(Property: string, Button: GuiButton, Highlightees: { { GuiObject | { [string]: Color3 | number } } })
        local exitedButton = false
        widgets.applyMouseEnter(Button, function()
            for _, Highlightee in Highlightees do
                Highlightee[1][Property .. "Color3"] = Highlightee[2].HoveredColor
                Highlightee[1][Property .. "Transparency"] = Highlightee[2].HoveredTransparency

                exitedButton = false
            end
        end)

        widgets.applyMouseLeave(Button, function()
            for _, Highlightee in Highlightees do
                Highlightee[1][Property .. "Color3"] = Highlightee[2].Color
                Highlightee[1][Property .. "Transparency"] = Highlightee[2].Transparency

                exitedButton = true
            end
        end)

        widgets.applyInputBegan(Button, function(input: InputObject)
            if not (input.UserInputType == Enum.UserInputType.MouseButton1 or input.UserInputType == Enum.UserInputType.Gamepad1) then
                return
            end
            for _, Highlightee in Highlightees do
                Highlightee[1][Property .. "Color3"] = Highlightee[2].ActiveColor
                Highlightee[1][Property .. "Transparency"] = Highlightee[2].ActiveTransparency
            end
        end)

        widgets.applyInputEnded(Button, function(input: InputObject)
            if not (input.UserInputType == Enum.UserInputType.MouseButton1 or input.UserInputType == Enum.UserInputType.Gamepad1) or exitedButton then
                return
            end
            for _, Highlightee in Highlightees do
                if input.UserInputType == Enum.UserInputType.MouseButton1 then
                    Highlightee[1][Property .. "Color3"] = Highlightee[2].HoveredColor
                    Highlightee[1][Property .. "Transparency"] = Highlightee[2].HoveredTransparency
                end
                if input.UserInputType == Enum.UserInputType.Gamepad1 then
                    Highlightee[1][Property .. "Color3"] = Highlightee[2].Color
                    Highlightee[1][Property .. "Transparency"] = Highlightee[2].Transparency
                end
            end
        end)

        Button.SelectionImageObject = Iris.SelectionImageObject
    end

    function widgets.applyFrameStyle(thisInstance: GuiObject, noPadding: boolean?, noCorner: boolean?)
        -- padding, border, and rounding
        -- optimized to only use what instances are needed, based on style
        local FrameBorderSize = Iris._config.FrameBorderSize
        local FrameRounding = Iris._config.FrameRounding
        thisInstance.BorderSizePixel = 0

        if FrameBorderSize > 0 then
            widgets.UIStroke(thisInstance, FrameBorderSize, Iris._config.BorderColor, Iris._config.BorderTransparency)
        end
        if FrameRounding > 0 and not noCorner then
            widgets.UICorner(thisInstance, FrameRounding)
        end
        if not noPadding then
            widgets.UIPadding(thisInstance, Iris._config.FramePadding)
        end
    end

    function widgets.applyButtonClick(thisInstance: GuiButton, callback: () -> ())
        thisInstance.MouseButton1Click:Connect(function()
            callback()
        end)
    end

    function widgets.applyButtonDown(thisInstance: GuiButton, callback: (x: number, y: number) -> ())
        thisInstance.MouseButton1Down:Connect(function(x: number, y: number)
            local position = Vector2.new(x, y) - widgets.MouseOffset
            callback(position.X, position.Y)
        end)
    end

    function widgets.applyMouseEnter(thisInstance: GuiObject, callback: (x: number, y: number) -> ())
        thisInstance.MouseEnter:Connect(function(x: number, y: number)
            local position = Vector2.new(x, y) - widgets.MouseOffset
            callback(position.X, position.Y)
        end)
    end

    function widgets.applyMouseMoved(thisInstance: GuiObject, callback: (x: number, y: number) -> ())
        thisInstance.MouseMoved:Connect(function(x: number, y: number)
            local position = Vector2.new(x, y) - widgets.MouseOffset
            callback(position.X, position.Y)
        end)
    end

    function widgets.applyMouseLeave(thisInstance: GuiObject, callback: (x: number, y: number) -> ())
        thisInstance.MouseLeave:Connect(function(x: number, y: number)
            local position = Vector2.new(x, y) - widgets.MouseOffset
            callback(position.X, position.Y)
        end)
    end

    function widgets.applyInputBegan(thisInstance: GuiButton, callback: (input: InputObject) -> ())
        thisInstance.InputBegan:Connect(function(...)
            callback(...)
        end)
    end

    function widgets.applyInputEnded(thisInstance: GuiButton, callback: (input: InputObject) -> ())
        thisInstance.InputEnded:Connect(function(...)
            callback(...)
        end)
    end

    function widgets.discardState(thisWidget: Types.StateWidget)
        for _, state in thisWidget.state do
            state.ConnectedWidgets[thisWidget.ID] = nil
        end
    end

    function widgets.registerEvent(event: string, callback: (...any) -> ())
        table.insert(Iris._initFunctions, function()
            table.insert(Iris._connections, widgets.UserInputService[event]:Connect(callback))
        end)
    end

    widgets.EVENTS = {
        hover = function(pathToHovered: (thisWidget: Types.Widget) -> GuiObject)
            return {
                ["Init"] = function(thisWidget: Types.Widget & Types.Hovered)
                    local hoveredGuiObject = pathToHovered(thisWidget)
                    widgets.applyMouseEnter(hoveredGuiObject, function()
                        thisWidget.isHoveredEvent = true
                    end)
                    widgets.applyMouseLeave(hoveredGuiObject, function()
                        thisWidget.isHoveredEvent = false
                    end)
                    thisWidget.isHoveredEvent = false
                end,
                ["Get"] = function(thisWidget: Types.Widget & Types.Hovered)
                    return thisWidget.isHoveredEvent
                end,
            }
        end,

        click = function(pathToClicked: (thisWidget: Types.Widget) -> GuiButton)
            return {
                ["Init"] = function(thisWidget: Types.Widget & Types.Clicked)
                    local clickedGuiObject = pathToClicked(thisWidget)
                    thisWidget.lastClickedTick = -1

                    widgets.applyButtonClick(clickedGuiObject, function()
                        thisWidget.lastClickedTick = Iris._cycleTick + 1
                    end)
                end,
                ["Get"] = function(thisWidget: Types.Widget & Types.Clicked)
                    return thisWidget.lastClickedTick == Iris._cycleTick
                end,
            }
        end,

        rightClick = function(pathToClicked: (thisWidget: Types.Widget) -> GuiButton)
            return {
                ["Init"] = function(thisWidget: Types.Widget & Types.RightClicked)
                    local clickedGuiObject = pathToClicked(thisWidget)
                    thisWidget.lastRightClickedTick = -1

                    clickedGuiObject.MouseButton2Click:Connect(function()
                        thisWidget.lastRightClickedTick = Iris._cycleTick + 1
                    end)
                end,
                ["Get"] = function(thisWidget: Types.Widget & Types.RightClicked)
                    return thisWidget.lastRightClickedTick == Iris._cycleTick
                end,
            }
        end,

        doubleClick = function(pathToClicked: (thisWidget: Types.Widget) -> GuiButton)
            return {
                ["Init"] = function(thisWidget: Types.Widget & Types.DoubleClicked)
                    local clickedGuiObject = pathToClicked(thisWidget)
                    thisWidget.lastClickedTime = -1
                    thisWidget.lastClickedPosition = Vector2.zero
                    thisWidget.lastDoubleClickedTick = -1

                    widgets.applyButtonDown(clickedGuiObject, function(x: number, y: number)
                        local currentTime = widgets.getTime()
                        local isTimeValid = currentTime - thisWidget.lastClickedTime < Iris._config.MouseDoubleClickTime
                        if isTimeValid and (Vector2.new(x, y) - thisWidget.lastClickedPosition).Magnitude < Iris._config.MouseDoubleClickMaxDist then
                            thisWidget.lastDoubleClickedTick = Iris._cycleTick + 1
                        else
                            thisWidget.lastClickedTime = currentTime
                            thisWidget.lastClickedPosition = Vector2.new(x, y)
                        end
                    end)
                end,
                ["Get"] = function(thisWidget: Types.Widget & Types.DoubleClicked)
                    return thisWidget.lastDoubleClickedTick == Iris._cycleTick
                end,
            }
        end,

        ctrlClick = function(pathToClicked: (thisWidget: Types.Widget) -> GuiButton)
            return {
                ["Init"] = function(thisWidget: Types.Widget & Types.CtrlClicked)
                    local clickedGuiObject = pathToClicked(thisWidget)
                    thisWidget.lastCtrlClickedTick = -1

                    widgets.applyButtonClick(clickedGuiObject, function()
                        if widgets.UserInputService:IsKeyDown(Enum.KeyCode.LeftControl) or widgets.UserInputService:IsKeyDown(Enum.KeyCode.RightControl) then
                            thisWidget.lastCtrlClickedTick = Iris._cycleTick + 1
                        end
                    end)
                end,
                ["Get"] = function(thisWidget: Types.Widget & Types.CtrlClicked)
                    return thisWidget.lastCtrlClickedTick == Iris._cycleTick
                end,
            }
        end,
    }

    Iris._utility = widgets

    require(script.Root)(Iris, widgets)
    require(script.Window)(Iris, widgets)

    require(script.Menu)(Iris, widgets)

    require(script.Format)(Iris, widgets)

    require(script.Text)(Iris, widgets)
    require(script.Button)(Iris, widgets)
    require(script.Checkbox)(Iris, widgets)
    require(script.RadioButton)(Iris, widgets)
    require(script.Image)(Iris, widgets)

    require(script.Tree)(Iris, widgets)
    require(script.Tab)(Iris, widgets)

    require(script.Input)(Iris, widgets)
    require(script.Combo)(Iris, widgets)
    require(script.Plot)(Iris, widgets)

    require(script.Table)(Iris, widgets)
end


local ReplicatedStorage = game:GetService("ReplicatedStorage")
local Iris = require(ReplicatedStorage.Iris).Init()

Iris:Connect(Iris.ShowDemoWindow)


--------------------------------------------------------------------------------
--               Batched Yield-Safe Signal Implementation                     --
-- This is a Signal class which has effectively identical behavior to a       --
-- normal RBXScriptSignal, with the only difference being a couple extra      --
-- stack frames at the bottom of the stack trace when an error is thrown.     --
-- This implementation caches runner coroutines, so the ability to yield in   --
-- the signal handlers comes at minimal extra cost over a naive signal        --
-- implementation that either always or never spawns a thread.                --
--                                                                            --
-- API:                                                                       --
--   local Signal = require(THIS MODULE)                                      --
--   local sig = Signal.new()                                                 --
--   local connection = sig:Connect(function(arg1, arg2, ...) ... end)        --
--   sig:Fire(arg1, arg2, ...)                                                --
--   connection:Disconnect()                                                  --
--   sig:DisconnectAll()                                                      --
--   local arg1, arg2, ... = sig:Wait()                                       --
--                                                                            --
-- Licence:                                                                   --
--   Licenced under the MIT licence.                                          --
--                                                                            --
-- Authors:                                                                   --
--   stravant - July 31st, 2021 - Created the file.                           --
--------------------------------------------------------------------------------

-- The currently idle thread to run the next handler on
local freeRunnerThread = nil

-- Function which acquires the currently idle handler runner thread, runs the
-- function fn on it, and then releases the thread, returning it to being the
-- currently idle one.
-- If there was a currently idle runner thread already, that's okay, that old
-- one will just get thrown and eventually GCed.
local function acquireRunnerThreadAndCallEventHandler(fn, ...)
    local acquiredRunnerThread = freeRunnerThread
    freeRunnerThread = nil
    fn(...)
    -- The handler finished running, this runner thread is free again.
    freeRunnerThread = acquiredRunnerThread
end

-- Coroutine runner that we create coroutines of. The coroutine can be
-- repeatedly resumed with functions to run followed by the argument to run
-- them with.
local function runEventHandlerInFreeThread()
    -- Note: We cannot use the initial set of arguments passed to
    -- runEventHandlerInFreeThread for a call to the handler, because those
    -- arguments would stay on the stack for the duration of the thread's
    -- existence, temporarily leaking references. Without access to raw bytecode
    -- there's no way for us to clear the "..." references from the stack.
    while true do
        acquireRunnerThreadAndCallEventHandler(coroutine.yield())
    end
end

-- Connection class
local Connection = {}
Connection.__index = Connection

function Connection.new(signal, fn)
    return setmetatable({
        _connected = true,
        _signal = signal,
        _fn = fn,
        _next = false,
    }, Connection)
end

function Connection:Disconnect()
    self._connected = false

    -- Unhook the node, but DON'T clear it. That way any fire calls that are
    -- currently sitting on this node will be able to iterate forwards off of
    -- it, but any subsequent fire calls will not hit it, and it will be GCed
    -- when no more fire calls are sitting on it.
    if self._signal._handlerListHead == self then
        self._signal._handlerListHead = self._next
    else
        local prev = self._signal._handlerListHead
        while prev and prev._next ~= self do
            prev = prev._next
        end
        if prev then
            prev._next = self._next
        end
    end
end

-- Make Connection strict
setmetatable(Connection, {
    __index = function(_, key)
        error(("Attempt to get Connection::%s (not a valid member)"):format(tostring(key)), 2)
    end,
    __newindex = function(_, key, _)
        error(("Attempt to set Connection::%s (not a valid member)"):format(tostring(key)), 2)
    end,
})

-- Signal class
local Signal = {}
Signal.__index = Signal

function Signal.new()
    return setmetatable({
        _handlerListHead = false,
    }, Signal)
end

function Signal:Connect(fn)
    local connection = Connection.new(self, fn)
    if self._handlerListHead then
        connection._next = self._handlerListHead
        self._handlerListHead = connection
    else
        self._handlerListHead = connection
    end
    return connection
end

-- Disconnect all handlers. Since we use a linked list it suffices to clear the
-- reference to the head handler.
function Signal:DisconnectAll()
    self._handlerListHead = false
end

-- Signal:Fire(...) implemented by running the handler functions on the
-- coRunnerThread, and any time the resulting thread yielded without returning
-- to us, that means that it yielded to the Roblox scheduler and has been taken
-- over by Roblox scheduling, meaning we have to make a new coroutine runner.
function Signal:Fire(...)
    local item = self._handlerListHead
    while item do
        if item._connected then
            if not freeRunnerThread then
                freeRunnerThread = coroutine.create(runEventHandlerInFreeThread)
                -- Get the freeRunnerThread to the first yield
                coroutine.resume(freeRunnerThread)
            end
            task.spawn(freeRunnerThread, item._fn, ...)
        end
        item = item._next
    end
end

-- Implement Signal:Wait() in terms of a temporary connection using
-- a Signal:Connect() which disconnects itself.
function Signal:Wait()
    local waitingCoroutine = coroutine.running()
    local cn
    cn = self:Connect(function(...)
        cn:Disconnect()
        task.spawn(waitingCoroutine, ...)
    end)
    return coroutine.yield()
end

-- Implement Signal:Once() in terms of a connection which disconnects
-- itself before running the handler.
function Signal:Once(fn)
    local cn
    cn = self:Connect(function(...)
        if cn._connected then
            cn:Disconnect()
        end
        fn(...)
    end)
    return cn
end

-- Make signal strict
setmetatable(Signal, {
    __index = function(_, key)
        error(("Attempt to get Signal::%s (not a valid member)"):format(tostring(key)), 2)
    end,
    __newindex = function(_, key, _)
        error(("Attempt to set Signal::%s (not a valid member)"):format(tostring(key)), 2)
    end,
})

return Signal


local Signal = require(script.Signal)

local UserInputService: UserInputService = game:GetService("UserInputService")

local Input = {}

Input.X = 0
Input.Y = 0
Input.KeyDown = {}
Input._connections = {}

-- This frame will act as our UserInputService detector. Most events should go through it.
-- it's not perfect, but there's not a better alternative (I think)
local SinkFrame: Frame = Instance.new("Frame")
SinkFrame.Name = "SinkFrame"
SinkFrame.AnchorPoint = Vector2.new(0.5, 0.5)
SinkFrame.Position = UDim2.fromScale(0.5, 0.5)
SinkFrame.Size = UDim2.fromScale(1, 1)
SinkFrame.BackgroundTransparency = 1
SinkFrame.ZIndex = 1024 ^ 2

Input.SinkFrame = SinkFrame

-- Methods and events

Input.InputBegan = Signal.new()
Input.InputChanged = Signal.new()
Input.InputEnded = Signal.new()
Input.MouseMoved = Signal.new()
Input.TouchTapInWorld = Signal.new()

function Input:GetMouseLocation()
    return Vector2.new(Input.X, Input.Y)
end

function Input:IsKeyDown(keyCode: Enum.KeyCode)
    if Input.KeyDown[keyCode] == true then
        return true
    end
    return false
end

-- UserInputService hooks

function inputBegan(input: InputObject)
    Input.KeyDown[input.KeyCode] = true
    Input.InputBegan:Fire(input, true)
end

function inputChanged(input: InputObject)
    Input.InputChanged:Fire(input, true)
end

function inputEnded(input: InputObject)
    Input.KeyDown[input.KeyCode] = nil
    Input.InputEnded:Fire(input, true)
end

function mouseMoved(x: number, y: number)
    Input.X = x
    Input.Y = y
end

SinkFrame.InputBegan:Connect(inputBegan)
SinkFrame.InputChanged:Connect(inputChanged)
SinkFrame.InputEnded:Connect(inputEnded)
SinkFrame.MouseMoved:Connect(mouseMoved)

table.insert(Input._connections, UserInputService.InputBegan:Connect(inputBegan))
table.insert(Input._connections, UserInputService.InputChanged:Connect(inputChanged))
table.insert(Input._connections, UserInputService.InputEnded:Connect(inputEnded))

return Input


local Iris = require(script.Parent.Iris)
local Input = require(script.Parent.UserInputService)

-- Create the plugin toolbar, button and dockwidget for Iris to work in.
local widgetInfo = DockWidgetPluginGuiInfo.new(Enum.InitialDockState.Float, false, false, 200, 300)

local Toolbar = plugin:CreateToolbar("Iris")
local ToggleButton = Toolbar:CreateButton("Toggle Iris", "Toggle Iris running in a plugin window.", "rbxasset://textures/AnimationEditor/icon_checkmark.png")
local IrisWidget = plugin:CreateDockWidgetPluginGuiAsync("IrisWidget", widgetInfo)

-- defein the widget and button properties
IrisWidget.Name = "Iris"
IrisWidget.Title = "Iris"
IrisWidget.ZIndexBehavior = Enum.ZIndexBehavior.Sibling
ToggleButton.ClickableWhenViewportHidden = true

local IrisEnabled = false
Input.SinkFrame.Parent = IrisWidget

-- configure a few things within Iris. We need to provide our own UserInputService and change the config.
Iris.Internal._utility.UserInputService = Input
Iris.UpdateGlobalConfig({
    UseScreenGUIs = false,
})
Iris.Disabled = true

Iris.Init(IrisWidget)

-- We can start defining our code. This just uses the demo window and then forces it to be the same size
-- as the Plugin Widget. You don't have to do it this way.
Iris:Connect(function()
    local window = Iris.ShowDemoWindow()
    window.state.size:set(IrisWidget.AbsoluteSize)
    window.state.position:set(Vector2.zero)
end)

IrisWidget:BindToClose(function()
    IrisEnabled = false
    IrisWidget.Enabled = false
    Iris.Disabled = true
    ToggleButton:SetActive(false)
end)

ToggleButton.Click:Connect(function()
    IrisEnabled = not IrisEnabled
    IrisWidget.Enabled = IrisEnabled
    Iris.Disabled = not IrisEnabled
    ToggleButton:SetActive(IrisEnabled)
end)

-- This is quite important. We need to ensure Iris properly shutdowns and closes any connections.
plugin.Unloading:Connect(function()
    Iris.Shutdown()

    for _, connection in Input._connections do
        connection:Disconnect()
    end

    Input.SinkFrame:Destroy()

    IrisEnabled = false
    IrisWidget.Enabled = false
    Iris.Disabled = true
    ToggleButton:SetActive(false)
end)


local ReplicatedStorage = game:GetService("ReplicatedStorage")
local Types = require(ReplicatedStorage.Iris.PubTypes)

return function(parent: GuiObject)
    local Iris: Types.Iris = require(ReplicatedStorage.Iris)
    local Input = require(script.Parent.UserInputService)

    Input.SinkFrame.Parent = parent

    Iris.Internal._utility.UserInputService = Input
    Iris.UpdateGlobalConfig({
        UseScreenGUIs = false,
    })
    Iris.Internal._utility.GuiOffset = Input.SinkFrame.AbsolutePosition
    Iris.Internal._utility.MouseOffset = Input.SinkFrame.AbsolutePosition
    Input.SinkFrame:GetPropertyChangedSignal("AbsolutePosition"):Connect(function()
        Iris.Internal._utility.GuiOffset = Input.SinkFrame.AbsolutePosition
        Iris.Internal._utility.MouseOffset = Input.SinkFrame.AbsolutePosition
    end)

    Iris.Init(parent)

    -- Actual Iris code here:
    Iris:Connect(Iris.ShowDemoWindow)

    return function()
        Iris.Shutdown()

        for _, connection in Input._connections do
            connection:Disconnect()
        end

        Input.SinkFrame:Destroy()
    end
end


--[=[
	@class Constants
	@private

	Module containing constant values used throughout the Lyra library.
]=]

-- Estimated bytes reserved within the DataStore value for Lyra's internal metadata
-- (like appliedMigrations, orphanedFiles list, and File object overhead when sharded).
-- This ensures that the actual user data shards stay safely within the DataStore limits.
local INTERNAL_SPACE_RESERVE_BYTES = 10_000

--[=[
	@interface Constants
	@within Constants

	.RECORD_SCOPE "lyra/records" -- Scope prefix used for keys storing the main DataStoreRecord objects.

	.TX_SCOPE "lyra/tx" -- Scope prefix used for keys storing transaction status markers in a DataStore.

	.SHARD_SCOPE "lyra/shards" -- Scope prefix used for keys storing data shards in a DataStore.

	.LOCK_SCOPE "lyra/locks" -- Scope prefix used for keys storing lock information in a MemoryStore HashMap.

	.MAX_CHUNK_SIZE number -- Maximum size (in bytes) for a single data shard stored in DataStore. Derived from the Roblox DataStore value limit (4MB) minus a reserved amount (INTERNAL_SPACE_RESERVE_BYTES) for Lyra's internal metadata within the main record.

	.LOCK_REFRESH_INTERVAL_SECONDS number -- How often (in seconds) a held lock should be refreshed in MemoryStore to prevent it from expiring while still actively being used. Should be significantly shorter than LOCK_DURATION_SECONDS.

	.LOCK_DURATION_SECONDS number -- The initial time-to-live (TTL) duration (in seconds) for a lock acquired in MemoryStore. If the lock isn't refreshed within this time, it will automatically expire. Must be longer than LOCK_REFRESH_INTERVAL_SECONDS to allow time for refreshes.

	.AUTOSAVE_INTERVAL_SECONDS number -- How often (in seconds) the automatic saving mechanism should attempt to save dirty session data.
]=]
local Constants = {
	RECORD_SCOPE = "lyra/records",
	TX_SCOPE = "lyra/tx",
	SHARD_SCOPE = "lyra/shards",
	LOCK_SCOPE = "lyra/locks",

	MAX_CHUNK_SIZE = 4_000_000 - INTERNAL_SPACE_RESERVE_BYTES,

	LOCK_REFRESH_INTERVAL_SECONDS = 60,

	LOCK_DURATION_SECONDS = 90,

	AUTOSAVE_INTERVAL_SECONDS = 5 * 60, -- 5 minutes
}

return Constants


--[=[
	@class Log
	@private

	Provides a structured logging implementation for the Lyra library.

	**Design:**
	- **Callback-based:** Instead of directly printing or sending logs, this module
	  uses a callback function (`logCallback`) provided during logger creation. This
	  allows the consuming application to decide how and where log messages are routed
	  (e.g., print to console, send to an external service, store in memory).
	- **Structured Context:** Log messages include a `context` table. Loggers can be
	  `extend`ed with additional context fields, which are automatically merged into
	  every subsequent log message created by that logger instance or its descendants.
	  This helps provide detailed, structured information for debugging and monitoring.
	- **Log Levels:** Supports standard log levels (`fatal`, `error`, `warn`, `info`,
	  `debug`, `trace`). A global log level can be set using `Log.setLevel` to filter
	  out messages below the desired severity.

	**Usage:**
	```lua
	local Log = require(script.Parent.Log)

	-- Set the global minimum log level (optional, defaults to "info")
	Log.setLevel("debug")

	-- Create a logger instance with a callback
	local myLogger = Log.createLogger(function(logMessage)
		print(`[{logMessage.level}] {logMessage.message}`, logMessage.context)
	end, { initialContext = "value" })

	-- Log messages
	myLogger:log("info", "User logged in", { userId = 123 })

	-- Create a logger with extended context
	local sessionLogger = myLogger:extend({ sessionId = "abc" })
	sessionLogger:log("debug", "Session data loaded")
	-- Output will include { initialContext = "value", sessionId = "abc", userId = 123 }
	-- if logged via myLogger, or { initialContext = "value", sessionId = "abc" }
	-- if logged via sessionLogger.
	```
]=]

--[=[
	Represents the different log levels available for logging messages.

	@type LogLevel "fatal" | "error" | "warn" | "info" | "debug" | "trace"
	@tag enum
	@within Log
]=]
export type LogLevel = "fatal" | "error" | "warn" | "info" | "debug" | "trace"

--[=[
	Represents a log message sent to the logger's callback function.

	@interface LogMessage
	@within Log
	.message string -- The main content of the log message.
	.level LogLevel -- The severity level of the message.
	.context { [string]: any }? -- Optional table containing additional structured context.
]=]
export type LogMessage = {
	message: string,
	level: LogLevel,
	context: { [string]: any }?,
}

export type LoggerImpl = {
	__index: LoggerImpl,

	log: (self: Logger, level: LogLevel, message: string, context: { [string]: any }?) -> (),
	extend: (self: Logger, context: { [string]: any }) -> Logger,
}

export type LoggerProps = {
	_logCallback: (logMessage: LogMessage) -> (),
	_context: { [string]: any },
}

export type Logger = typeof(setmetatable({} :: LoggerProps, {} :: LoggerImpl))

-- Ordered list of log levels from most severe to least severe.
local levels = {
	"fatal",
	"error",
	"warn",
	"info",
	"debug",
	"trace",
}

-- The main Log module table, holds global settings like the current log level.
local Log = {
	level = "info", -- Default global log level.
}

--[=[
	Sets the global minimum log level.
	Messages with a severity lower than this level will be ignored by all loggers.

	@within Log
	@param level LogLevel -- The minimum log level to allow.
	@error string -- Throws an error if the provided level is invalid.
]=]
function Log.setLevel(level: LogLevel)
	if table.find(levels, level) == nil then
		error(`Invalid log level: '{level}'`)
	end
	Log.level = level
end

-- Metatable implementing the Logger methods. (Internal)
local Logger: LoggerImpl = {} :: LoggerImpl
Logger.__index = Logger

--[=[
	Logs a message if its level is at or above the globally set log level.

	Merges the provided `context` table with the logger's persistent context
	before calling the configured `_logCallback`.

	@within Log
	@param level LogLevel -- The severity level of the message.
	@param message string -- The log message content.
	@param context { [string]: any }? -- Optional additional context specific to this log call.
]=]
function Logger:log(level: LogLevel, message: string, context: { [string]: any }?)
	-- Filter messages based on the global log level.
	if table.find(levels, level) > table.find(levels, Log.level) then
		return
	end

	-- Merge instance context with call-specific context.
	local finalContext = table.clone(self._context)
	if context then
		for key, value in context do
			finalContext[key] = value
		end
	end

	-- Call the configured log callback within a protected call to catch errors.
	local ok, result = pcall(function()
		self._logCallback({
			level = level,
			message = message,
			context = finalContext,
		})
	end)

	if not ok then
		-- If the logging callback itself errors, print a warning.
		warn(`Error in log callback: {result}`)
	end
end

--[=[
	Creates a new Logger instance that inherits the parent's callback
	but has an extended context.

	The new logger's context is a merged table containing the parent's context
	and the additional `context` provided here.

	@within Log
	@param context { [string]: any } -- The additional context fields to add.
	@return Logger -- A new Logger instance with the extended context.
]=]
function Logger:extend(context: { [string]: any }): Logger
	-- Merge existing context with the new context.
	local finalContext = table.clone(self._context)
	for key, value in context do
		finalContext[key] = value
	end

	-- Create and return a new logger instance sharing the callback but with the new context.
	return setmetatable({
		_logCallback = self._logCallback,
		_context = finalContext,
	}, Logger)
end

--[=[
	Factory function to create a new root Logger instance.

	@within Log
	@param logCallback (logMessage: LogMessage) -> () -- The function that will be called for each log message that passes the level filter. This function receives the complete `LogMessage` object including merged context.
	@param context { [string]: any }? -- Optional initial context for this logger.
	@return Logger -- A new Logger instance.
]=]
function Log.createLogger(logCallback: (logMessage: LogMessage) -> (), context: { [string]: any }?): Logger
	return setmetatable({
		_logCallback = logCallback,
		_context = context or {},
	}, Logger)
end

return Log


--[=[
	@class Migrations
	@private

	Provides functionality for managing and applying data migrations.

	**Concept:** As data schemas evolve over time, older data stored
	in DataStores needs to be updated to match the new structure. This module
	allows defining a series of migration steps (`Types.MigrationStep`) that
	transform data from one version of the schema to the next.

	**Workflow:**
	1. Define migration steps using helpers like `makeAddFieldsStep` or `makeTransformStep`.
	   Each step must have a unique `name`.
	2. Provide the list of steps to the `Store` via `StoreContext`.
	3. When data is loaded (`Session:load`), the `apply` function in this module is called.
	4. `apply` compares the list of all defined steps against the `appliedMigrations`
	   list stored within the `DataStoreRecord`.
	5. It executes the `apply` function of any step that hasn't been applied yet, in the
	   order they are defined.
	6. The `apply` function of each step receives the current data and returns the transformed data.
	7. The names of successfully applied steps are added to the `appliedMigrations` list,
	   which is then saved back to the `DataStoreRecord`.

	**Idempotency:** The system ensures migrations are idempotent (applying them multiple
	times has the same effect as applying them once) by checking the `appliedMigrations` list
	before running a step.

	**Error Handling:** Each step's `apply` function is executed within a `pcall` to catch
	errors. If a step fails, the migration process stops, and the error is propagated.
]=]

local Log = require(script.Parent.Log)
local Promise = require(script.Parent.Promise)
local Tables = require(script.Parent.Tables)
local Types = require(script.Parent.Types)

--[=[
	Alias for a generic data table type used in migrations.

	@type Data { [string]: any }
	@within Migrations
	@private
]=]
type Data = { [string]: any }

--[=[
	Validates that the provided migration steps adhere to the expected structure.

	@within Migrations
	@param steps {Types.MigrationStep} -- The list of migration steps to validate.
	@error string -- Throws an error if any step is malformed.
]=]
local function validate(steps: { Types.MigrationStep })
	assert(typeof(steps) == "table", "steps must be a table")
	for _, step in steps do
		assert(typeof(step) == "table", "step must be a table")
		assert(typeof(step.name) == "string", "step.name must be a string")
		assert(typeof(step.apply) == "function", "step.apply must be a function")
	end
end

--[=[
	Helper function to create a common type of migration step: adding new fields
	with default values to existing data. Uses a deep merge strategy.

	@within Migrations
	@param name string -- The unique name for this migration step.
	@param fields Data -- A table containing the new fields and their default values.
	@return Types.MigrationStep -- A migration step object.
]=]
local function makeAddFieldsStep(name: string, fields: Data): Types.MigrationStep
	return {
		name = name,
		apply = function(data)
			-- Merges the default `fields` into the existing `data`.
			-- Existing keys in `data` are preserved unless they are tables themselves,
			-- in which case they are recursively merged.
			return Tables.mergeDeep(fields, data)
		end,
	}
end

--[=[
	Helper function to create a migration step that applies a custom transformation
	function to the data.

	@within Migrations
	@param name string -- The unique name for this migration step.
	@param transformFunc (currentValue: Data) -> Data -- The function that takes the current data and returns the transformed data.
	@return Types.MigrationStep -- A migration step object.
]=]
local function makeTransformStep(name: string, transformFunc: (currentValue: Data) -> Data): Types.MigrationStep
	return {
		name = name,
		apply = transformFunc,
	}
end

--[=[
	Parameters for the `apply` function.

	@interface ApplyParams
	@within Migrations
	.logger Log.Logger -- Logger instance for logging migration progress and errors.
	.data Data -- The current data loaded from the DataStore record's File.
	.steps {Types.MigrationStep} -- The list of all defined migration steps for the store.
	.appliedMigrations {string} -- The list of names of migration steps already applied to this specific data record, loaded from the DataStoreRecord.
]=]
export type ApplyParams = {
	logger: Log.Logger,
	data: Data,
	steps: { Types.MigrationStep },
	appliedMigrations: { string },
}

--[=[
	Result returned by the `apply` function upon successful completion.

	@interface ApplyResult
	@within Migrations
	.data Data -- The potentially modified data after applying necessary migration steps.
	.appliedMigrations {string} -- The updated list of applied migration names, including any newly applied steps. This should be saved back to the DataStoreRecord.
]=]
export type ApplyResult = {
	data: Data,
	appliedMigrations: { string },
}

--[=[
	Applies pending migration steps to the data.

	Iterates through the defined `steps` and applies any step whose name is not
	present in the `appliedMigrations` list. Ensures idempotency and uses `pcall`
	for safe execution of each step's `apply` function.

	@within Migrations
	@param params ApplyParams -- The parameters for applying migrations.
	@return Promise<ApplyResult> -- A Promise that resolves with the updated data and the new list of applied migration names.
	@error string -- Rejects if any migration step fails during `pcall`.
]=]
local function apply(params: ApplyParams): Promise.TPromise<ApplyResult>
	local logger = params.logger
	local currentData = params.data
	-- Clone the list to avoid modifying the original list passed in params directly within this function scope.
	local appliedMigrations = table.clone(params.appliedMigrations)

	-- Create a set for quick lookup of already applied migrations.
	local appliedSet = {}
	for _, name in appliedMigrations do
		appliedSet[name] = true
	end

	return Promise.new(function(resolve, reject)
		for _, step in params.steps do
			-- Idempotency Check: Only apply if the step name is not in the applied set.
			if not appliedSet[step.name] then
				logger:log("trace", "applying migration step", { stepName = step.name })

				-- Deep copy the data before applying the step to avoid modifying the original
				-- `currentData` reference in case of an error within `step.apply`.
				-- TODO: Do we actually care about atomicity here? This might be a performance hit.
				local staged = Tables.copyDeep(currentData)
				-- Safely execute the migration step's apply function.
				local ok, result = pcall(step.apply, staged)

				if not ok then
					-- If pcall failed, log the error and reject the promise.
					logger:log("error", "failed to apply migration step", { stepName = step.name, error = result })
					return reject(`Failed migration step '{step.name}': {result}`)
				end

				-- If pcall succeeded, update currentData with the result.
				-- Deep copy again to ensure `currentData` holds a distinct copy for the next step.
				-- TODO: We should probably just warn about not modifying the data after returning - this is expensive.
				currentData = Tables.copyDeep(result)
				-- Add the step name to the list of applied migrations for this record.
				table.insert(appliedMigrations, step.name)
				-- Also update the lookup set for the next iteration.
				appliedSet[step.name] = true
			end
		end

		-- All applicable steps applied successfully. Resolve with the final data and updated list.
		return resolve({
			data = currentData,
			appliedMigrations = appliedMigrations,
		})
	end)
end

--[=[
	Utility function to extract just the names from a list of migration steps.

	@within Migrations
	@param migrations {Types.MigrationStep} -- The list of migration steps.
	@return {string} -- A list containing only the names of the migration steps.
]=]
local function getStepNames(migrations: { Types.MigrationStep }): { string }
	local names = {}
	for _, step in migrations do
		table.insert(names, step.name)
	end
	return names
end

return {
	makeAddFieldsStep = makeAddFieldsStep,
	makeTransformStep = makeTransformStep,
	validate = validate,
	apply = apply,
	getStepNames = getStepNames,
}


local HttpService = game:GetService("HttpService")

local Tables = require(script.Parent.Tables)

local MAX_QUEUE_SIZE = 30
local RATE_LIMITS = {
	GetAsync = { Base = 60, PlayerMultiplier = 10 },
	SetAsync = { Base = 60, PlayerMultiplier = 10 },
	UpdateAsync = { Base = 60, PlayerMultiplier = 10 },
	RemoveAsync = { Base = 60, PlayerMultiplier = 10 },
	GetVersion = { Base = 5, PlayerMultiplier = 2 },
}

local MAX_KEY_LENGTH = 50
local MAX_DATA_SIZE = 4_194_304
local MAX_METADATA_SIZE = 300
local CACHE_DURATION = 4

local function getLatencyForRequest(service: any, requestType: string): number
	if service._latencyByRequestType[requestType] then
		return service._latencyByRequestType[requestType]
	end

	if service._randomLatencyRange then
		local min = service._randomLatencyRange.min
		local max = service._randomLatencyRange.max
		return min + math.random() * (max - min)
	end

	return service._globalLatency or 0.1
end

local function shouldSimulateError(service: any, requestType: string): (boolean, string?)
	if #service.queuedErrors[requestType] > 0 then
		return true, table.remove(service.queuedErrors[requestType], 1)
	end

	if service._errorRates[requestType] and math.random() < service._errorRates[requestType] then
		return true, "DataStore request failed due to random error simulation"
	end

	if service._simulateThrottling then
		return true, "DataStore request was throttled, try again later"
	end

	if service._simulateTimeout then
		return true, "DataStore request timed out"
	end

	return false, nil
end

local function doLatency(service: any, requestType: string)
	local latency = getLatencyForRequest(service, requestType)
	if latency > 0 then
		task.wait(latency)
	end
end

local function getNow()
	return DateTime.now().UnixTimestampMillis
end

local function msToHours(ms: number)
	return math.floor(ms / (60 * 60 * 1000))
end

local function getUtcHour()
	return msToHours(getNow())
end

local function createVersion(value: any, userIds: any, metadata: any)
	return {
		Value = value,
		UserIds = userIds or {},
		Metadata = metadata or {},
		CreatedTime = getNow(),
		Version = tostring(getNow()),
	}
end

local function createMockKeyInfo(userIds: any, metadata: any)
	local info = {
		CreatedTime = getNow(),
		UpdatedTime = getNow(),
		Version = tostring(getNow()),
		userIds = userIds or {},
		metadata = metadata or {},
		GetUserIds = function(self)
			return self.userIds
		end,
		GetMetadata = function(self)
			return self.metadata
		end,
	}
	return info
end

local function validateMetadata(metadata: any)
	if metadata == nil then
		return
	end

	if typeof(metadata) ~= "table" then
		error("DataStoreService: Metadata must be a table", 0)
	end

	local serialized = game:GetService("HttpService"):JSONEncode(metadata)
	if #serialized > MAX_METADATA_SIZE then
		error(`DataStoreService: Metadata size exceeds {MAX_METADATA_SIZE} limit`, 0)
	end
end

local function getScopedData(self: any)
	local datastore = self._service.mockData[self.datastoreName]
	if not datastore then
		datastore = {}
		self._service.mockData[self.datastoreName] = datastore
	end

	local scopedData = datastore[self.scope]
	if not scopedData then
		scopedData = {}
		datastore[self.scope] = scopedData
	end

	return scopedData
end

local function getVersionsForKey(self: any, key: string)
	local versionKey = `{self.datastoreName}|{self.scope}|{key}`
	if not self._service.mockVersions[versionKey] then
		self._service.mockVersions[versionKey] = {}
	end
	return self._service.mockVersions[versionKey]
end

local function getCacheKey(self: any, key: string)
	return `{self.datastoreName}|{self.scope}|{key}`
end

local function getCachedValue(self: any, key: string)
	local cacheKey = getCacheKey(self, key)
	local cached = self._service.mockCache[cacheKey]
	if cached and (os.time() - cached.timestamp) < CACHE_DURATION then
		return cached.value, cached.keyInfo
	end
	return nil, nil :: any
end

local function setCachedValue(self: any, key: string, value: any, keyInfo: any)
	local cacheKey = getCacheKey(self, key)
	self._service.mockCache[cacheKey] = {
		value = value,
		keyInfo = keyInfo,
		timestamp = os.time(),
	}
end

local function enforceKeyLimits(key: any)
	if typeof(key) ~= "string" then
		error("DataStoreService: Key must be a string", 0)
	end

	if #key == 0 then
		error("DataStoreService: Key name can't be empty", 0)
	end

	if #key > MAX_KEY_LENGTH then
		error(`DataStoreService: Key name exceeds the {MAX_KEY_LENGTH} character limit`, 0)
	end
end

local function enforceDataLimits(value: any)
	if value == nil then
		return
	end

	local success, result = pcall(function()
		return game:GetService("HttpService"):JSONEncode(value)
	end)
	if not success then
		error(`DataStoreService: Cannot store {typeof(value)} in DataStore`, 0)
	end

	if #result > MAX_DATA_SIZE then
		error(`DataStoreService: Serialized value exceeds {MAX_DATA_SIZE} limit`, 0)
	end
end

local function doRateLimitAndQueue(self: any, requestType: string)
	local queue = self._service.requestQueues[requestType]
	local limitData = RATE_LIMITS[requestType]
	if not limitData then
		limitData = RATE_LIMITS.GetAsync
	end

	local budget = self._service._customBudgets[requestType]
	if not budget then
		local currentPlayers = 15
		budget = limitData.Base + (currentPlayers * limitData.PlayerMultiplier)
	end

	if #queue >= budget then
		return false
	end

	if #queue >= MAX_QUEUE_SIZE then
		error(
			`DataStoreService:{requestType}() request dropped. Request was throttled, but throttled request queue was full`,
			0
		)
	end

	if self._service._forcedThrottles[requestType] then
		local throttleInfo = self._service._forcedThrottles[requestType]
		if throttleInfo.endTime > os.time() then
			error("DataStoreService: Request is currently throttled", 0)
		else
			self._service._forcedThrottles[requestType] = nil
		end
	end

	table.insert(queue, true)
	doLatency(self._service, requestType)
	local fn = table.remove(queue, 1)

	local shouldError, errorMessage = shouldSimulateError(self._service, requestType)
	if shouldError then
		error(errorMessage, 0)
	end

	if typeof(fn) == "function" then
		fn()
	end

	return true
end

local function sanitize(value: any)
	return Tables.copyDeep(value)
end

local MockDataStore = {}

function MockDataStore:SetAsync(key, value, userIds, options)
	value = sanitize(value)
	enforceKeyLimits(key)
	enforceDataLimits(value)
	if options then
		validateMetadata(options:GetMetadata())
	end

	doRateLimitAndQueue(self, "SetAsync")

	local scopedData = getScopedData(self)
	local versions = getVersionsForKey(self, key)
	local currentHour = getUtcHour()

	if #versions == 0 or msToHours(versions[#versions].CreatedTime) ~= currentHour then
		table.insert(versions, createVersion(value, userIds, options and options:GetMetadata()))
	else
		local currentVersion = versions[#versions]
		currentVersion.Value = value
		if userIds then
			currentVersion.UserIds = userIds
		end
		if options then
			currentVersion.Metadata = options:GetMetadata() or {}
		end
	end

	local latestVersion = versions[#versions].Version
	scopedData[key] = value
	local keyInfo = createMockKeyInfo(userIds, options and options:GetMetadata())
	self._service.mockKeyInfo[`{self.datastoreName}|{self.scope}|{key}`] = keyInfo

	setCachedValue(self, key, value, keyInfo)

	return latestVersion
end

function MockDataStore:GetAsync(key, options)
	enforceKeyLimits(key)

	if options == nil or options.UseCache == true then
		local cachedValue, cachedKeyInfo = getCachedValue(self, key)
		if cachedValue ~= nil then
			return cachedValue, cachedKeyInfo
		end
	end

	doRateLimitAndQueue(self, "GetAsync")

	local scopedData = getScopedData(self)
	local data = scopedData[key]
	local keyInfo = self._service.mockKeyInfo[`{self.datastoreName}|{self.scope}|{key}`]

	if data ~= nil then
		setCachedValue(self, key, data, keyInfo)
	end

	return data, keyInfo
end

function MockDataStore:UpdateAsync(key, transformFunction)
	enforceKeyLimits(key)

	doRateLimitAndQueue(self, "UpdateAsync")

	while true do
		local scopedData = getScopedData(self)
		local keyInfo = self._service.mockKeyInfo[`{self.datastoreName}|{self.scope}|{key}`] or createMockKeyInfo()
		local currentValue = scopedData[key]
		local valueForCallback = currentValue ~= nil and Tables.copyDeep(currentValue)

		local newValue, newUserIds, newMetadata = transformFunction(valueForCallback, keyInfo)
		if newValue == nil then
			return currentValue, keyInfo
		end
		newValue = sanitize(newValue)

		enforceDataLimits(newValue)
		validateMetadata(newMetadata)

		if not Tables.equalsDeep(currentValue, scopedData[key]) then
			continue
		end

		local versions = getVersionsForKey(self, key)
		local currentHour = getUtcHour()

		if #versions == 0 or versions[#versions].CreatedTime ~= currentHour then
			table.insert(
				versions,
				createVersion(newValue, newUserIds or keyInfo:GetUserIds(), newMetadata or keyInfo:GetMetadata())
			)
		else
			local currentVersion = versions[#versions]
			currentVersion.Value = newValue
			if newUserIds then
				currentVersion.UserIds = newUserIds
			end
			if newMetadata then
				currentVersion.Metadata = newMetadata
			end
		end

		scopedData[key] = newValue
		if newUserIds ~= nil then
			keyInfo.userIds = newUserIds
		end
		if newMetadata ~= nil then
			keyInfo.metadata = newMetadata
		end
		keyInfo.UpdatedTime = getNow()
		keyInfo.Version = tostring(getNow())

		setCachedValue(self, key, newValue, keyInfo)
		return newValue, keyInfo
	end
end

function MockDataStore:RemoveAsync(key)
	enforceKeyLimits(key)

	doRateLimitAndQueue(self, "RemoveAsync")

	local scopedData = getScopedData(self)
	local data = scopedData[key]
	local keyInfo = self._service.mockKeyInfo[`{self.datastoreName}|{self.scope}|{key}`]

	if data ~= nil then
		local versions = getVersionsForKey(self, key)
		table.insert(versions, createVersion(nil, nil, nil))
	end

	scopedData[key] = nil

	local cacheKey = getCacheKey(self, key)
	self._service.mockCache[cacheKey] = nil

	return data, keyInfo
end

function MockDataStore:GetVersionAsync(key, version)
	enforceKeyLimits(key)

	doRateLimitAndQueue(self, "GetVersion")

	local versions = getVersionsForKey(self, key)
	for _, v in versions do
		if v.Version == version then
			return v.Value, createMockKeyInfo(v.UserIds, v.Metadata)
		end
	end
	return nil, nil :: any
end

function MockDataStore:ListVersionsAsync(key, sortDirection, minDate, maxDate)
	enforceKeyLimits(key)

	doRateLimitAndQueue(self, "GetVersion")

	local versions = getVersionsForKey(self, key)
	local results = {}

	for _, v in versions do
		if (not minDate or v.CreatedTime >= minDate) and (not maxDate or v.CreatedTime <= maxDate) then
			table.insert(results, {
				Version = v.Version,
				CreatedTime = v.CreatedTime,
			})
		end
	end

	if sortDirection == Enum.SortDirection.Descending then
		table.sort(results, function(a, b)
			return a.CreatedTime > b.CreatedTime
		end)
	else
		table.sort(results, function(a, b)
			return a.CreatedTime < b.CreatedTime
		end)
	end

	return {
		IsFinished = true,
		GetCurrentPage = function()
			return results
		end,
	}
end

local function createMockStore(service: any, datastoreName: string, scope: string?)
	return setmetatable({
		datastoreName = datastoreName,
		scope = scope or "global",
		_service = service,
	}, service.mockStoreMeta)
end

local function new()
	local mockStoreMeta = table.clone(MockDataStore)
	mockStoreMeta.__index = mockStoreMeta

	local tbl = {
		dataStores = {},
		mockData = {},
		mockKeyInfo = {},
		mockVersions = {},
		mockCache = {},
		requestQueues = {
			GetAsync = {},
			SetAsync = {},
			UpdateAsync = {},
			RemoveAsync = {},
			GetVersion = {},
		},
		queuedErrors = {
			GetAsync = {},
			SetAsync = {},
			UpdateAsync = {},
			RemoveAsync = {},
			GetVersion = {},
		},
		mockStoreMeta = mockStoreMeta,

		_globalLatency = nil,
		_latencyByRequestType = {},
		_randomLatencyRange = nil,
		_errorRates = {},
		_simulateThrottling = false,
		_simulateTimeout = false,
		_customBudgets = {},
		_forcedThrottles = {},
	}

	tbl.GetDataStore = function(self, datastoreName: string, scope: string?)
		local store = self.dataStores[`{datastoreName}{scope or ""}`]
		if not store then
			store = createMockStore(self, datastoreName, scope)
			self.dataStores[`{datastoreName}{scope or ""}`] = store
		end
		return store
	end

	tbl.GetRequestBudgetForRequestType = function(self, requestType: Enum.DataStoreRequestType)
		local requestTypeStr = requestType.Name

		local queueKey = requestTypeStr
		if requestTypeStr == "GetVersionAsync" then
			queueKey = "GetVersion"
		end
		if requestTypeStr == "GetIncrementAsync" then
			queueKey = "GetAsync"
		end
		if requestTypeStr == "SetIncrementAsync" then
			queueKey = "SetAsync"
		end

		if self._customBudgets[queueKey] then
			return self._customBudgets[queueKey]
		end

		local limitData = RATE_LIMITS[queueKey]
		if not limitData then
			limitData = RATE_LIMITS.GetAsync
		end

		local currentPlayers = 15
		local totalBudget = limitData.Base + (currentPlayers * limitData.PlayerMultiplier)

		local currentQueueSize = #self.requestQueues[queueKey]
		return math.max(0, totalBudget - currentQueueSize)
	end

	return (tbl :: any) :: DataStoreService
end

return {
	new = new,

	mockStoreMethod = function(service, method: string, jest: any)
		local real = service.mockStoreMeta[method]
		local mock, fn = jest.fn(real)
		service.mockStoreMeta[method] = fn
		return mock, real
	end,

	setGlobalLatency = function(service: any, seconds: number)
		service._globalLatency = seconds
	end,
	setLatencyForRequestType = function(service: any, requestType: string, seconds: number)
		service._latencyByRequestType[requestType] = seconds
	end,
	setRandomLatency = function(service: any, minSeconds: number, maxSeconds: number)
		service._randomLatencyRange = { min = minSeconds, max = maxSeconds }
	end,
	queueError = function(service: any, requestType: string, errorMessage: string, count: number?)
		for _ = 1, count or 1 do
			table.insert(service.queuedErrors[requestType], errorMessage)
		end
	end,
	setErrorRate = function(service: any, requestType: string, rate: number)
		service._errorRates[requestType] = rate
	end,
	simulateThrottling = function(service: any, enabled: boolean)
		service._simulateThrottling = enabled
	end,
	simulateTimeout = function(service: any, enabled: boolean)
		service._simulateTimeout = enabled
	end,
	setRequestBudget = function(service: any, requestType: string, budget: number)
		service._customBudgets[requestType] = budget
	end,
	forceThrottle = function(service: any, requestType: string, duration: number)
		service._forcedThrottles[requestType] = {
			endTime = os.time() + duration,
		}
	end,
	snapshot = function(service: any)
		return HttpService:JSONEncode({
			mockData = service.mockData,
			mockKeyInfo = service.mockKeyInfo,
			mockVersions = service.mockVersions,
			mockCache = service.mockCache,
		})
	end,
	restore = function(service: any, snapshot: string)
		local decoded = HttpService:JSONDecode(snapshot)
		service.mockData = decoded.mockData
		service.mockKeyInfo = decoded.mockKeyInfo
		service.mockVersions = decoded.mockVersions
		service.mockCache = decoded.mockCache
	end,
}


local HttpService = game:GetService("HttpService")

local Tables = require(script.Parent.Tables)

local MAX_QUEUE_SIZE = 30
local RATE_LIMITS = {
	GetAsync = { Base = 100, PlayerMultiplier = 10 },
	SetAsync = { Base = 100, PlayerMultiplier = 10 },
	UpdateAsync = { Base = 100, PlayerMultiplier = 10 },
	RemoveAsync = { Base = 100, PlayerMultiplier = 10 },
}

local MAX_KEY_LENGTH = 128
local MAX_VALUE_SIZE = 32 * 1024 -- 32KB
local MAX_UPDATE_RETRIES = 3
local MAX_EXPIRATION = 3888000 -- 45 days in seconds

local function getNow()
	return DateTime.now().UnixTimestampMillis
end

local function getLatencyForRequest(service: any, requestType: string): number
	if service._latencyByRequestType[requestType] then
		return service._latencyByRequestType[requestType]
	end

	if service._randomLatencyRange then
		local min = service._randomLatencyRange.min
		local max = service._randomLatencyRange.max
		return min + math.random() * (max - min)
	end

	return service._globalLatency or 0.1
end

local function shouldSimulateError(service: any, requestType: string): (boolean, string?)
	if #service.queuedErrors[requestType] > 0 then
		return true, table.remove(service.queuedErrors[requestType], 1)
	end

	if service._errorRates[requestType] and math.random() < service._errorRates[requestType] then
		return true, "InternalError: random error simulation"
	end

	if service._simulateThrottling then
		return true, "Throttled: MemoryStore request was throttled, try again later"
	end

	if service._simulateTimeout then
		return true, "Timeout: MemoryStore request timed out"
	end

	return false, nil
end

local function doLatency(service: any, requestType: string)
	local latency = getLatencyForRequest(service, requestType)
	if latency > 0 then
		task.wait(latency)
	end
end

local function doRateLimitAndQueue(service: any, requestType: string)
	local queue = service.requestQueues[requestType]
	local limitData = RATE_LIMITS[requestType]
	local budget = service._customBudgets[requestType]
	if not budget then
		local currentPlayers = 15
		budget = limitData.Base + (currentPlayers * limitData.PlayerMultiplier)
	end

	if #queue >= budget then
		error("TotalRequestsOverLimit", 0)
	end

	if #queue >= MAX_QUEUE_SIZE then
		error("RequestThrottled", 0)
	end

	if service._forcedThrottles[requestType] then
		local throttleInfo = service._forcedThrottles[requestType]
		if throttleInfo.endTime > os.time() then
			error("MemoryStore request is currently throttled", 0)
		else
			service._forcedThrottles[requestType] = nil
		end
	end

	table.insert(queue, true)
	doLatency(service, requestType)
	local fn = table.remove(queue, 1)

	local shouldError, errorMessage = shouldSimulateError(service, requestType)
	if shouldError then
		error(errorMessage, 0)
	end

	if typeof(fn) == "function" then
		fn()
	end
end

local function enforceKeyLimits(key: string)
	if typeof(key) ~= "string" then
		error("InvalidRequest: key must be a string", 0)
	end
	if #key == 0 then
		error("InvalidRequest: key cannot be empty", 0)
	end
	if #key > MAX_KEY_LENGTH then
		error("InvalidRequest: key length exceeds limit", 0)
	end
end

local function enforceValueLimits(value: any)
	if value == nil then
		return
	end

	local success, encoded = pcall(function()
		return game:GetService("HttpService"):JSONEncode(value)
	end)

	if not success then
		error("InvalidRequest: value cannot be encoded to JSON", 0)
	end

	if #encoded > MAX_VALUE_SIZE then
		error("ItemValueSizeTooLarge", 0)
	end
end

local function enforceExpirationLimits(expiration: number)
	if typeof(expiration) ~= "number" then
		error("InvalidRequest: expiration must be a number", 0)
	end
	if expiration < 0 then
		error("InvalidRequest: expiration must be >= 0", 0)
	end
	if expiration > MAX_EXPIRATION then
		error("InvalidRequest: expiration exceeds maximum", 0)
	end
end

local function getScopedData(self: any)
	local hashMap = self.service.mockData[self.name]
	if not hashMap then
		hashMap = {}
		self.service.mockData[self.name] = hashMap
	end
	return hashMap
end

local function isExpired(entry: any): boolean
	if not entry or not entry.expiration then
		return true
	end
	return getNow() > entry.expiration
end

local function sanitize(value: any): any
	return Tables.copyDeep(value)
end

local MockHashMap = {}

function MockHashMap:GetAsync(key: string): any?
	enforceKeyLimits(key)

	doRateLimitAndQueue(self.service, "GetAsync")

	local scopedData = getScopedData(self)
	local entry = scopedData[key]

	if isExpired(entry) then
		scopedData[key] = nil
		return nil
	end

	return entry.value
end

function MockHashMap:SetAsync(key: string, value: any, expiration: number): boolean
	value = sanitize(value)
	enforceKeyLimits(key)
	enforceValueLimits(value)
	enforceExpirationLimits(expiration)

	doRateLimitAndQueue(self.service, "SetAsync")

	local scopedData = getScopedData(self)
	local serialized = HttpService:JSONEncode(value)
	if serialized == nil then
		return false
	end

	scopedData[key] = {
		value = value,
		expiration = getNow() + expiration * 1000,
	}

	return true
end

function MockHashMap:UpdateAsync(key: string, transformFunction: (any?) -> any?, expiration: number): any?
	enforceKeyLimits(key)
	enforceExpirationLimits(expiration)

	doRateLimitAndQueue(self.service, "UpdateAsync")

	local scopedData = getScopedData(self)
	local retries = 0

	while retries < MAX_UPDATE_RETRIES do
		local entry = scopedData[key]
		local currentValue = nil

		if not isExpired(entry) then
			currentValue = Tables.copyDeep(entry.value)
		end

		local success, newValue = pcall(transformFunction, currentValue)
		if not success then
			error("TransformCallbackFailed", 0)
		end

		if newValue == nil then
			return nil
		end
		newValue = sanitize(newValue)

		enforceValueLimits(newValue)

		if entry and not isExpired(entry) and not Tables.equalsDeep(entry.value, currentValue) then
			retries += 1
			if retries >= MAX_UPDATE_RETRIES then
				error("DataUpdateConflict", 0)
			end
			continue
		end

		scopedData[key] = {
			value = newValue,
			expiration = getNow() + expiration * 1000,
		}

		return newValue
	end

	error("UpdateConflict: Exceeded max number of retries", 0)
end

function MockHashMap:RemoveAsync(key: string)
	enforceKeyLimits(key)

	doRateLimitAndQueue(self.service, "RemoveAsync")

	local scopedData = getScopedData(self)
	scopedData[key] = nil
end

local function createMockHashMap(service: any, name: string): any
	return setmetatable({
		service = service,
		name = name,
	}, { __index = service.mockHashMapMeta })
end

local function new(): MemoryStoreService
	local mockHashMapMeta = table.clone(MockHashMap)
	mockHashMapMeta.__index = mockHashMapMeta

	local tbl = {
		mockData = {},
		hashMaps = {},
		mockHashMapMeta = mockHashMapMeta,
		requestQueues = {
			GetAsync = {},
			SetAsync = {},
			UpdateAsync = {},
			RemoveAsync = {},
		},
		queuedErrors = {
			GetAsync = {},
			SetAsync = {},
			UpdateAsync = {},
			RemoveAsync = {},
		},

		_globalLatency = nil,
		_latencyByRequestType = {},
		_randomLatencyRange = nil,
		_errorRates = {},
		_simulateThrottling = false,
		_simulateTimeout = false,
		_customBudgets = {},
		_forcedThrottles = {},
	}

	tbl.GetHashMap = function(self, name: string): any
		if not self.hashMaps[name] then
			self.hashMaps[name] = createMockHashMap(self, name)
		end
		return self.hashMaps[name]
	end

	return (tbl :: any) :: MemoryStoreService
end

return {
	new = new,

	mockHashMapMethod = function(service: any, method: string, jest: any)
		local real = service.mockHashMapMeta[method]
		local mock, fn = jest.fn(real)
		service.mockHashMapMeta[method] = fn
		return mock, real
	end,

	setGlobalLatency = function(service: any, seconds: number)
		service._globalLatency = seconds
	end,
	setLatencyForRequestType = function(service: any, requestType: string, seconds: number)
		service._latencyByRequestType[requestType] = seconds
	end,
	setRandomLatency = function(service: any, minSeconds: number, maxSeconds: number)
		service._randomLatencyRange = { min = minSeconds, max = maxSeconds }
	end,
	queueError = function(service: any, requestType: string, errorMessage: string, count: number?)
		for _ = 1, count or 1 do
			table.insert(service.queuedErrors[requestType], errorMessage)
		end
	end,
	setErrorRate = function(service: any, requestType: string, rate: number)
		service._errorRates[requestType] = rate
	end,
	simulateThrottling = function(service: any, enabled: boolean)
		service._simulateThrottling = enabled
	end,
	simulateTimeout = function(service: any, enabled: boolean)
		service._simulateTimeout = enabled
	end,
	setRequestBudget = function(service: any, requestType: string, budget: number)
		service._customBudgets[requestType] = budget
	end,
	forceThrottle = function(service: any, requestType: string, duration: number)
		service._forcedThrottles[requestType] = {
			endTime = os.time() + duration,
		}
	end,

	snapshot = function(service: any)
		return HttpService:JSONEncode({
			mockData = service.mockData,
		})
	end,
	restore = function(service: any, snapshot: string)
		local decoded = HttpService:JSONDecode(snapshot)
		service.mockData = decoded.mockData
	end,
}


--[=[
	A PlayerStore wraps a regular Store to provide a more convenient API for working with Player data.
	It automatically converts Players to UserId keys and handles player kicks on data errors.

	```lua
	local playerStore = PlayerStore.create({
		name = "PlayerData",
		template = {
			coins = 0,
			items = {},
		},
		schema = function(data)
			return typeof(data.coins) == "number" and typeof(data.items) == "table",
			"Invalid data format"
		end,
	})

	-- Load data when player joins
	Players.PlayerAdded:Connect(function(player)
		playerStore:loadAsync(player)
	end)

	-- Unload data when player leaves
	Players.PlayerRemoving:Connect(function(player)
		playerStore:unloadAsync(player)
	end)
	```

	@class PlayerStore
]=]

local Players = game:GetService("Players")

local Log = require(script.Parent.Log)
local Promise = require(script.Parent.Promise)
local Store = require(script.Parent.Store)
local Types = require(script.Parent.Types)

type PlayerStoreImpl<T> = {
	__index: PlayerStoreImpl<T>,
	_kickPlayer: (self: PlayerStore<T>, keyOrPlayer: string | Player, message: string) -> (),

	get: (self: PlayerStore<T>, player: Player) -> Promise.TPromise<T>,
	load: (self: PlayerStore<T>, player: Player) -> Promise.Promise,
	unload: (self: PlayerStore<T>, player: Player) -> Promise.Promise,
	update: (
		self: PlayerStore<T>,
		player: Player,
		transformFunction: (data: T) -> boolean
	) -> Promise.TPromise<boolean>,
	updateImmutable: (
		self: PlayerStore<T>,
		player: Player,
		transformFunction: (data: T) -> T | false
	) -> Promise.TPromise<boolean>,
	tx: (
		self: PlayerStore<T>,
		players: { Player },
		transformFunction: (state: { [Player]: T }) -> boolean
	) -> Promise.TPromise<boolean>,
	txImmutable: (
		self: PlayerStore<T>,
		players: { Player },
		transformFunction: (state: { [Player]: T }) -> { [Player]: T } | false
	) -> Promise.TPromise<boolean>,
	save: (self: PlayerStore<T>, player: Player) -> Promise.Promise,
	close: (self: PlayerStore<T>) -> Promise.Promise,
	peek: (self: PlayerStore<T>, userId: number) -> T,

	getAsync: (self: PlayerStore<T>, player: Player) -> T,
	loadAsync: (self: PlayerStore<T>, player: Player) -> (),
	unloadAsync: (self: PlayerStore<T>, player: Player) -> (),
	updateAsync: (self: PlayerStore<T>, player: Player, transformFunction: (data: T) -> boolean) -> boolean,
	updateImmutableAsync: (
		self: PlayerStore<T>,
		player: Player,
		transformFunction: (data: T) -> T | false
	) -> boolean,
	txAsync: (
		self: PlayerStore<T>,
		players: { Player },
		transformFunction: (state: { [Player]: T }) -> boolean
	) -> (),
	txImmutableAsync: (
		self: PlayerStore<T>,
		players: { Player },
		transformFunction: (state: { [Player]: T }) -> { [Player]: T } | false
	) -> boolean,
	saveAsync: (self: PlayerStore<T>, player: Player) -> (),
	closeAsync: (self: PlayerStore<T>) -> (),
	peekAsync: (self: PlayerStore<T>, userId: number) -> T,
}

type PlayerStoreProps<T> = {
	_store: Store.Store<T>,
}

--[=[
	Configuration for creating a new Store.

	@interface PlayerStoreConfig
	.name string -- The name of the store
	.template T -- The template data for new keys
	.schema (value: any) -> (boolean, string?) -- A function to validate data
	.migrationSteps { MigrationStep }? -- Optional migration steps
	.importLegacyData ((key: string) -> any?)? -- Optional function to import legacy data
	.changedCallbacks { (key: string, newData: T, oldData: T?) -> () }? -- Optional callbacks for data changes
	.logCallback ((logMessage: LogMessage) -> ())? -- Optional callback for log messages
	.memoryStoreService MemoryStoreService? -- Optional MemoryStoreService instance for mocking
	.dataStoreService DataStoreService? -- Optional DataStoreService instance for mocking
	.useMock boolean? -- If true (and running in Studio), uses mock in-memory implementations of DataStoreService and MemoryStoreService instead of the actual Roblox services. Useful for testing in a controlled environment. Defaults to `false`.

	@within PlayerStore
]=]
type PlayerStoreConfig<T> = {
	name: string,
	template: T,
	useDSSLocking: boolean?,
	schema: (value: any) -> (boolean, string?),
	migrationSteps: { Types.MigrationStep }?,
	importLegacyData: ((key: string) -> any?)?,
	changedCallbacks: { (key: string, newData: T, oldData: T?) -> () }?,
	logCallback: ((logMessage: Log.LogMessage) -> ())?,
	memoryStoreService: MemoryStoreService?,
	dataStoreService: DataStoreService?,
	useMock: boolean?,
}

type PlayerStore<T> = typeof(setmetatable({} :: PlayerStoreProps<T>, {} :: PlayerStoreImpl<T>))

local function getUserIdKey(player: Player): string
	return tostring(player.UserId)
end

local PlayerStore: PlayerStoreImpl<any> = {} :: PlayerStoreImpl<any>
PlayerStore.__index = PlayerStore

--[=[
	Creates a new PlayerStore with the given configuration.
	Configuration is similar to Store.createStore, but automatically adds player kick handling.

	```lua
	local playerStore = PlayerStore.create({
		name = "PlayerData",
		template = { coins = 0 },
		schema = function(data)
			return typeof(data.coins) == "number", "coins must be a number"
		end,

		-- Optional: Runs whenever data changes
		changedCallbacks = {
			function(key, newData, oldData)
				print(key, "changed from", oldData.coins, "to", newData.coins)
			end,
		},
	})
	```

	Players will be automatically kicked with an error message if:
	- Their data fails to load
	- The DataStore lock is lost during their session

	@param config PlayerStoreConfig<T> -- Configuration for the store
	@return PlayerStore<T>
	@within PlayerStore
]=]
local function createPlayerStore<T>(config: PlayerStoreConfig<T>): PlayerStore<T>
	local self: PlayerStore<T>

	local storeConfig: Store.StoreConfig<T> = {
		name = config.name,
		template = config.template,
		schema = config.schema,
		useDSSLocking = config.useDSSLocking,
		migrationSteps = config.migrationSteps,
		importLegacyData = config.importLegacyData,
		changedCallbacks = config.changedCallbacks,
		logCallback = config.logCallback,
		onLockLost = function(key: string)
			self:_kickPlayer(key, "DataStore lock lost, please rejoin the game.")
		end,
		memoryStoreService = config.memoryStoreService,
		dataStoreService = config.dataStoreService,
		useMock = config.useMock,
	}

	local store = Store.createStore(storeConfig)
	self = setmetatable({ _store = store }, PlayerStore)
	return self
end

--[=[
	Internal helper to kick players when data errors occur.

	@within PlayerStore
	@private
]=]
function PlayerStore:_kickPlayer(keyOrPlayer: string | Player, message: string): ()
	if typeof(keyOrPlayer) ~= "string" then
		keyOrPlayer:Kick(message)
	else
		local userId = tonumber(keyOrPlayer)
		if userId ~= nil then
			local player = Players:GetPlayerByUserId(userId)
			if player ~= nil then
				player:Kick(message)
			end
		end
	end
end

--[=[
	Gets the current data for the given player.

	```lua
	playerStore:get(player):andThen(function(data)
		print(player.Name, "has", data.coins, "coins")
	end)
	```

	@error "Key not loaded" -- The player's data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@return Promise<T> -- Resolves with the player's data
	@within PlayerStore
]=]
function PlayerStore:get<T>(player: Player): Promise.TPromise<T>
	local userIdKey = getUserIdKey(player)
	return self._store:get(userIdKey)
end

--[=[
	Syntactic sugar for `playerStore:get(player):expect()`.

	See [PlayerStore:get]
	@yields
]=]
function PlayerStore:getAsync<T>(player: Player): Promise.TPromise<T>
	return self:get(player):expect()
end

--[=[
	Loads data for the given player. Must be called before using other methods.

	```lua
	playerStore:load(player):andThen(function()
		print("Data loaded for", player.Name)
	end)
	```

	:::caution
	If loading fails, the player will be kicked from the game.
	:::

	@error "Load already in progress" -- Another load is in progress for this player
	@error "Store is closed" -- The store has been closed
	@return Promise -- Resolves when data is loaded
	@within PlayerStore
]=]
function PlayerStore:load(player: Player)
	local userIdKey = getUserIdKey(player)
	return self._store:load(userIdKey, { player.UserId }):catch(function(e)
		self:_kickPlayer(player, "DataStore load failed, please rejoin the game.")
		return Promise.reject(e)
	end)
end

--[=[
	Syntactic sugar for `playerStore:load(player):expect()`.

	See [PlayerStore:load]
	@yields
]=]
function PlayerStore:loadAsync(player: Player)
	return self:load(player):expect()
end

--[=[
	Unloads data for the given player.

	```lua
	playerStore:unload(player):andThen(function()
		print("Data unloaded for", player.Name)
	end)
	```

	@error "Store is closed" -- The store has been closed
	@return Promise<boolean> -- Resolves when the update is complete, with a boolean indicating success
	@within PlayerStore
]=]
function PlayerStore:unload(player: Player): Promise.TPromise<boolean>
	local userIdKey = getUserIdKey(player)
	return self._store:unload(userIdKey)
end

--[=[
	Syntactic sugar for `playerStore:unload(player):expect()`.

	See [PlayerStore:unload]
	@yields
]=]
function PlayerStore:unloadAsync(player: Player)
	return self:unload(player):expect()
end

--[=[
	Updates data for the given player using a transform function.
	The transform function must return true to commit changes, or false to abort.

	```lua
	playerStore:update(player, function(data)
		if data.coins < 100 then
			data.coins += 50
			return true -- Commit changes
		end
		return false -- Don't commit changes
	end)
	```

	@error "Key not loaded" -- The player's data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@error "Schema validation failed" -- The transformed data failed schema validation
	@return Promise<boolean> -- Resolves when the update is complete, with the boolean returned by the transformFunction
	@within PlayerStore
]=]
function PlayerStore:update<T>(player: Player, transformFunction: (data: T) -> boolean): Promise.TPromise<boolean>
	local userIdKey = getUserIdKey(player)
	return self._store:update(userIdKey, transformFunction)
end

--[=[
	Syntactic sugar for `playerStore:update(player, transformFunction):expect()`.

	See [PlayerStore:update]
	@yields
]=]
function PlayerStore:updateAsync<T>(player: Player, transformFunction: (data: T) -> boolean): boolean
	return self:update(player, transformFunction):expect()
end

--[=[
	Updates data for the given player using a transform function that does not mutate the original data.
	The transform function must return the new data or false to abort.

	```lua
	playerStore:updateImmutable(player, function(data)
		if data.coins < 100 then
			return { coins = data.coins + 50 } -- Return new data to commit changes
		end
		return false -- Don't commit changes
	end)
	```

	@error "Key not loaded" -- The player's data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@error "Schema validation failed" -- The transformed data failed schema validation
	@return Promise<boolean> -- Resolves when the update is complete, with the boolean returned by the transformFunction
	@within PlayerStore
]=]
function PlayerStore:updateImmutable<T>(
	player: Player,
	transformFunction: (data: T) -> T | false
): Promise.TPromise<boolean>
	local userIdKey = getUserIdKey(player)
	return self._store:updateImmutable(userIdKey, transformFunction)
end

--[=[
	Syntactic sugar for `playerStore:updateImmutable(player, transformFunction):expect()`.

	See [PlayerStore:updateImmutable]
	@yields
]=]
function PlayerStore:updateImmutableAsync<T>(player: Player, transformFunction: (data: T) -> T | false): boolean
	return self:updateImmutable(player, transformFunction):expect()
end

--[=[
	Performs a transaction across multiple players' data atomically.
	All players' data must be loaded first. Either all changes apply or none do.

	```lua
	playerStore:tx({player1, player2}, function(state)
		-- Transfer coins between players
		if state[player1].coins >= 100 then
			state[player1].coins -= 100
			state[player2].coins += 100
			return true -- Commit transaction
		end
		return false -- Abort transaction
	end)
	```

	@error "Key not loaded" -- One or more players' data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@error "Schema validation failed" -- The transformed data failed schema validation
	@return Promise<boolean> -- Resolves with `true` if the transaction was successful, or `false` if it was aborted. Rejects on error.
	@within PlayerStore
]=]
function PlayerStore:tx<T>(
	players: { Player },
	transformFunction: (state: { [Player]: T }) -> boolean
): Promise.TPromise<boolean>
	local userIdKeys = table.create(#players) :: { string }
	local userIdKeyToPlayer: { [string]: Player } = {}
	for i, player in players do
		local key = getUserIdKey(player)
		userIdKeys[i] = key
		userIdKeyToPlayer[key] = player
	end

	local function wrapped(state: { [string]: any }): boolean
		local userIdState: { [Player]: T } = {}
		for key, value in state do
			local player = userIdKeyToPlayer[key]
			userIdState[player] = value
		end

		local success = transformFunction(userIdState)
		if success == false then
			return false
		end

		for key in state do
			state[key] = nil
		end

		for playerVariant, value in userIdState do
			local stringKey = getUserIdKey(playerVariant)
			state[stringKey] = value
		end

		return success
	end

	return self._store:tx(userIdKeys, wrapped)
end

--[=[
	Syntactic sugar for `playerStore:tx(players, transformFunction):expect()`.

	See [PlayerStore:tx]
	@yields
]=]
function PlayerStore:txAsync<T>(players: { Player }, transformFunction: (state: { [Player]: T }) -> boolean): boolean
	return self:tx(players, transformFunction):expect()
end

--[=[
	Performs a transaction across multiple players' data atomically using immutable updates.
	All players' data must be loaded first. Either all changes apply or none do.

	```lua
	playerStore:txImmutable({player1, player2}, function(state)
		-- Transfer coins between players
		if state[player1].coins >= 100 then
			return {
				[player1] = { coins = state[player1].coins - 100 },
				[player2] = { coins = state[player2].coins + 100 },
			} -- Commit transaction with new data
		end
		return false -- Abort transaction
	end)
	```

	@error "Key not loaded" -- One or more players' data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@error "Schema validation failed" -- The transformed data failed schema validation
	@return Promise<boolean> -- Resolves with `true` if the transaction was successful, or `false` if it was aborted. Rejects on error.
	@within PlayerStore
]=]
function PlayerStore:txImmutable<T>(
	players: { Player },
	transformFunction: (state: { [Player]: T }) -> { [Player]: T } | false
): Promise.TPromise<boolean>
	local userIdKeys = table.create(#players) :: { string }
	local userIdKeyToPlayer: { [string]: Player } = {}
	for i, player in players do
		local key = getUserIdKey(player)
		userIdKeys[i] = key
		userIdKeyToPlayer[key] = player
	end

	local function wrapped(state: { [string]: any }): { [string]: any } | false
		local userIdState: { [Player]: T } = {}
		for key, value in state do
			local player = userIdKeyToPlayer[key]
			userIdState[player] = value
		end

		local newData = transformFunction(userIdState)
		if newData == false then
			return false
		end

		for key in state do
			state[key] = nil
		end

		for playerVariant, value in newData :: { [Player]: T } do
			local stringKey = getUserIdKey(playerVariant)
			state[stringKey] = value
		end

		return state
	end

	return self._store:txImmutable(userIdKeys, wrapped)
end

--[=[
	Syntactic sugar for `playerStore:txImmutable(players, transformFunction):expect()`.

	See [PlayerStore:txImmutable]
	@yields
]=]
function PlayerStore:txImmutableAsync<T>(
	players: { Player },
	transformFunction: (state: { [Player]: T }) -> { [Player]: T } | false
): boolean
	return self:txImmutable(players, transformFunction):expect()
end

--[=[
	Forces an immediate save of the given player's data.

	:::info
	Data is automatically saved periodically, so manual saves are usually unnecessary.
	:::

	@error "Key not loaded" -- The player's data hasn't been loaded
	@error "Store is closed" -- The store has been closed
	@return Promise -- Resolves when the save is complete
	@within PlayerStore
]=]
function PlayerStore:save(player: Player)
	local userIdKey = getUserIdKey(player)
	return self._store:save(userIdKey)
end

--[=[
	Syntactic sugar for `playerStore:save(player):expect()`.

	See [PlayerStore:save]
	@yields
]=]
function PlayerStore:saveAsync(player: Player)
	return self:save(player):expect()
end

--[=[
	Closes the store and unloads all active sessions.
	The store cannot be used after closing.

	@return Promise -- Resolves when the store is closed
	@within PlayerStore
]=]
function PlayerStore:close()
	return self._store:close()
end

--[=[
	Syntactic sugar for `playerStore:close():expect()`.

	See [PlayerStore:close]
	@yields
]=]
function PlayerStore:closeAsync()
	return self:close():expect()
end

--[=[
	Returns the current data for the given key without loading it into the store.

	```lua
	playerStore:peek(userId):andThen(function(data)
		print("Current coins:", data.coins)
	end)
	```

	@return Promise<T> -- Resolves with the current data
	@within PlayerStore
]=]
function PlayerStore:peek(userId: number)
	local userIdKey = tostring(userId)
	return self._store:peek(userIdKey)
end

--[=[
	Syntactic sugar for `playerStore:peek(userId):expect()`.

	See [PlayerStore:peek]
	@yields
]=]
function PlayerStore:peekAsync(userId: number)
	return self:peek(userId):expect()
end

return {
	createPlayerStore = createPlayerStore,
}