-- Per-wheel load sampling for Mud System Physics.
-- Uses the engine contact load first and falls back to vehicle/payload mass.

WheelLoadSystem = {}

WheelLoadSystem.enabled = true
WheelLoadSystem.debug = false

WheelLoadSystem.cacheMs = 150
WheelLoadSystem.fallbackCacheMs = 250
WheelLoadSystem.referenceWheelLoadKg = 2500
WheelLoadSystem.minimumLoadFactor = 0.60
WheelLoadSystem.maximumLoadFactor = 1.65
WheelLoadSystem.loadExponent = 0.55
WheelLoadSystem.minimumMeasuredLoadKg = 25
WheelLoadSystem.maximumMeasuredLoadKg = 60000
WheelLoadSystem.measuredLoadSmoothing = 0.45

WheelLoadSystem.raisedWheeledTransferFactor = 0.35
WheelLoadSystem.attachmentAxleBias = 0.85
WheelLoadSystem.attachmentTreeMaxDepth = 8

WheelLoadSystem._wheelCache = setmetatable({}, { __mode = "k" })
WheelLoadSystem._vehicleCache = setmetatable({}, { __mode = "k" })
WheelLoadSystem._layoutCache = setmetatable({}, { __mode = "k" })
WheelLoadSystem._lastMeasuredLoad = setmetatable({}, { __mode = "k" })

local function wlsClamp(value, minimum, maximum)
    value = tonumber(value) or minimum
    return math.max(minimum, math.min(maximum, value))
end

local function wlsNowMs()
    if g_time ~= nil then
        return g_time
    end

    if g_currentMission ~= nil and g_currentMission.time ~= nil then
        return g_currentMission.time
    end

    if getTimeSec ~= nil then
        return getTimeSec() * 1000
    end

    return 0
end

local function wlsResolveWheelPhysics(vehicle, wheelOrPhysics)
    if wheelOrPhysics == nil then
        return vehicle, nil, nil
    end

    if wheelOrPhysics.physics ~= nil then
        return vehicle or wheelOrPhysics.physics.vehicle,
               wheelOrPhysics,
               wheelOrPhysics.physics
    end

    local wheel = wheelOrPhysics.wheel
    return vehicle or wheelOrPhysics.vehicle,
           wheel,
           wheelOrPhysics
end

local function wlsHasWheelContact(wheelPhysics)
    if wheelPhysics == nil then
        return false
    end

    if wheelPhysics.hasGroundContact == true
        or wheelPhysics.hasSoilContact == true then
        return true
    end

    if wheelPhysics.contact ~= nil and WheelContactType ~= nil then
        return wheelPhysics.contact ~= WheelContactType.NONE
    end

    return false
end

local function wlsGetWheelNode(wheel)
    if wheel == nil then
        return nil
    end

    return wheel.node or wheel.driveNode or wheel.repr
end

function WheelLoadSystem:getPhysicalWheels(vehicle)
    local result = {}
    local wheelsSpec = vehicle ~= nil and vehicle.spec_wheels or nil

    if wheelsSpec == nil or wheelsSpec.wheels == nil then
        return result
    end

    for _, wheel in pairs(wheelsSpec.wheels) do
        if wheel ~= nil and wheel.physics ~= nil then
            result[#result + 1] = wheel
        end
    end

    return result
end

function WheelLoadSystem:getFillPayloadMassKg(vehicle)
    local fillSpec = vehicle ~= nil and vehicle.spec_fillUnit or nil
    if fillSpec == nil or fillSpec.fillUnits == nil then
        return 0
    end

    local payloadMassKg = 0

    for fillUnitIndex, fillUnit in ipairs(fillSpec.fillUnits) do
        if fillUnit ~= nil then
            local fillLevel = tonumber(fillUnit.fillLevel) or 0
            local fillType = fillUnit.fillType

            if vehicle.getFillUnitFillLevel ~= nil then
                local ok, value = pcall(
                    vehicle.getFillUnitFillLevel,
                    vehicle,
                    fillUnitIndex
                )
                if ok and value ~= nil then
                    fillLevel = tonumber(value) or fillLevel
                end
            end

            if vehicle.getFillUnitFillType ~= nil then
                local ok, value = pcall(
                    vehicle.getFillUnitFillType,
                    vehicle,
                    fillUnitIndex
                )
                if ok and value ~= nil then
                    fillType = value
                end
            end

            if fillLevel > 0
                and fillType ~= nil
                and (FillType == nil or fillType ~= FillType.UNKNOWN)
                and g_fillTypeManager ~= nil then
                local desc = g_fillTypeManager:getFillTypeByIndex(fillType)
                if desc ~= nil then
                    local massPerLiterKg = math.max(
                        0,
                        tonumber(desc.massPerLiter) or 0
                    ) * 1000
                    payloadMassKg = payloadMassKg
                        + fillLevel * massPerLiterKg
                end
            end
        end
    end

    return math.max(0, payloadMassKg)
end

function WheelLoadSystem:getOwnMassData(vehicle)
    local result = {
        totalMassKg = 0,
        payloadMassKg = 0,
        source = "none"
    }

    if vehicle == nil then
        return result
    end

    local measuredMassKg = 0

    -- onlyGivenVehicle=true avoids counting the complete attachment tree.
    if vehicle.getTotalMass ~= nil then
        local ok, value = pcall(vehicle.getTotalMass, vehicle, true)
        if ok and value ~= nil then
            measuredMassKg = math.max(0, tonumber(value) or 0) * 1000
            if measuredMassKg > 0 then
                result.source = "getTotalMass"
            end
        end
    end

    if measuredMassKg <= 0 and vehicle.getTotalMassInKg ~= nil then
        local ok, value = pcall(vehicle.getTotalMassInKg, vehicle)
        if ok and value ~= nil then
            measuredMassKg = math.max(0, tonumber(value) or 0)
            if measuredMassKg > 0 then
                result.source = "getTotalMassInKg"
            end
        end
    end

    local defaultMassKg = 0
    if vehicle.getDefaultMass ~= nil then
        local ok, value = pcall(vehicle.getDefaultMass, vehicle)
        if ok and value ~= nil then
            defaultMassKg = math.max(0, tonumber(value) or 0) * 1000
        end
    end

    if defaultMassKg <= 0 and vehicle.defaultMass ~= nil then
        defaultMassKg = math.max(0, tonumber(vehicle.defaultMass) or 0) * 1000
    end

    local payloadMassKg = self:getFillPayloadMassKg(vehicle)
    local estimatedMassKg = defaultMassKg > 0
        and (defaultMassKg + payloadMassKg)
        or payloadMassKg

    result.totalMassKg = math.max(measuredMassKg, estimatedMassKg)
    result.payloadMassKg = math.min(result.totalMassKg, payloadMassKg)

    if estimatedMassKg > measuredMassKg + 1 then
        result.source = "defaultMass+fillUnits"
    end

    return result
end

function WheelLoadSystem:getImplementSupportState(object)
    local physicalWheels = self:getPhysicalWheels(object)
    local wheelContact = false

    for _, wheel in ipairs(physicalWheels) do
        if wlsHasWheelContact(wheel.physics) then
            wheelContact = true
            break
        end
    end

    local isLowered = false
    if object ~= nil and object.getIsLowered ~= nil then
        local ok, value = pcall(object.getIsLowered, object, false)
        if not ok then
            ok, value = pcall(object.getIsLowered, object)
        end
        isLowered = ok and value == true
    end

    local groundReferenceActive = false
    local groundSpec = object ~= nil and object.spec_groundReference or nil
    if groundSpec ~= nil and groundSpec.groundReferenceNodes ~= nil then
        for _, nodeData in ipairs(groundSpec.groundReferenceNodes) do
            if nodeData ~= nil and nodeData.isActive == true then
                groundReferenceActive = true
                break
            end
        end
    end

    local workAreaActive = false
    local workSpec = object ~= nil and object.spec_workArea or nil
    if workSpec ~= nil and workSpec.workAreas ~= nil then
        for _, workArea in pairs(workSpec.workAreas) do
            local active = false

            if object.getIsWorkAreaProcessing ~= nil then
                local ok, value = pcall(
                    object.getIsWorkAreaProcessing,
                    object,
                    workArea
                )
                active = ok and value == true
            end

            if not active and object.getIsWorkAreaActive ~= nil then
                local ok, value = pcall(
                    object.getIsWorkAreaActive,
                    object,
                    workArea
                )
                active = ok and value == true
            end

            if active then
                workAreaActive = true
                break
            end
        end
    end

    local isGroundWorkingTool = object ~= nil and (
        object.spec_workArea ~= nil
        or object.spec_groundReference ~= nil
        or object.spec_plow ~= nil
        or object.spec_cultivator ~= nil
        or object.spec_sowingMachine ~= nil
        or object.spec_roller ~= nil
        or object.spec_mower ~= nil
        or object.spec_stonePicker ~= nil
        or object.spec_powerHarrow ~= nil
        or object.spec_subsoiler ~= nil
    )

    return {
        wheelContact = wheelContact,
        physicalWheelCount = #physicalWheels,
        groundSupported = wheelContact
            or groundReferenceActive
            or workAreaActive
            or (isLowered and isGroundWorkingTool)
    }
end

function WheelLoadSystem:getAttachmentSide(carrier, object, implement)
    if carrier == nil or object == nil or localToLocal == nil then
        return "center"
    end

    local carrierNode = (
        carrier.components ~= nil
        and carrier.components[1] ~= nil
        and carrier.components[1].node
    ) or carrier.rootNode

    local sourceNode = nil
    if carrier.spec_attacherJoints ~= nil
        and carrier.spec_attacherJoints.attacherJoints ~= nil
        and implement ~= nil
        and implement.jointDescIndex ~= nil then
        local joint = carrier.spec_attacherJoints.attacherJoints[
            implement.jointDescIndex
        ]
        if joint ~= nil then
            sourceNode = joint.jointTransform or joint.node or joint.rootNode
        end
    end

    sourceNode = sourceNode or (
        object.components ~= nil
        and object.components[1] ~= nil
        and object.components[1].node
    ) or object.rootNode

    if carrierNode == nil or sourceNode == nil then
        return "center"
    end

    local ok, _, _, z = pcall(
        localToLocal,
        sourceNode,
        carrierNode,
        0,
        0,
        0
    )

    if not ok or z == nil then
        return "center"
    end

    if z > 0.50 then
        return "front"
    elseif z < -0.50 then
        return "rear"
    end

    return "center"
end

function WheelLoadSystem:getTransferredAttachmentData(carrier, visited, depth)
    local result = {
        totalMassKg = 0,
        frontMassKg = 0,
        rearMassKg = 0,
        centerMassKg = 0
    }

    if carrier == nil or carrier.getAttachedImplements == nil then
        return result
    end

    depth = tonumber(depth) or 0
    if depth >= math.max(1, math.floor(self.attachmentTreeMaxDepth or 8)) then
        return result
    end

    local ok, attached = pcall(carrier.getAttachedImplements, carrier)
    if not ok or attached == nil then
        return result
    end

    for _, implement in pairs(attached) do
        local object = implement ~= nil and implement.object or nil
        if object ~= nil and visited[object] ~= true then
            local child = self:getFallbackVehicleMassData(
                object,
                visited,
                depth + 1
            )
            local support = self:getImplementSupportState(object)
            local transferFactor = 1

            if support.groundSupported then
                transferFactor = 0
            elseif support.physicalWheelCount > 0 then
                transferFactor = wlsClamp(
                    self.raisedWheeledTransferFactor,
                    0,
                    1
                )
            end

            local transferredMassKg = math.max(
                0,
                tonumber(child.totalMassKg) or 0
            ) * transferFactor

            if transferredMassKg > 0.5 then
                local side = self:getAttachmentSide(carrier, object, implement)
                result.totalMassKg = result.totalMassKg + transferredMassKg
                if side == "front" then
                    result.frontMassKg = result.frontMassKg + transferredMassKg
                elseif side == "rear" then
                    result.rearMassKg = result.rearMassKg + transferredMassKg
                else
                    result.centerMassKg = result.centerMassKg + transferredMassKg
                end
            end
        end
    end

    return result
end

function WheelLoadSystem:getFallbackVehicleMassData(vehicle, visited, depth)
    if vehicle == nil then
        return {
            totalMassKg = 0,
            ownMassKg = 0,
            frontMassKg = 0,
            rearMassKg = 0,
            centerMassKg = 0,
            source = "none"
        }
    end

    local now = wlsNowMs()
    local cached = self._vehicleCache[vehicle]
    if cached ~= nil
        and now - (cached.timestamp or 0) < math.max(50, self.fallbackCacheMs or 250) then
        return cached
    end

    visited = visited or {}
    if visited[vehicle] == true then
        local own = self:getOwnMassData(vehicle)
        return {
            totalMassKg = own.totalMassKg,
            ownMassKg = own.totalMassKg,
            frontMassKg = 0,
            rearMassKg = 0,
            centerMassKg = 0,
            source = own.source,
            timestamp = now
        }
    end

    visited[vehicle] = true
    local own = self:getOwnMassData(vehicle)
    local attached = self:getTransferredAttachmentData(
        vehicle,
        visited,
        depth or 0
    )
    visited[vehicle] = nil

    local result = {
        totalMassKg = own.totalMassKg + attached.totalMassKg,
        ownMassKg = own.totalMassKg,
        payloadMassKg = own.payloadMassKg,
        frontMassKg = attached.frontMassKg,
        rearMassKg = attached.rearMassKg,
        centerMassKg = attached.centerMassKg,
        source = attached.totalMassKg > 0.5
            and (tostring(own.source) .. "+attachments")
            or own.source,
        timestamp = now
    }

    self._vehicleCache[vehicle] = result
    return result
end

function WheelLoadSystem:getWheelLayout(vehicle)
    if vehicle == nil then
        return nil
    end

    local cached = self._layoutCache[vehicle]
    if cached ~= nil then
        return cached
    end

    local wheels = self:getPhysicalWheels(vehicle)
    local baseNode = (
        vehicle.components ~= nil
        and vehicle.components[1] ~= nil
        and vehicle.components[1].node
    ) or vehicle.rootNode

    if #wheels == 0 or baseNode == nil or localToLocal == nil then
        return nil
    end

    local entries = {}
    local minimumZ = math.huge
    local maximumZ = -math.huge

    for _, wheel in ipairs(wheels) do
        local wheelNode = wlsGetWheelNode(wheel)
        if wheelNode ~= nil then
            local ok, _, _, z = pcall(
                localToLocal,
                wheelNode,
                baseNode,
                0,
                0,
                0
            )
            if ok and z ~= nil then
                z = tonumber(z) or 0
                entries[#entries + 1] = { wheel = wheel, z = z }
                minimumZ = math.min(minimumZ, z)
                maximumZ = math.max(maximumZ, z)
            end
        end
    end

    if #entries == 0 then
        return nil
    end

    local positions = {}
    local frontCount = 0
    local rearCount = 0
    local splitZ = (minimumZ + maximumZ) * 0.5

    if maximumZ - minimumZ < 0.50 then
        for _, entry in ipairs(entries) do
            positions[entry.wheel] = "center"
        end
    else
        for _, entry in ipairs(entries) do
            if entry.z >= splitZ then
                positions[entry.wheel] = "front"
                frontCount = frontCount + 1
            else
                positions[entry.wheel] = "rear"
                rearCount = rearCount + 1
            end
        end
    end

    local result = {
        positions = positions,
        totalCount = #entries,
        frontCount = frontCount,
        rearCount = rearCount
    }
    self._layoutCache[vehicle] = result
    return result
end

function WheelLoadSystem:getFallbackWheelLoadKg(vehicle, wheel)
    local mass = self:getFallbackVehicleMassData(vehicle)
    local layout = self:getWheelLayout(vehicle)
    local totalCount = math.max(
        1,
        layout ~= nil and layout.totalCount or #self:getPhysicalWheels(vehicle)
    )

    if layout == nil
        or layout.frontCount <= 0
        or layout.rearCount <= 0
        or wheel == nil then
        return mass.totalMassKg / totalCount, mass.source
    end

    local axle = layout.positions[wheel]
    if axle ~= "front" and axle ~= "rear" then
        return mass.totalMassKg / totalCount, mass.source
    end

    local bias = wlsClamp(self.attachmentAxleBias, 0.50, 1.00)
    local frontAxleMassKg = mass.ownMassKg * 0.50
        + mass.frontMassKg * bias
        + mass.rearMassKg * (1 - bias)
        + mass.centerMassKg * 0.50
    local rearAxleMassKg = mass.ownMassKg * 0.50
        + mass.rearMassKg * bias
        + mass.frontMassKg * (1 - bias)
        + mass.centerMassKg * 0.50

    if axle == "front" then
        return frontAxleMassKg / math.max(1, layout.frontCount),
               tostring(mass.source) .. ":front"
    end

    return rearAxleMassKg / math.max(1, layout.rearCount),
           tostring(mass.source) .. ":rear"
end

function WheelLoadSystem:getMeasuredWheelLoadKg(wheelPhysics)
    if wheelPhysics == nil or wheelPhysics.getTireLoad == nil then
        return nil
    end

    local ok, value = pcall(wheelPhysics.getTireLoad, wheelPhysics)
    if not ok or value == nil then
        return nil
    end

    local loadKg = wlsClamp(
        value,
        0,
        self.maximumMeasuredLoadKg or 60000
    )

    if loadKg < math.max(1, self.minimumMeasuredLoadKg or 25) then
        return nil
    end

    local previous = self._lastMeasuredLoad[wheelPhysics]
    if previous ~= nil and previous > 0 then
        -- Limit one-sample suspension/contact spikes, then smooth the load.
        loadKg = wlsClamp(loadKg, previous * 0.45, previous * 1.80)
        local alpha = wlsClamp(self.measuredLoadSmoothing, 0.05, 1.00)
        loadKg = previous + (loadKg - previous) * alpha
    end

    self._lastMeasuredLoad[wheelPhysics] = loadKg
    return loadKg
end

function WheelLoadSystem:getLoadFactorFromKg(loadKg)
    if self.enabled ~= true then
        return 1
    end

    local reference = math.max(
        250,
        tonumber(self.referenceWheelLoadKg) or 2500
    )
    local exponent = wlsClamp(self.loadExponent, 0.20, 1.00)
    local factor = math.pow(
        math.max(0.02, (tonumber(loadKg) or 0) / reference),
        exponent
    )

    return wlsClamp(
        factor,
        self.minimumLoadFactor or 0.60,
        self.maximumLoadFactor or 1.65
    )
end

function WheelLoadSystem:getWheelLoadData(vehicle, wheelOrPhysics)
    local resolvedVehicle, wheel, wheelPhysics = wlsResolveWheelPhysics(
        vehicle,
        wheelOrPhysics
    )
    vehicle = resolvedVehicle

    if self.enabled ~= true or wheelPhysics == nil then
        return {
            loadKg = tonumber(self.referenceWheelLoadKg) or 2500,
            loadFactor = 1,
            source = "disabled",
            measured = false
        }
    end

    local now = wlsNowMs()
    local cached = self._wheelCache[wheelPhysics]
    if cached ~= nil
        and now - (cached.timestamp or 0) < math.max(50, self.cacheMs or 150) then
        return cached
    end

    -- WheelPhysics may still expose a small wheel-mass force while airborne.
    -- Treat the engine value as authoritative only for a real contact wheel;
    -- otherwise use the deterministic mass/axle fallback.
    local hasContact = wlsHasWheelContact(wheelPhysics)
    local loadKg = hasContact
        and self:getMeasuredWheelLoadKg(wheelPhysics)
        or nil
    local source = "getTireLoad"
    local measured = loadKg ~= nil

    if loadKg == nil then
        loadKg, source = self:getFallbackWheelLoadKg(vehicle, wheel)
        source = "fallback:" .. tostring(source)
    end

    loadKg = math.max(0, tonumber(loadKg) or 0)
    if loadKg <= 0 then
        loadKg = tonumber(self.referenceWheelLoadKg) or 2500
        source = "reference"
    end

    local result = {
        loadKg = loadKg,
        loadFactor = self:getLoadFactorFromKg(loadKg),
        source = source,
        measured = measured,
        hasContact = hasContact,
        timestamp = now
    }

    self._wheelCache[wheelPhysics] = result

    if self.debug == true then
        wheelPhysics.__wlsNextDebugMs = wheelPhysics.__wlsNextDebugMs or 0
        if now >= wheelPhysics.__wlsNextDebugMs then
            wheelPhysics.__wlsNextDebugMs = now + 5000
            print(string.format(
                "[WheelLoadSystem] load=%.0fkg factor=%.2f source=%s contact=%s",
                result.loadKg,
                result.loadFactor,
                tostring(result.source),
                tostring(result.hasContact)
            ))
        end
    end

    return result
end

function WheelLoadSystem:getWheelLoadKg(vehicle, wheelOrPhysics)
    return self:getWheelLoadData(vehicle, wheelOrPhysics).loadKg
end

function WheelLoadSystem:getWheelLoadFactor(vehicle, wheelOrPhysics)
    return self:getWheelLoadData(vehicle, wheelOrPhysics).loadFactor
end

function WheelLoadSystem:getVehicleLoadFactor(vehicle)
    if self.enabled ~= true
        or vehicle == nil
        or vehicle.spec_wheels == nil
        or vehicle.spec_wheels.wheels == nil then
        return 1
    end

    local total = 0
    local count = 0
    local fallbackTotal = 0
    local fallbackCount = 0

    for _, wheel in pairs(vehicle.spec_wheels.wheels) do
        if wheel ~= nil and wheel.physics ~= nil then
            local data = self:getWheelLoadData(vehicle, wheel.physics)
            fallbackTotal = fallbackTotal + data.loadFactor
            fallbackCount = fallbackCount + 1

            if data.hasContact then
                total = total + data.loadFactor
                count = count + 1
            end
        end
    end

    if count > 0 then
        return total / count
    end

    if fallbackCount > 0 then
        return fallbackTotal / fallbackCount
    end

    return 1
end

function WheelLoadSystem:clearVehicle(vehicle)
    if vehicle == nil then
        return
    end

    self._vehicleCache[vehicle] = nil
    self._layoutCache[vehicle] = nil

    for _, wheel in ipairs(self:getPhysicalWheels(vehicle)) do
        if wheel.physics ~= nil then
            self._wheelCache[wheel.physics] = nil
            self._lastMeasuredLoad[wheel.physics] = nil
        end
    end
end

function WheelLoadSystem:deleteMap()
    self._wheelCache = setmetatable({}, { __mode = "k" })
    self._vehicleCache = setmetatable({}, { __mode = "k" })
    self._layoutCache = setmetatable({}, { __mode = "k" })
    self._lastMeasuredLoad = setmetatable({}, { __mode = "k" })
end

addModEventListener(WheelLoadSystem)
