HDUpgradeSavegame = {}
HDUpgradeSavegame.FILE_NAME = "hdBuildingUpgrades.xml"
HDUpgradeSavegame.pendingLevelsByUniqueId = {}
HDUpgradeSavegame.settingsLoadedForMission = nil
HDUpgradeSavegame.settingsTotal = 0
HDUpgradeSavegame.appliedEarly = 0
HDUpgradeSavegame.appliedLate = 0

function HDUpgradeSavegame.getPlaceables()
    if g_currentMission == nil or g_currentMission.placeableSystem == nil then
        return {}
    end
    return g_currentMission.placeableSystem.placeables or {}
end

function HDUpgradeSavegame.getFilePath()
    if g_currentMission == nil
        or g_currentMission.missionInfo == nil
        or g_currentMission.missionInfo.savegameDirectory == nil then
        return nil
    end
    return g_currentMission.missionInfo.savegameDirectory .. "/" .. HDUpgradeSavegame.FILE_NAME
end

function HDUpgradeSavegame.save()
    if g_server == nil then
        return
    end

    local filePath = HDUpgradeSavegame.getFilePath()
    if filePath == nil then
        return
    end

    local xmlFile = XMLFile.create("hdBuildingUpgrades", filePath, "hdBuildingUpgrades")
    if xmlFile == nil then
        Logging.error("[HD Building Upgrades] Could not create %s", filePath)
        return
    end

    local index = 0
    for _, placeable in pairs(HDUpgradeSavegame.getPlaceables()) do
        local spec = placeable[HDPlaceableUpgrade.SPEC_TABLE_NAME]
        local hasProductivities = false
        for _, productivity in pairs(spec ~= nil and spec.productivities or {}) do
            if productivity ~= 100 then
                hasProductivities = true
                break
            end
        end
        if spec ~= nil
            and (spec.level > 1 or hasProductivities)
            and placeable.uniqueId ~= nil
            and not HDPlaceableUpgrade.isRelayProductionPoint(placeable) then
            local key = string.format("hdBuildingUpgrades.building(%d)", index)
            xmlFile:setString(key .. "#uniqueId", placeable.uniqueId)
            xmlFile:setInt(key .. "#level", spec.level)
            local productionIndex = 0
            for productionId, productivity in pairs(spec.productivities or {}) do
                if productivity ~= 100 then
                    local productionKey = string.format("%s.production(%d)", key, productionIndex)
                    xmlFile:setString(productionKey .. "#id", productionId)
                    xmlFile:setInt(productionKey .. "#productivity", productivity)
                    productionIndex = productionIndex + 1
                end
            end
            index = index + 1
        end
    end

    xmlFile:save()
    xmlFile:delete()
    print(string.format("[HD Building Upgrades] Saved %d upgrade levels", index))
end

-- Idempotent per g_currentMission: reads hdBuildingUpgrades.xml into pendingLevelsByUniqueId at
-- most once per mission instance, then every later caller (the early per-placeable load wrappers
-- below, and the onStartMission catch-up pass) just consumes whatever is still pending.
--
-- Returns true only once initialization for the CURRENT g_currentMission has actually completed
-- (either a settings file was read, or none exists -- both are a valid, finished state); returns
-- false while g_server/g_currentMission/the savegame path are not yet available, so callers can
-- retry later instead of proceeding against an incomplete or stale cache. settingsLoadedForMission
-- is only stamped with g_currentMission at the very end, right before returning true: stamping it
-- any earlier (e.g. before the path check) would make every later call this mission short-circuit
-- via the identity check above without ever having actually loaded anything.
function HDUpgradeSavegame.ensureSettingsLoaded()
    if g_server == nil or g_currentMission == nil then
        return false
    end
    if HDUpgradeSavegame.settingsLoadedForMission == g_currentMission then
        return true
    end

    -- Mission changed (or first call ever): drop the previous mission's pending settings right
    -- away, before we even know whether the new mission's savegame path is available yet, so
    -- stale entries from another mission can never be consumed during this window.
    HDUpgradeSavegame.pendingLevelsByUniqueId = {}
    HDUpgradeSavegame.settingsTotal = 0
    HDUpgradeSavegame.appliedEarly = 0
    HDUpgradeSavegame.appliedLate = 0

    local filePath = HDUpgradeSavegame.getFilePath()
    if filePath == nil then
        return false
    end

    local settingsByUniqueId = {}
    local loadedCount = 0
    local xmlFile = XMLFile.loadIfExists("hdBuildingUpgrades", filePath, "hdBuildingUpgrades")
    if xmlFile ~= nil then
        local index = 0
        while true do
            local key = string.format("hdBuildingUpgrades.building(%d)", index)
            if not xmlFile:hasProperty(key) then
                break
            end
            local uniqueId = xmlFile:getString(key .. "#uniqueId")
            local level = xmlFile:getInt(key .. "#level", 1)
            if uniqueId ~= nil then
                local settings = {
                    level = math.clamp(level, 1, HDUpgrade.MAX_LEVEL),
                    productivities = {}
                }
                local productionIndex = 0
                while true do
                    local productionKey = string.format("%s.production(%d)", key, productionIndex)
                    if not xmlFile:hasProperty(productionKey) then
                        break
                    end
                    local productionId = xmlFile:getString(productionKey .. "#id")
                    local productivity = xmlFile:getInt(productionKey .. "#productivity", 100)
                    if productionId ~= nil and HDPlaceableUpgrade.isValidProductivity(productivity) then
                        settings.productivities[productionId] = productivity
                    end
                    productionIndex = productionIndex + 1
                end
                settingsByUniqueId[uniqueId] = settings
                loadedCount = loadedCount + 1
            end
            index = index + 1
        end
        xmlFile:delete()
    end
    -- xmlFile == nil (no settings file yet) is a valid, completed state: an empty settings table.

    HDUpgradeSavegame.pendingLevelsByUniqueId = settingsByUniqueId
    HDUpgradeSavegame.settingsTotal = loadedCount
    HDUpgradeSavegame.settingsLoadedForMission = g_currentMission
    return true
end

-- Runs at FSBaseMission.onStartMission, after every placeable has already loaded (and therefore
-- after the early ProductionPoint/PlaceableSilo wrappers below have already consumed whatever
-- pending entries they could reach). This pass is the safety net for anything those two wrappers
-- don't cover (e.g. productivity-only entries on buildings without a silo/production storage);
-- applyPendingLevel() only acts on entries that are still pending, so nothing is reloaded or
-- applied twice.
function HDUpgradeSavegame.load()
    if g_server == nil then
        return
    end

    local settingsReady = HDUpgradeSavegame.ensureSettingsLoaded()

    local appliedLate = 0
    if settingsReady then
        for _, placeable in pairs(HDUpgradeSavegame.getPlaceables()) do
            if not HDPlaceableUpgrade.isRelayProductionPoint(placeable) then
                HDPlaceableUpgrade.ensureRuntimeSpec(placeable)
                if HDUpgradeSavegame.applyPendingLevel(placeable) then
                    appliedLate = appliedLate + 1
                end
            end
        end
    end
    HDUpgradeSavegame.appliedLate = appliedLate

    print(string.format(
        "[HD Building Upgrades] Startup: %d setting(s) loaded, %d applied before storage loading, %d applied at mission start",
        HDUpgradeSavegame.settingsTotal,
        HDUpgradeSavegame.appliedEarly,
        HDUpgradeSavegame.appliedLate
    ))
end

function HDUpgradeSavegame.applyPendingLevel(placeable)
    if placeable == nil or placeable.uniqueId == nil then
        return false
    end

    if HDPlaceableUpgrade.isRelayProductionPoint(placeable) then
        HDUpgradeSavegame.pendingLevelsByUniqueId[placeable.uniqueId] = nil
        return false
    end

    local settings = HDUpgradeSavegame.pendingLevelsByUniqueId[placeable.uniqueId]
    if settings == nil then
        return false
    end
    local spec = placeable[HDPlaceableUpgrade.SPEC_TABLE_NAME]
    if spec == nil or not HDPlaceableUpgrade.hasUpgradeableBase(spec.base) then
        return false
    end

    HDUpgradeSavegame.pendingLevelsByUniqueId[placeable.uniqueId] = nil

    spec.productivities = settings.productivities or {}
    HDPlaceableUpgrade.applyLevel(placeable, settings.level)
    return true
end

-- Overwrites ProductionPoint:loadFromXMLFile so the upgraded capacity is applied before the
-- original method runs. The original calls self.storage:loadFromXMLFile(), whose setFillLevel()
-- clamps the saved fill level to whatever capacity is set on the storage at that moment; without
-- this, capacity is still at its un-upgraded XML default and a fill level above it gets clamped
-- down and lost. Client calls skip straight to superFunc without touching the savegame file.
function HDUpgradeSavegame.overwriteProductionPointLoadFromXMLFile(productionPoint, superFunc, xmlFile, key)
    if g_server ~= nil then
        local placeable = productionPoint ~= nil and productionPoint.owningPlaceable or nil
        local settingsReady = HDUpgradeSavegame.ensureSettingsLoaded()
        HDPlaceableUpgrade.ensureRuntimeSpec(placeable, productionPoint)
        if settingsReady and HDUpgradeSavegame.applyPendingLevel(placeable) then
            HDUpgradeSavegame.appliedEarly = HDUpgradeSavegame.appliedEarly + 1
        end
    end
    return superFunc(productionPoint, xmlFile, key)
end

ProductionPoint.loadFromXMLFile = Utils.overwrittenFunction(
    ProductionPoint.loadFromXMLFile,
    HDUpgradeSavegame.overwriteProductionPointLoadFromXMLFile
)

-- Same fix as above, for silo-only buildings: PlaceableSilo:loadFromXMLFile iterates all of a
-- placeable's storages in one call and would otherwise clamp each of them the same way.
function HDUpgradeSavegame.overwritePlaceableSiloLoadFromXMLFile(placeable, superFunc, xmlFile, key)
    if g_server ~= nil then
        local settingsReady = HDUpgradeSavegame.ensureSettingsLoaded()
        HDPlaceableUpgrade.ensureRuntimeSpec(placeable)
        if settingsReady and HDUpgradeSavegame.applyPendingLevel(placeable) then
            HDUpgradeSavegame.appliedEarly = HDUpgradeSavegame.appliedEarly + 1
        end
    end
    return superFunc(placeable, xmlFile, key)
end

PlaceableSilo.loadFromXMLFile = Utils.overwrittenFunction(
    PlaceableSilo.loadFromXMLFile,
    HDUpgradeSavegame.overwritePlaceableSiloLoadFromXMLFile
)

FSBaseMission.onStartMission = Utils.appendedFunction(
    FSBaseMission.onStartMission,
    HDUpgradeSavegame.load
)

FSCareerMissionInfo.saveToXMLFile = Utils.appendedFunction(
    FSCareerMissionInfo.saveToXMLFile,
    HDUpgradeSavegame.save
)
