-- InfStore.lua
-- boatbomber

-- A module to have DataStores hold an inf size dictionary
-- by automagically chunking the data behind the scenes

-- Example:
-- local store = InfStore.new("Global_v0.1.0", "Tutorials")
-- store:Get()
-- store:Add("UniqueTutorialId", TutorialData)
--
-- local submissions = InfStore.new(plr.UserId .. "_v0.1.0", "Submissions")
-- submissions:Add("UniqueSubmissionId", SubmissionData)

local DataStoreService = game:GetService("DataStoreService")
local HttpService = game:GetService("HttpService")

local LIMIT = 2_000_000 -- ~2Mb
local VERBOSE = true

local function verCall(callback)
	-- Callback instead of print string so when VERBOSE is false, expensive computations are skipped entirely
	if VERBOSE then
		callback()
	end
end

local module = {}
local existingStores = {}

function module.new(identifier: string, name: string)
	if existingStores[identifier] == nil then
		existingStores[identifier] = {}
	elseif existingStores[identifier][name] then
		return existingStores[identifier][name]
	end

	local InfStore = {
		_store = DataStoreService:GetDataStore(identifier),
		_name = name,
		_chunks = {},
		_cached = nil,
	}
	existingStores[identifier][name] = InfStore

	local debugId = identifier.."|"..name
	verCall(function()
		print("Created InfStore:", debugId)
	end)

	function InfStore:Get(skipCache: boolean?)
		if (not skipCache) and (self._cached ~= nil) then
			return self._cached
		end

		local result = {}

		local i = 0
		while true do
			i += 1

			-- Get next chunk
			local chunk = self:GetChunk(i, skipCache)

			-- Exit if we've reached the end
			if chunk == nil then
				break
			end

			-- Copy the chunk's values into our result
			for k, v in pairs(chunk) do
				result[k] = v
			end

			verCall(function()
				print(debugId, string.format("  :Get() chunk %d (%.2f Mb)", i, #HttpService:JSONEncode(chunk)/1024/1024))
			end)
		end

		verCall(function()
			print(debugId, string.format(":Get() returned a combined %d chunks (%.2f Mb)", i-1, #HttpService:JSONEncode(result)/1024/1024))
		end)

		self._cached = result
		return result
	end

	function InfStore:GetChunk(chunk: number, skipCache: boolean?)
		-- Used for manual lazy loading to get chunks on demand instead of all initially
		if (not skipCache) and (self._chunks[chunk]) then
			return self._chunks[chunk]
		else
			local c = self._store:GetAsync(self._name .. chunk)
			self._chunks[chunk] = c
			return c
		end
	end

	function InfStore:SetChunk(chunk: number, value: {})
		verCall(function()
			warn(debugId, string.format(":SetChunk() overwriting chunk %d to a %.2f Mb table", chunk, #HttpService:JSONEncode(value)/1024/1024))
		end)
		-- Used for overwriting stores
		self._store:SetAsync(self._name .. chunk, value)
	end

	function InfStore:Add(key: string, value: any)
		local cost = #HttpService:JSONEncode(value)

		local i = 0
		while true do
			i += 1

			-- Try next chunk
			local complete = false
			self._store:UpdateAsync(self._name .. i, function(chunk)
				-- New chunk
				if chunk == nil then
					complete = true
					return {
						[key] = value,
					}
				end

				-- No space in this chunk
				if LIMIT - #HttpService:JSONEncode(chunk) <= cost then
					return nil
				end

				-- Add to chunk
				complete = true
				chunk[key] = value
				self._chunks[i] = chunk

				verCall(function()
					print(debugId, string.format(":Add() inserted a value (%.3f Kb) into chunk %d (%.2f Mb) at ['%s']", cost/1024, i, #HttpService:JSONEncode(chunk)/1024/1024, key))
				end)
				return chunk
			end)

			if complete then
				self._cached = nil
				break
			end

		end

	end

	function InfStore:Remove(key: string)
		local i = 0
		while true do
			i += 1

			local complete, empty = false, false
			self._store:UpdateAsync(self._name .. i, function(chunk)
				-- Reached last chunk, exit
				if chunk == nil then
					complete = true
					return nil
				end

				-- Chunk doesn't contain key, skip
				if chunk[key] == nil then
					return nil
				end

				-- Remove key from chunk
				chunk[key] = nil
				self._chunks[i] = chunk

				verCall(function()
					print(debugId, ":Remove() cleared ['" .. key .. "'] from chunk " .. i)
				end)

				if next(chunk) == nil then
					-- Removing this key has emptied the chunk
					empty = true
				end

				return chunk
			end)

			-- Move chunks down if this one got emptied by the removal
			if empty then
				-- Clear this chunk
				self._store:RemoveAsync(self._name .. i)
				self._chunks[i] = nil

				-- Push down subsequent chunks to fill the whole
				local n = 0
				while true do
					n += 1
					-- Get next chunk
					local chunk = self._store:GetAsync(self._name .. (i+n))

					-- Exit if we've reached the end
					if chunk == nil then
						break
					end

					-- Move chunk down
					local a, b = (i+n-1), (i+n)

					self._store:SetAsync(self._name .. a, chunk)
					self._chunks[a] = chunk

					self._store:RemoveAsync(self._name .. b)
					self._chunks[b] = nil
				end
			end

			-- Exit once we've done the last chunk
			if complete then
				self._cached = nil
				break
			end
		end
	end

	function InfStore:RemoveBulk(keys: {string})
		local i = 0
		while true do
			i += 1

			local complete, empty = false, false
			self._store:UpdateAsync(self._name .. i, function(chunk)
				-- Reached last chunk, exit
				if chunk == nil then
					complete = true
					return nil
				end

				-- Remove keys from chunk
				local changed = false
				for _, key in keys do
					if chunk[key] ~= nil then
						chunk[key] = nil
						changed = true
					end
				end
				self._chunks[i] = chunk

				-- If we didn't change anything, skip
				if not changed then
					return nil
				end

				if next(chunk) == nil then
					-- Removing this key has emptied the chunk
					empty = true
				end

				return chunk
			end)

			-- Move chunks down if this one got emptied by the removal
			if empty then
				-- Clear this chunk
				self._store:RemoveAsync(self._name .. i)
				self._chunks[i] = nil

				-- Push down subsequent chunks to fill the whole
				local n = 0
				while true do
					n += 1
					-- Get next chunk
					local chunk = self._store:GetAsync(self._name .. (i+n))

					-- Exit if we've reached the end
					if chunk == nil then
						break
					end

					-- Move chunk down
					local a, b = (i+n-1), (i+n)

					self._store:SetAsync(self._name .. a, chunk)
					self._chunks[a] = chunk

					self._store:RemoveAsync(self._name .. b)
					self._chunks[b] = nil
				end
			end

			-- Exit once we've done the last chunk
			if complete then
				self._cached = nil
				break
			end
		end

		verCall(function()
			print(debugId, ":RemoveBulk() cleared {" .. table.concat(keys, ", ") .. "} from the store")
		end)
	end

	function InfStore:Replace(key: string, value: any)
		local i = 0
		while true do
			i += 1

			local complete = false
			self._store:UpdateAsync(self._name .. i, function(chunk)
				-- Empty chunk, exit
				if chunk == nil then
					complete = true
					return nil
				end

				-- Doesn't have our key, skip
				if chunk[key] == nil then
					return nil
				end

				-- Update key in chunk
				chunk[key] = value
				complete = true
				self._chunks[i] = chunk

				verCall(function()
					print(debugId, ":Replace() replaced value for ['" .. key .. "'] in chunk", i)
				end)

				return chunk
			end)

			-- Exit if we've reached the end
			if complete then
				self._cached = nil
				break
			end
		end
	end

	function InfStore:Update(key: string, transformer: (any) -> any?)
		local i = 0
		while true do
			i += 1

			local complete = false
			self._store:UpdateAsync(self._name .. i, function(chunk)
				-- Empty chunk, exit
				if chunk == nil then
					complete = true
					return nil
				end

				-- Doesn't have our key, skip
				if chunk[key] == nil then
					return nil
				end

				-- Update key in chunk
				local success, replace = pcall(transformer, chunk[key])
				if success and replace then
					chunk[key] = replace
					self._chunks[i] = chunk
				end
				complete = true

				if success and replace then -- Only use datastore budget if we actually changed chunk
					verCall(function()
						print(debugId, ":Update() transformed value for ['" .. key .. "'] in chunk", i)
					end)
					return chunk
				else
					return nil
				end
			end)

			-- Exit if we've reached the end
			if complete then
				self._cached = nil
				break
			end
		end
	end

	-- Aliases for API style consistency in various projects
	InfStore.Retrieve = InfStore.Get
	InfStore.Pull = InfStore.Get
	InfStore.Set = InfStore.Add
	InfStore.Insert = InfStore.Add
	InfStore.Push = InfStore.Add
	InfStore.Delete = InfStore.Remove
	InfStore.Patch = InfStore.Replace
	InfStore.Transform = InfStore.Update

	return InfStore
end

return module