local function NewStateBag(Element)
    return setmetatable({}, {
        __index = function(_, key)
            if (key == 'set') then
                return function(_, k, value, Sync, ChangesPolicy)
                    if (localPlayer) then 
                        setElementData(Element, k, value, Sync)
                    else
                        setElementData(Element, k, value, Sync, ChangesPolicy)
                    end
                    return value
                end
            elseif (key == 'get') then
                return function(_, k)
                    return getElementData(Element, k)
                end
            end

            return getElementData(Element, key)
        end,
        
        __newindex = function(_, key, value)
            setElementData(Element, key, value)
        end
    })
end

local StateBagsMT = {
    Entity = {
        __index = function(t, key)
            if (key == 'state') then
                return NewStateBag(t.__element)
            end

            return nil
        end,
        
        __newindex = function(t, key, value)
            if (key == 'state') then
                error('Cannot set state directly, use Entity(element).state.key = value', 2)
            end
        end
    },
    Player = {
        __index = function(t, key)
            if (key == 'state') then
                return NewStateBag(t.__element)
            end

            return nil
        end,
    
        __newindex = function(t, key, value)
            if (key == 'state') then
                error('Cannot set state directly, use Player(element).state.key = value', 2)
            end
        end
    },
}

function Entity(Element)
    if (type(Element) ~= 'userdata' or not isElement(Element)) then
        error('Entity() requires a valid element', 2)
    end
    
    return setmetatable({
        __element = Element
    }, StateBagsMT.Entity)
end

function _Player(Player)
    if (type(Player) ~= 'userdata' or not isElement(Player)) then
        error('Player() requires a valid player', 2)
     end
    
    return setmetatable({
        __element = Player
    }, StateBagsMT.Player)
end

if isClient or isElement(localPlayer) then
    LocalPlayer = _Player(localPlayer)
end

return true