-- Copyright ©2024-2025 by Todd Hundersmarck (ThundR) 
-- All Rights Reserved

--[[

Unauthorized use and/or distribution of this work entitles myself, the author, to unlimited free and unrestricted use, access, and distribution of any works related to the unauthorized user and/or distributor.

--]]

local g_thModPath = g_currentModDirectory
THDensityMapManager = {}
local THDensityMapManager_mt = Class(THDensityMapManager)
THDensityMapManager.CLASS_NAME = "THDensityMapManager"
local THInfoLayer = {}
local THInfoLayer_mt = Class(THInfoLayer)
local function initScript()
    local self = THDensityMapManager.new()
    if self ~= nil then
        _G.g_thDensityMapManager = self
        THUtils.prependFunction("DensityMapHeightManager", "saveCollisionMap", false, false, self, THDensityMapManager.prepend_saveCollisionMap)
        THUtils.prependFunction("DensityMapHeightManager", "prepareSaveCollisionMap", false, false, self, THDensityMapManager.prepend_prepareSaveCollisionMap)
        THUtils.overwriteFunction("DensityMapHeightManager", "savePreparedCollisionMap", false, false, self, THDensityMapManager.overwrite_savePreparedCollisionMap)
    end
end
function THDensityMapManager.new(customMt)
    customMt = customMt or THDensityMapManager_mt
    if THUtils.argIsValid(type(customMt) == THValueType.TABLE, "customMt", customMt) then
        local self = setmetatable({}, customMt)
        self.infoLayers = {
            byName = {},
            byIndex = {},
            byBitVectorId = {},
            toSave = {},
            toSync = {}
        }
        if g_thEventManager ~= nil then
            g_thEventManager:addEventListener(self)
        end
        return self
    end
end
function THDensityMapManager.onInitTerrain(self, mission, terrainNode, mapI3DFilename)
    local dmSyncer = mission.densityMapSyncer
    local numInfoLayersToSync = #self.infoLayers.toSync
    if dmSyncer ~= nil and numInfoLayersToSync > 0 then
        for layerIndex = 1, numInfoLayersToSync do
            local layerData = self.infoLayers.toSync[layerIndex]
            if layerData.desc.bitVectorId ~= nil and layerData.desc.bitVectorId > 0 then
                dmSyncer:addDensityMap(layerData.desc.bitVectorId)
            end
        end
    end
end
function THDensityMapManager.onDelete(self)
    local numInfoLayers = #self.infoLayers.byIndex
    if numInfoLayers > 0 then
        for layerIndex = 1, numInfoLayers do
            local layerData = self.infoLayers.byIndex[layerIndex]
            if layerData ~= nil then
                layerData:delete()
                layerData.desc.bitVectorId = nil
                layerData.desc.isResetting = false
                layerData.desc.isSavePrepared = false
            end
        end
    end
    for _, value in pairs(self.infoLayers) do
        if type(value) == THValueType.TABLE then
            THUtils.clearTable(value)
        end
    end
end
function THDensityMapManager.preparedInfoLayerSaveCallback(_, ...)
end
function THDensityMapManager.createInfoLayer(self, name, size, numChannels, basePath, isSaved, isSynced, terrainNode)
    terrainNode = terrainNode or g_terrainNode
    if THUtils.argIsValid(THUtils.isValidName(name), "name", name)
        and THUtils.argIsValid(type(size) == THValueType.NUMBER and size > 0 and math.log(size, 2) % 1 == 0, "size", size)
        and THUtils.argIsValid(type(numChannels) == THValueType.NUMBER and numChannels > 0 and numChannels % 1 == 0, "numChannels", numChannels)
        and THUtils.argIsValid(type(basePath) == THValueType.STRING, "basePath", basePath)
        and THUtils.argIsValid(not isSaved or isSaved == true, "isSaved", isSaved)
        and THUtils.argIsValid(not isSynced or isSynced == true, "isSynced", isSynced)
        and THUtils.argIsValid(THUtils.getIsType(terrainNode, THValueType.INTEGER) and terrainNode > 0, "terrainNode", terrainNode)
    then
        local mission = g_currentMission
        local nameUpper = string.upper(name)
        local layerData = self.infoLayers.byName[nameUpper]
        if layerData == nil then
            basePath = string.gsub(basePath, "/$", "")
            local bitVectorId = createBitVectorMap(name)
            local absFilename = nil
            if mission == nil or mission.missionInfo == nil then
                THUtils.errorMsg(nil, "Trying to load infoLayer %q too early!", name)
            else
                local savegamePath = THUtils.getNormalizedFilePath(mission.missionInfo.savegameDirectory)
                if savegamePath ~= nil and savegamePath ~= "" then
                    absFilename = savegamePath .. name .. ".grle"
                    if not fileExists(absFilename) then
                        absFilename = nil
                    end
                end
            end
            if absFilename == nil then
                absFilename = string.format("%s%s/%s.png", g_thModPath, basePath, name)
                if not fileExists(absFilename) then
                    absFilename = nil
                end
            end
            if absFilename == nil then
                absFilename = string.format("%s%s/%s.grle", g_thModPath, basePath, name)
                if not fileExists(absFilename) then
                    absFilename = nil
                end
            end
            if absFilename ~= nil then
                local isFileValid = false
                if loadBitVectorMapFromFile(bitVectorId, absFilename, numChannels) then
                    local curBitVectorSize = getBitVectorMapSize(bitVectorId) or 0
                    local curBitVectorNumChannels = getBitVectorMapNumChannels(bitVectorId) or 0
                    if curBitVectorSize ~= size then
                        THUtils.errorMsg(false, "Info layer size changed (%s >> %s): %s", curBitVectorSize, size, absFilename)
                    elseif curBitVectorNumChannels ~= numChannels then
                        THUtils.errorMsg(false, "Info layer numChannels changed (%s >> %s): %s", curBitVectorNumChannels, numChannels, absFilename)
                    else
                        isFileValid = true
                    end
                end
                if not isFileValid then
                    delete(bitVectorId)
                    bitVectorId = createBitVectorMap(name)
                    absFilename = nil
                end
            end
            if absFilename == nil then
                loadBitVectorMapNew(bitVectorId, size, size, numChannels, false)
            end
            if type(bitVectorId) ~= THValueType.NUMBER or bitVectorId <= 0 then
                THUtils.errorMsg(nil, "Could not create info layer: %s", name)
            else
                if absFilename ~= nil then
                    THUtils.displayMsg("Loading %d x %d info layer: %s", size, size, absFilename)
                else
                    THUtils.displayMsg("Creating %d x %d info layer: %s", size, size, name)
                end
                layerData = setmetatable({}, THInfoLayer_mt)
                layerData.desc = {
                    owner = self,
                    name = name,
                    size = size,
                    bitVectorId = bitVectorId,
                    numChannels = numChannels,
                    isSaved = isSaved == true,
                    isSynced = isSynced == true,
                    isSavePrepared = false,
                    isResetting = false
                }
                layerData.desc.dmModifier = layerData:createModifier()
                layerData.desc.dmFilter = layerData:createFilter()
                table.insert(self.infoLayers.byIndex, layerData)
                layerData.desc.index = #self.infoLayers.byIndex
                self.infoLayers.byName[nameUpper] = layerData
                if layerData.desc.isSaved then
                    table.insert(self.infoLayers.toSave, layerData)
                end
                if layerData.desc.isSynced then
                    table.insert(self.infoLayers.toSync, layerData)
                end
            end
        end
        return layerData
    end
end
function THDensityMapManager.getInfoLayer(self, infoLayerId, verbose)
    local idVarType = type(infoLayerId)
    local layerData = nil
    if idVarType == THValueType.STRING then
        layerData = self.infoLayers.byName[infoLayerId:upper()]
    elseif idVarType == THValueType.NUMBER then
        layerData = self.infoLayers.byIndex[infoLayerId]
    elseif infoLayerId ~= nil then
        verbose = true
    end
    if layerData == nil and verbose then
        THUtils.errorMsg(true, THMessage.ARGUMENT_INVALID, "infoLayerId", infoLayerId)
    end
    return layerData
end
function THDensityMapManager.prepend_saveCollisionMap(self, superFunc, parent, savegamePath, ...)
    if type(savegamePath) == THValueType.STRING and savegamePath ~= "" then
        local numInfoLayers = #self.infoLayers.toSave
        if numInfoLayers > 0 then
            for layerIndex = 1, numInfoLayers do
                local thInfoLayer = self.infoLayers.toSave[layerIndex]
                if thInfoLayer ~= nil then
                    if thInfoLayer.desc.bitVectorId ~= nil and thInfoLayer.desc.bitVectorId > 0 then
                        local absFilename = THUtils.getNormalizedFilePath(savegamePath)
                        absFilename = absFilename .. thInfoLayer.desc.name .. ".grle"
                        saveBitVectorMapToFile(thInfoLayer.desc.bitVectorId, absFilename)
                    end
                    thInfoLayer.desc.isSavePrepared = false
                end
            end
        end
    end
end
function THDensityMapManager.prepend_prepareSaveCollisionMap(self, superFunc, parent, savegamePath, ...)
    if type(savegamePath) == THValueType.STRING and savegamePath ~= "" then
        local numInfoLayers = #self.infoLayers.toSave
        if numInfoLayers > 0 then
            for layerIndex = 1, numInfoLayers do
                local thInfoLayer = self.infoLayers.toSave[layerIndex]
                if thInfoLayer ~= nil then
                    thInfoLayer.desc.isSavePrepared = false
                    if thInfoLayer.desc.bitVectorId ~= nil and thInfoLayer.desc.bitVectorId > 0 then
                        local absFilename = THUtils.getNormalizedFilePath(savegamePath)
                        absFilename = absFilename .. thInfoLayer.desc.name .. ".grle"
                        prepareSaveBitVectorMapToFile(thInfoLayer.desc.bitVectorId, absFilename)
                        thInfoLayer.desc.isSavePrepared = true
                    end
                end
            end
        end
    end
end
function THDensityMapManager.overwrite_savePreparedCollisionMap(self, superFunc, parent, callbackName, callbackTarget, p1, p2, p3, p4, p5, p6, p7, p8, p9, p10)
    local numInfoLayers = #self.infoLayers.toSave
    if numInfoLayers > 0 then
        local tempTarget = {}
        local layerIndex = 0
        tempTarget.saveNextPreparedBitVectorMap = function()
            layerIndex = layerIndex + 1
            local thInfoLayer = self.infoLayers.toSave[layerIndex]
            if thInfoLayer ~= nil then
                local isBitVectorSaved = false
                if thInfoLayer.desc.isSavePrepared then
                    thInfoLayer.desc.isSavePrepared = false
                    if thInfoLayer.desc.bitVectorId ~= nil and thInfoLayer.desc.bitVectorId > 0 then
                        isBitVectorSaved = true
                        savePreparedBitVectorMapToFile(thInfoLayer.desc.bitVectorId, "saveNextPreparedBitVectorMap", tempTarget)
                    end
                end
                if not isBitVectorSaved then
                    tempTarget.saveNextPreparedBitVectorMap()
                end
            else
                return superFunc(parent, callbackName, callbackTarget, p1, p2, p3, p4, p5, p6, p7, p8, p9, p10)
            end
        end
        tempTarget.saveNextPreparedBitVectorMap()
    else
        return superFunc(parent, callbackName, callbackTarget, p1, p2, p3, p4, p5, p6, p7, p8, p9, p10)
    end
end
function THInfoLayer.createModifier(self, startChannel, numChannels)
    local terrainNode = g_terrainNode
    local dmModifier = nil
    if terrainNode ~= nil
        and self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0
    then
        startChannel = startChannel or 0
        numChannels = numChannels or self.desc.numChannels
        if THUtils.argIsValid(THUtils.getIsType(startChannel, THValueType.INTEGER) and startChannel >= 0, "startChannel", startChannel)
            and THUtils.argIsValid(THUtils.getIsType(numChannels, THValueType.INTEGER) and numChannels > 0, "numChannels", numChannels)
        then
            if startChannel > numChannels then
                THUtils.errorMsg(true, "%s >> Density modifier startChannel (%s) is greater than numChannels (%s)", self.desc.name, startChannel, numChannels)
                startChannel = numChannels
            end
            dmModifier = DensityMapModifier.new(self.desc.bitVectorId, startChannel, numChannels, terrainNode)
        end
    end
    return dmModifier
end
function THInfoLayer.createFilter(self, startChannel, numChannels)
    local dmFilter = nil
    if self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0 then
        startChannel = startChannel or 0
        numChannels = numChannels or self.desc.numChannels
        if THUtils.argIsValid(THUtils.getIsType(startChannel, THValueType.INTEGER) and startChannel >= 0, "startChannel", startChannel)
            and THUtils.argIsValid(THUtils.getIsType(numChannels, THValueType.INTEGER) and numChannels > 0, "numChannels", numChannels)
        then
            if startChannel > numChannels then
                THUtils.errorMsg(true, "%s >> Density filter startChannel (%s) is greater than numChannels (%s)", self.desc.name, startChannel, numChannels)
                startChannel = numChannels
            end
            dmFilter = DensityMapFilter.new(self.desc.bitVectorId, startChannel, numChannels)
        end
    end
    return dmFilter
end
function THInfoLayer.resetArea(self, sx, sz, wx, wz, hx, hz)
    if self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0 then
        local dmModifier = self.desc.dmModifier
        if dmModifier ~= nil then
            dmModifier:setParallelogramWorldCoords(sx, sz, wx, wz, hx, hz, DensityCoordType.POINT_POINT_POINT)
            dmModifier:executeSet(0)
        end
        return true
    end
    return false
end
function THInfoLayer.resetPolygon(self, dmPolygon)
    if self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0 then
        local dmModifier = self.desc.dmModifier
        if dmModifier ~= nil then
            dmPolygon:applyToModifier(dmModifier)
            dmModifier:executeSet(0)
        end
        return true
    end
    return false
end
function THInfoLayer.resetMap(self, dmFilter1, dmFilter2, dmFilter3)
    if self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0 then
        if not self.desc.isResetting then
            local dmModifier = self.desc.dmModifier
            local startPos = -(self.desc.size / 2)
            local resetChunkSize = 4096
            if self.desc.size >= resetChunkSize then
                local numResetSteps = THUtils.round(self.desc.size / resetChunkSize)
                for widthStep = 1, numResetSteps do
                    for lengthStep = 1, numResetSteps do
                        local stepPos = startPos + (resetChunkSize * (widthStep - 1)) + (resetChunkSize * (lengthStep - 1))
                        local sx = stepPos
                        local sz = stepPos
                        local wx = sx + resetChunkSize
                        local wz = sz
                        local hx = sx
                        local hz = sz + resetChunkSize
                        THUtils.addTask(true, function()
                            dmModifier:setParallelogramWorldCoords(sx, sz, wx, wz, hx, hz, DensityCoordType.POINT_POINT_POINT)
                            dmModifier:executeSet(0, dmFilter1, dmFilter2, dmFilter3)
                        end)
                    end
                end
            else
                local sx = startPos
                local sz = startPos
                local wx = sx + self.desc.size
                local wz = sz
                local hx = sx
                local hz = sz + self.desc.size
                THUtils.addTask(true, function()
                    dmModifier:setParallelogramWorldCoords(sx, sz, wx, wz, hx, hz, DensityCoordType.POINT_POINT_POINT)
                    dmModifier:executeSet(0, dmFilter1, dmFilter2, dmFilter3)
                end)
            end
            THUtils.addTask(function()
                self.desc.isResetting = false
                if THUtils.getIsDebugEnabled() then
                    THUtils.displayMsg("InfoLayer %q has been reset", self.desc.name)
                end
            end)
            self.desc.isResetting = true
            return true
        end
    end
    return false
end
function THInfoLayer.delete(self)
    if self.desc.bitVectorId ~= nil and self.desc.bitVectorId > 0 then
        delete(self.desc.bitVectorId)
    end
end
THUtils.pcall(initScript)