--!native
--!optimize 2

local ChunkManager = {}

local RunService = game:GetService("RunService")

local Chunk = require("./ChunkManager/Chunk")
local BlockManager = require("./ChunkManager/BlockManager")
local ChunkBuilder = require("./ChunkManager/ChunkBuilder")
local Util = require("./Util")
local Globals = require(script.Parent:WaitForChild("Globals"))

local remote = game:GetService("ReplicatedStorage"):WaitForChild("RecieveChunkPacket")
local tickremote = game:GetService("ReplicatedStorage"):WaitForChild("Tick")
local Players = game:GetService("Players")

local ChunkFolder = Instance.new("Folder")
ChunkFolder.Name = "$blockscraft_client"

ChunkManager.ChunkFolder = ChunkFolder

local CHUNK_RADIUS = Globals.RenderDistance or 5
local LOAD_BATCH = Globals.LoadBatch or 8
local RESYNC_INTERVAL = Globals.ResyncInterval or 5
local RESYNC_RADIUS = Globals.ResyncRadius or 2
local DEBUG_RESYNC = false
local FORCELOAD_CHUNKS = {
	{0, 1, 0}
}

local unloadingChunks = {}
local pendingChunkRequests = {}

local lastChunkKey: string? = nil
local lastHeavyTick = 0
local HEAVY_TICK_INTERVAL = 1.5
local lastUnloadSweep = 0
local UNLOAD_SWEEP_INTERVAL = 3 -- slower sweep cadence
local MAX_LOADED_CHUNKS = 0
local SPAWN_CHUNK_KEY: string? = nil
local playerFrozen = false
local storedMovementState = nil

local function worldToChunkCoord(v: number): number
	return math.floor((v + 16) / 32)
end

local CHUNK_OFFSETS = {}
do
	for y = -CHUNK_RADIUS, CHUNK_RADIUS do
		for x = -CHUNK_RADIUS, CHUNK_RADIUS do
			for z = -CHUNK_RADIUS, CHUNK_RADIUS do
				table.insert(CHUNK_OFFSETS, {x, y, z, (x * x) + (y * y) + (z * z)})
			end
		end
	end
	table.sort(CHUNK_OFFSETS, function(a, b)
		return a[4] < b[4]
	end)
	MAX_LOADED_CHUNKS = math.max(1, math.floor(#CHUNK_OFFSETS * 2)) -- tighter cap than full render cube
	if FORCELOAD_CHUNKS[1] then
		local forced = FORCELOAD_CHUNKS[1]
		SPAWN_CHUNK_KEY = `{forced[1]},{forced[2]},{forced[3]}`
	end
end

function ChunkManager:UnloadAllNow()
	for key, chunk in pairs(Chunk.AllChunks) do
		unloadingChunks[key] = true
		pcall(function()
			if chunk.loaded then
				chunk:UnloadImmediate()
			end
		end)
		pcall(function()
			chunk:Destroy()
		end)
		Chunk.AllChunks[key] = nil
		unloadingChunks[key] = nil
	end
end

local function Swait(l)
	for _ = 1, l do
		RunService.Stepped:Wait()
	end
end

local function setCharacterFrozen(shouldFreeze: boolean)
	local player = Players.LocalPlayer
	if not player then
		return
	end
	local character = player.Character
	if not character then
		return
	end
	local humanoid = character:FindFirstChildOfClass("Humanoid")
	local root = character:FindFirstChild("HumanoidRootPart")
	if not humanoid or not root then
		return
	end
	if shouldFreeze == playerFrozen then
		return
	end
	if shouldFreeze then
		if not storedMovementState then
			storedMovementState = {
				walkSpeed = humanoid.WalkSpeed,
				autoRotate = humanoid.AutoRotate,
			}
			if humanoid.UseJumpPower then
				storedMovementState.jumpPower = humanoid.JumpPower
			else
				storedMovementState.jumpHeight = humanoid.JumpHeight
			end
		end
		humanoid.AutoRotate = false
		humanoid.WalkSpeed = 0
		if humanoid.UseJumpPower then
			humanoid.JumpPower = 0
		else
			humanoid.JumpHeight = 0
		end
		root.Anchored = true
	else
		root.Anchored = false
		if storedMovementState then
			humanoid.AutoRotate = storedMovementState.autoRotate
			humanoid.WalkSpeed = storedMovementState.walkSpeed
			if humanoid.UseJumpPower and storedMovementState.jumpPower then
				humanoid.JumpPower = storedMovementState.jumpPower
			elseif storedMovementState.jumpHeight then
				humanoid.JumpHeight = storedMovementState.jumpHeight
			end
		end
		storedMovementState = nil
	end
	playerFrozen = shouldFreeze
end

local function getLocalPlayerChunkPos()
	local player = Players.LocalPlayer
	if not player then
		return nil
	end
	local character = player.Character
	if not character then
		return nil
	end
	local root = character:FindFirstChild("HumanoidRootPart")
	if not root then
		return nil
	end
	local pos = root.Position
	return {
		x = worldToChunkCoord(pos.X),
		y = worldToChunkCoord(pos.Y),
		z = worldToChunkCoord(pos.Z)
	}
end

local function isWithinRenderDistance(chunkPos: Vector3, centerChunkPos): boolean
	if not centerChunkPos then
		return false
	end
	return math.abs(chunkPos.X - centerChunkPos.x) <= CHUNK_RADIUS
		and math.abs(chunkPos.Y - centerChunkPos.y) <= CHUNK_RADIUS
		and math.abs(chunkPos.Z - centerChunkPos.z) <= CHUNK_RADIUS
end

local function shouldSkipUnload(key: string): boolean
	return SPAWN_CHUNK_KEY ~= nil and key == SPAWN_CHUNK_KEY
end

local function scheduleChunkUnload(key: string, chunk)
	if not chunk or unloadingChunks[key] then
		return
	end
	unloadingChunks[key] = true
	task.defer(function()
		chunk:Unload()
		chunk:Destroy()
		Chunk.AllChunks[key] = nil
		unloadingChunks[key] = nil
	end)
end

local function evictOutOfRangeChunks(centerChunkPos)
	if not centerChunkPos then
		return
	end
	local loadedCount = 0
	for key, loadedChunk in pairs(Chunk.AllChunks) do
		if loadedChunk.loaded and not shouldSkipUnload(key) then
			local inRange = isWithinRenderDistance(loadedChunk.pos, centerChunkPos)
			if inRange then
				loadedCount += 1
				if loadedCount > MAX_LOADED_CHUNKS then
					scheduleChunkUnload(key, loadedChunk)
				end
			else
				scheduleChunkUnload(key, loadedChunk)
			end
		end
	end
end

function ChunkManager:GetChunk(x, y, z)
	local key = `{x},{y},{z}`
	if Chunk.AllChunks[key] then
		return Chunk.AllChunks[key]
	end

	if pendingChunkRequests[key] then
		while pendingChunkRequests[key] do
			task.wait()
		end
		return Chunk.AllChunks[key]
	end

	pendingChunkRequests[key] = true
	local ok, data = pcall(function()
		return remote:InvokeServer(x, y, z)
	end)
	if not ok then
		data = {}
	end
	local ch = Chunk.from(x, y, z, data)
	Chunk.AllChunks[key] = ch
	pendingChunkRequests[key] = nil
	return ch
end

local function ensureNeighboringChunksLoaded(x, y, z)
	local offsets = {
		{1, 0, 0}, {-1, 0, 0},
		{0, 1, 0}, {0, -1, 0},
		{0, 0, 1}, {0, 0, -1}
	}

	for _, offset in ipairs(offsets) do
		local nx, ny, nz = x + offset[1], y + offset[2], z + offset[3]
		local neighbor = ChunkManager:GetChunk(nx, ny, nz)
		if neighbor then
			neighbor:Tick()
		end
	end
end

function ChunkManager:LoadChunk(x, y, z)
	local key = `{x},{y},{z}`
	if unloadingChunks[key] or not Chunk.AllChunks[key] or Chunk.AllChunks[key].loaded then
		return
	end

	unloadingChunks[key] = true
	task.defer(function()
		ensureNeighboringChunksLoaded(x, y, z)

		local chunk = Chunk.AllChunks[key]
		if not chunk then
			chunk = ChunkManager:GetChunk(x, y, z)
			Chunk.AllChunks[key] = chunk
		end

		local instance = ChunkBuilder:BuildChunk(chunk, ChunkFolder)
		chunk.instance = instance
		chunk.loaded = true
		unloadingChunks[key] = nil
		evictOutOfRangeChunks(getLocalPlayerChunkPos())
	end)
end

function ChunkManager:RefreshChunk(x, y, z)
	local key = `{x},{y},{z}`
	local chunk = Chunk.AllChunks[key]
	if not chunk then
		pendingChunkRequests[key] = nil
		ChunkManager:GetChunk(x, y, z)
		ChunkManager:LoadChunk(x, y, z)
		return
	end

	local ok, newData = pcall(function()
		return remote:InvokeServer(x, y, z)
	end)
	if not ok or typeof(newData) ~= "table" then
		if DEBUG_RESYNC then
			warn("[CHUNKMANAGER][RESYNC] Failed to fetch chunk data", key, ok, typeof(newData))
		end
		return
	end

	local function sameState(a, b)
		if a == b then
			return true
		end
		if not a or not b then
			return false
		end
		local count = 0
		for k, v in pairs(a) do
			count += 1
			if b[k] ~= v then
				return false
			end
		end
		for _ in pairs(b) do
			count -= 1
		end
		return count == 0
	end

	local function sameBlockData(a, b)
		if a == b then
			return true
		end
		if not a or not b then
			return false
		end
		if a.id ~= b.id then
			return false
		end
		return sameState(a.state, b.state)
	end

	local changed = 0
	local removed = 0
	for keyStr, newBlock in pairs(newData) do
		local oldBlock = chunk.data[keyStr]
		if not sameBlockData(oldBlock, newBlock) then
			chunk.data[keyStr] = newBlock
			local coords = Util.BlockPosStringToCoords(keyStr)
			chunk:PropogateChanges(coords.X, coords.Y, coords.Z, newBlock)
			changed += 1
		end
	end

	for keyStr in pairs(chunk.data) do
		if not newData[keyStr] then
			chunk.data[keyStr] = nil
			local coords = Util.BlockPosStringToCoords(keyStr)
			chunk:PropogateChanges(coords.X, coords.Y, coords.Z, 0)
			removed += 1
		end
	end
	if chunk.loaded and chunk.instance then
		local pruned = 0
		for _, child in ipairs(chunk.instance:GetChildren()) do
			if not newData[child.Name] then
				pruned += 1
				child:Destroy()
			end
		end
		if DEBUG_RESYNC and pruned > 0 then
			print("[CHUNKMANAGER][RESYNC] Pruned ghost instances", key, "count", pruned)
		end
	end
	if DEBUG_RESYNC and (changed > 0 or removed > 0) then
		print("[CHUNKMANAGER][RESYNC] Applied diff", key, "changed", changed, "removed", removed)
	end
end

function ChunkManager:ForceTick()
	for _, coords in ipairs(FORCELOAD_CHUNKS) do
		local key = `{coords[1]},{coords[2]},{coords[3]}`
		local chunk = Chunk.AllChunks[key]
		if not chunk then
			ChunkManager:LoadChunk(coords[1], coords[2], coords[3])
		else
			chunk:Tick()
		end
	end
end

function ChunkManager:TickI()
	for key, chunk in pairs(Chunk.AllChunks) do
		if tick() - chunk.inhabitedTime <= 5 then
			tickremote:FireServer(key)
		end
	end
end

function ChunkManager:Tick()
	ChunkManager:ForceTick()
	local player = Players.LocalPlayer
	if not player.Character then
		return
	end

	local pos = player.Character:GetPivot().Position
	local chunkPos = {
		x = worldToChunkCoord(pos.X),
		y = worldToChunkCoord(pos.Y),
		z = worldToChunkCoord(pos.Z)
	}
	local ck = `{chunkPos.x},{chunkPos.y},{chunkPos.z}`
	local currentChunk = Chunk.AllChunks[ck]
	local now = tick()
	local shouldHeavyTick = (ck ~= lastChunkKey) or (now - lastHeavyTick >= HEAVY_TICK_INTERVAL)
	lastChunkKey = ck
	if shouldHeavyTick then
		lastHeavyTick = now
	end

	setCharacterFrozen(not (currentChunk and currentChunk.loaded))

	if shouldHeavyTick then
		task.defer(function()
			local processed = 0
			for _, offset in ipairs(CHUNK_OFFSETS) do
				local cx, cy, cz = chunkPos.x + offset[1], chunkPos.y + offset[2], chunkPos.z + offset[3]
				local chunk = ChunkManager:GetChunk(cx, cy, cz)
				if chunk then
					chunk.inhabitedTime = now
					if not chunk.loaded then
						ChunkManager:LoadChunk(cx, cy, cz)
						processed += 1
						if processed % LOAD_BATCH == 0 then
							Swait(1)
						end
					end
				end
			end
		end)
	else
		if currentChunk then
			currentChunk.inhabitedTime = now
		end
	end
	
	--[[
	task.defer(function()
		for y = 0, 2 do
			task.defer(function()
				for x = -CHUNK_RADIUS, CHUNK_RADIUS do
					for z = -CHUNK_RADIUS, CHUNK_RADIUS do
						local cx, cy, cz = chunkPos.x + x, y, chunkPos.z + z
						local key = `{cx},{cy},{cz}`
						local chunk = ChunkManager:GetChunk(cx, cy, cz)
						chunk.inhabitedTime = tick()
						if not chunk.loaded then
							ChunkManager:LoadChunk(cx, cy, cz)
							Swait(2)
						end
					end
				end
			end)
			Swait(10)
		end
	end)
	--]]

	if now - lastUnloadSweep >= UNLOAD_SWEEP_INTERVAL then
		lastUnloadSweep = now
		for key, loadedChunk in pairs(Chunk.AllChunks) do
			if now - loadedChunk.inhabitedTime > 30 and not unloadingChunks[key] and not shouldSkipUnload(key) then -- keep chunks around longer before unloading
				scheduleChunkUnload(key, loadedChunk)
			end
		end
	end
end

function ChunkManager:ResyncAroundPlayer(radius: number)
	local player = game:GetService("Players").LocalPlayer
	if not player.Character then
		return
	end
	local pos = player.Character:GetPivot().Position
	local chunkPos = {
		x = worldToChunkCoord(pos.X),
		y = worldToChunkCoord(pos.Y),
		z = worldToChunkCoord(pos.Z)
	}
	for y = -radius, radius do
		for x = -radius, radius do
			for z = -radius, radius do
				local cx, cy, cz = chunkPos.x + x, chunkPos.y + y, chunkPos.z + z
				ChunkManager:RefreshChunk(cx, cy, cz)
			end
		end
	end
end

function ChunkManager:ResyncAroundChunk(cx: number, cy: number, cz: number, radius: number)
	for y = -radius, radius do
		for x = -radius, radius do
			for z = -radius, radius do
				ChunkManager:RefreshChunk(cx + x, cy + y, cz + z)
			end
		end
	end
end

function ChunkManager:Init()
	if not RunService:IsClient() then
		error("ChunkManager:Init can only be called on the client")
	end

	ChunkFolder.Parent = game:GetService("Workspace")
	ChunkManager:ForceTick()

	tickremote.OnClientEvent:Connect(function(m)
		if m == "U_ALL" then
			ChunkManager:UnloadAllNow()
		end
	end)

	task.defer(function()
		while true do
			wait(2)
			ChunkManager:TickI()
		end
	end)

	task.defer(function()
		while true do
			task.defer(function()
				local success, err = pcall(function()
					ChunkManager:Tick()
				end)
				if not success then
					warn("[CHUNKMANAGER]", err)
				end
			end)
			Swait(20)
		end
	end)

	task.defer(function()
		while true do
			wait(RESYNC_INTERVAL)
			local ok, err = pcall(function()
				ChunkManager:ResyncAroundPlayer(RESYNC_RADIUS)
			end)
			if not ok then
				warn("[CHUNKMANAGER][RESYNC]", err)
			end
		end
	end)
end

return ChunkManager