---
-- TrailerDriveAxle
--
-- Trailer drive axle controller: applies torque to selected wheels based on the attacher's PTO.
-- Source: Torque formula  P = T * omega  =>  T = P / omega:
-- https://de.wikipedia.org/wiki/Drehmoment  "Performance" section
--
-- @author      nilBrain
-- @date        2026-03-22
--
-- Copyright (C) nilBrain | Philipp Jaeger, Confidential, All Rights Reserved.

TrailerDriveAxle = {}

TrailerDriveAxle.MOD_NAME = g_currentModName
TrailerDriveAxle.SPEC_NAME = TrailerDriveAxle.MOD_NAME .. ".trailerDriveAxle"
TrailerDriveAxle.SPEC_TABLE_NAME = "spec_" .. TrailerDriveAxle.MOD_NAME .. ".trailerDriveAxle"

TrailerDriveAxle.CONFIG_NAME = "trailerDriveAxle"
TrailerDriveAxle.XML_CONFIG_BASE_KEY = "vehicle.trailerDriveAxleConfigurations.trailerDriveAxleConfiguration"

TrailerDriveAxle.DEFAULT_PTO_RPM = 540
TrailerDriveAxle.DEFAULT_MAX_POWER_KW = 50
TrailerDriveAxle.DEFAULT_TORQUE_SCALE = 0.85
TrailerDriveAxle.DEFAULT_MAX_FORWARD_SPEED = 15
TrailerDriveAxle.DEFAULT_MAX_BACKWARD_SPEED = 5
TrailerDriveAxle.DEFAULT_GRIP_FACTOR = 1
TrailerDriveAxle.DEFAULT_ROLLING_CIRCUMFERENCE_M = 1.8
TrailerDriveAxle.OMEGA_BOOST_FACTOR = 1.10
TrailerDriveAxle.OMEGA_MIN_FACTOR = 0.6
TrailerDriveAxle.OMEGA_MAX_FACTOR = 1.2
TrailerDriveAxle.BRAKE_INPUT_THRESHOLD = 0.01
TrailerDriveAxle.DRIVING_DIRECTION_THRESHOLD = 0.01
TrailerDriveAxle.DIFFERENTIAL_RATIO = 0.5
TrailerDriveAxle.DIFFERENTIAL_BIAS = 1.3
TrailerDriveAxle.KMH_TO_MS = 1 / 3.6


---Called while initializing the specialization
function TrailerDriveAxle.initSpecialization()
    g_vehicleConfigurationManager:addConfigurationType(TrailerDriveAxle.CONFIG_NAME, g_i18n:getText("configuration_trailerDriveAxlePower"), nil, VehicleConfigurationItem)

    local schema = Vehicle.xmlSchema
    schema:setXMLSpecializationType("TrailerDriveAxle")

    TrailerDriveAxle.registerXMLPaths(schema)

    schema:setXMLSpecializationType()
end

---Register XML paths for the specialization
-- @param XMLSchema schema xml schema
function TrailerDriveAxle.registerXMLPaths(schema)
    TrailerDriveAxle.registerTrailerDriveAxleXMLPaths(schema, TrailerDriveAxle.XML_CONFIG_BASE_KEY .. "(?).axle")
    TrailerDriveAxle.registerTrailerDriveAxleXMLPaths(schema, "vehicle.trailerDriveAxle.axle")
end

---Register the axle-specific XML attributes under the given base path
-- @param XMLSchema schema xml schema
-- @param string basePath base xml path
function TrailerDriveAxle.registerTrailerDriveAxleXMLPaths(schema, basePath)
    schema:register(XMLValueType.FLOAT, basePath .. "#maxPowerKW", "Power in kW", TrailerDriveAxle.DEFAULT_MAX_POWER_KW)
    schema:register(XMLValueType.STRING, basePath .. "#driveWheelIndices", "Wheel indices, space-separated, 0-based")
    schema:register(XMLValueType.FLOAT, basePath .. "#brakeForce", "Brake torque when PTO is off in Nm", 0)
    schema:register(XMLValueType.FLOAT, basePath .. "#gripFactor", "Grip factor 0..1", TrailerDriveAxle.DEFAULT_GRIP_FACTOR)
    schema:register(XMLValueType.FLOAT, basePath .. "#torqueScale", "Torque scaling 0-1", TrailerDriveAxle.DEFAULT_TORQUE_SCALE)
    schema:register(XMLValueType.FLOAT, basePath .. "#maxForwardSpeed", "Maximum forward speed in km/h", TrailerDriveAxle.DEFAULT_MAX_FORWARD_SPEED)
    schema:register(XMLValueType.FLOAT, basePath .. "#maxBackwardSpeed", "Maximum reverse speed in km/h", TrailerDriveAxle.DEFAULT_MAX_BACKWARD_SPEED)
end

---Register all functions from the specialization that can be called on vehicle level
-- @param table vehicleType vehicle type
function TrailerDriveAxle.registerFunctions(vehicleType)
    SpecializationUtil.registerFunction(vehicleType, "getTrailerDriveAxleIsActive", TrailerDriveAxle.getTrailerDriveAxleIsActive)
    SpecializationUtil.registerFunction(vehicleType, "getTrailerDriveAxleTorque", TrailerDriveAxle.getTrailerDriveAxleTorque)
    SpecializationUtil.registerFunction(vehicleType, "setTrailerDriveAxleBrakeForce", TrailerDriveAxle.setTrailerDriveAxleBrakeForce)
    SpecializationUtil.registerFunction(vehicleType, "clearDriveDifferentials", TrailerDriveAxle.clearDriveDifferentials)
    SpecializationUtil.registerFunction(vehicleType, "setDriveControl", TrailerDriveAxle.setDriveControl)
end

---
function TrailerDriveAxle.registerOverwrittenFunctions(vehicleType)
    SpecializationUtil.registerOverwrittenFunction(vehicleType, "getCanBeTurnedOn", TrailerDriveAxle.getCanBeTurnedOn)
end

---Register all events that should be called for this specialization
-- @param table vehicleType vehicle type
function TrailerDriveAxle.registerEventListeners(vehicleType)
    SpecializationUtil.registerEventListener(vehicleType, "onLoad", TrailerDriveAxle)
    SpecializationUtil.registerEventListener(vehicleType, "onDelete", TrailerDriveAxle)
    SpecializationUtil.registerEventListener(vehicleType, "onUpdate", TrailerDriveAxle)
    SpecializationUtil.registerEventListener(vehicleType, "onDetach", TrailerDriveAxle)
end

---Checks if all prerequisite specializations are loaded
-- @param table specializations specializations
-- @return boolean hasPrerequisite true if all prerequisite specializations are loaded
function TrailerDriveAxle.prerequisitesPresent(specializations)
    local hasPrerequisite = SpecializationUtil.hasSpecialization(Attachable, specializations)
        and SpecializationUtil.hasSpecialization(Wheels, specializations)
        and SpecializationUtil.hasSpecialization(TurnOnVehicle, specializations)
    return hasPrerequisite
end

---Called on load
-- @param table savegame savegame
function TrailerDriveAxle:onLoad(savegame)
    local function hasWheelDriveNode(wheel)
        if wheel == nil then
            return false
        end

        return wheel.driveNode ~= nil and wheel.driveNode ~= 0
    end

    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    local xmlFile = self.xmlFile

    spec.isValid = false
    spec.currentTorqueNm = 0
    spec.maxPowerKW = TrailerDriveAxle.DEFAULT_MAX_POWER_KW
    spec.torqueScale = TrailerDriveAxle.DEFAULT_TORQUE_SCALE
    spec.maxForwardSpeed = TrailerDriveAxle.DEFAULT_MAX_FORWARD_SPEED
    spec.maxBackwardSpeed = TrailerDriveAxle.DEFAULT_MAX_BACKWARD_SPEED
    spec.brakeForce = 0
    spec.gripFactor = TrailerDriveAxle.DEFAULT_GRIP_FACTOR
    spec.rollingCircumferenceM = TrailerDriveAxle.DEFAULT_ROLLING_CIRCUMFERENCE_M
    spec.driveWheelIndices = nil
    spec.driveWheels = {}
    spec.currentPowerKW = 0
    spec.driveNode = nil
    spec.differentialsAdded = false
    spec.requestedDirection = 1
    spec.torques = {0, 0}
    spec.rotationSpeeds = {0, 0}
    spec.wheelShapes = {}

    local configIdx = self.configurations[TrailerDriveAxle.CONFIG_NAME] or 1
    local itemPath = string.format("%s(%d).axle", TrailerDriveAxle.XML_CONFIG_BASE_KEY, configIdx - 1)
    local legacyItemPath = "vehicle.trailerDriveAxle.axle"
    if not xmlFile:hasProperty(itemPath) and xmlFile:hasProperty(legacyItemPath) then
        itemPath = legacyItemPath
    end

    if xmlFile:hasProperty(itemPath) then
        spec.maxPowerKW = xmlFile:getFloat(itemPath .. "#maxPowerKW", spec.maxPowerKW)
        spec.gripFactor = xmlFile:getFloat(itemPath .. "#gripFactor", spec.gripFactor)
        spec.torqueScale = xmlFile:getFloat(itemPath .. "#torqueScale", spec.torqueScale)
        spec.maxForwardSpeed = xmlFile:getFloat(itemPath .. "#maxForwardSpeed", spec.maxForwardSpeed)
        spec.maxBackwardSpeed = xmlFile:getFloat(itemPath .. "#maxBackwardSpeed", spec.maxBackwardSpeed)
        spec.brakeForce = xmlFile:getFloat(itemPath .. "#brakeForce", spec.brakeForce)

        local wheelIndicesStr = xmlFile:getString(itemPath .. "#driveWheelIndices", nil)
        if wheelIndicesStr ~= nil then
            spec.driveWheelIndices = {}
            for _, token in ipairs(string.split(wheelIndicesStr, " ")) do
                local index = tonumber(token)
                if index ~= nil then
                    table.insert(spec.driveWheelIndices, index)
                end
            end
        end
    end

    local wheels = self.spec_wheels.wheels

    if spec.driveWheelIndices == nil or #spec.driveWheelIndices == 0 then
        SpecializationUtil.removeEventListener(self, "onUpdate", TrailerDriveAxle)
        SpecializationUtil.removeEventListener(self, "onDetach", TrailerDriveAxle)
        return
    end

    local hasConfiguredDrivenWheel = false
    for _, wheelIndex in ipairs(spec.driveWheelIndices) do
        local wheel = wheels[wheelIndex + 1]
        if hasWheelDriveNode(wheel) then
            hasConfiguredDrivenWheel = true
            break
        end
    end

    if not hasConfiguredDrivenWheel then
        local fallbackIndices = {}
        for i, wheel in ipairs(wheels) do
            if hasWheelDriveNode(wheel) then
                table.insert(fallbackIndices, i - 1)
            end
        end

        if #fallbackIndices == 0 then
            SpecializationUtil.removeEventListener(self, "onUpdate", TrailerDriveAxle)
            SpecializationUtil.removeEventListener(self, "onDetach", TrailerDriveAxle)
            Logging.xmlWarning(self.xmlFile, "[TrailerDriveAxle] No driven wheels available, disabling update/detach listeners.")
            return
        end

        spec.driveWheelIndices = fallbackIndices
    end

    spec.driveWheels = {}
    for _, wheelIndex in ipairs(spec.driveWheelIndices) do
        local wheel = wheels[wheelIndex + 1]
        if wheel ~= nil and wheel.physics ~= nil then
            table.insert(spec.driveWheels, wheel)
        end
    end

    if #spec.driveWheels == 0 then
        SpecializationUtil.removeEventListener(self, "onUpdate", TrailerDriveAxle)
        SpecializationUtil.removeEventListener(self, "onDetach", TrailerDriveAxle)
        Logging.xmlWarning(self.xmlFile, "[TrailerDriveAxle] No valid driven wheel physics found, disabling update/detach listeners.")
        return
    end

    if self.components == nil or self.components[1] == nil then
        SpecializationUtil.removeEventListener(self, "onUpdate", TrailerDriveAxle)
        SpecializationUtil.removeEventListener(self, "onDetach", TrailerDriveAxle)
        Logging.xmlWarning(self.xmlFile, "[TrailerDriveAxle] No components found, disabling update/detach listeners.")
        return
    end

    self.components[1].motorized = true
    spec.driveNode = self.components[1].node

    local driveWheelPhysics = spec.driveWheels[1].physics
    if driveWheelPhysics.radius ~= nil then
        spec.rollingCircumferenceM = 2 * math.pi * driveWheelPhysics.radius
    end

    spec.isValid = true
end

---Called on delete
function TrailerDriveAxle:onDelete()
    if not self.isServer then
        return
    end

    self:clearDriveDifferentials()
end

---Called on update
-- @param float dt time since last call in ms
function TrailerDriveAxle:onUpdate(dt)
    if not self.isServer then
        return
    end

    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    local attacherVehicle = self:getAttacherVehicle() or self:getRootVehicle()

    if attacherVehicle == nil then
        spec.currentPowerKW = 0
        return
    end

    local brakeInput = math.abs(self.spec_wheels.brakePedal or 0)
    if attacherVehicle.spec_wheels ~= nil then
        brakeInput = math.max(brakeInput, math.abs(attacherVehicle.spec_wheels.brakePedal or 0))
    end

    local isBraking = brakeInput > TrailerDriveAxle.BRAKE_INPUT_THRESHOLD
    local speedKMH = math.abs(self:getLastSpeed())
    local maxSpeed = self.movingDirection >= 0 and spec.maxForwardSpeed or spec.maxBackwardSpeed

    if isBraking or speedKMH > maxSpeed then
        spec.currentTorqueNm = 0
        spec.currentPowerKW = 0
        self:clearDriveDifferentials()
        self:setDriveControl(0, 0, 0, 0, math.huge, 0, 1, 0)
        self:setTrailerDriveAxleBrakeForce(spec.brakeForce * brakeInput)
        return
    end

    if not attacherVehicle:getIsMotorStarted() then
        spec.currentTorqueNm = 0
        spec.currentPowerKW = 0
        self:clearDriveDifferentials()
        self:setDriveControl(0, 0, 0, 0, math.huge, 0, 1, 0)
        self:setTrailerDriveAxleBrakeForce(spec.brakeForce)
        return
    end

    if not self:getIsTurnedOn() then
        spec.currentTorqueNm = 0
        spec.currentPowerKW = 0
        self:clearDriveDifferentials()
        self:setDriveControl(0, 0, 0, 0, math.huge, 0, 1, 0)
        self:setTrailerDriveAxleBrakeForce(spec.brakeForce)
        return
    end

    self:setTrailerDriveAxleBrakeForce(0)

    local ptoRPM = TrailerDriveAxle.DEFAULT_PTO_RPM
    if self.getPtoRpm ~= nil then
        ptoRPM = self:getPtoRpm() or ptoRPM
    elseif attacherVehicle.getPtoRpm ~= nil then
        ptoRPM = attacherVehicle:getPtoRpm() or ptoRPM
    end
    ptoRPM = math.max(ptoRPM, 1)

    local omegaRadS = (ptoRPM / 60) * 2 * math.pi
    local wheelOmegaRadS = 0
    if spec.rollingCircumferenceM > 0.01 then
        local speedMS = speedKMH * TrailerDriveAxle.KMH_TO_MS
        wheelOmegaRadS = speedMS * ((2 * math.pi) / spec.rollingCircumferenceM)
    end

    local boostedOmegaRadS = math.max(1, wheelOmegaRadS * TrailerDriveAxle.OMEGA_BOOST_FACTOR)
    local powerW = spec.maxPowerKW * 1000 * spec.torqueScale
    local baseTorqueNm = (powerW / omegaRadS) * spec.gripFactor

    local motor = nil
    if attacherVehicle.spec_motorized ~= nil then
        motor = attacherVehicle.spec_motorized.motor
    end

    local acceleratorPedal = 0
    local direction = spec.requestedDirection

    if motor ~= nil then
        local loadPercentage = 0
        if attacherVehicle.getMotorLoadPercentage ~= nil then
            loadPercentage = math.max(0, attacherVehicle:getMotorLoadPercentage())
        end

        local clutchFactor = 1 - math.min(1, math.max(0, motor:getClutchPedal()))
        local drivingDirection = motor:getDrivingDirection()
        local driveRequest = math.abs(drivingDirection)
        acceleratorPedal = math.max(driveRequest, math.min(1, loadPercentage * clutchFactor))

        if drivingDirection < -TrailerDriveAxle.DRIVING_DIRECTION_THRESHOLD then
            direction = -1
            spec.requestedDirection = -1
        elseif drivingDirection > TrailerDriveAxle.DRIVING_DIRECTION_THRESHOLD then
            direction = 1
            spec.requestedDirection = 1
        end
    end

    local torqueNm = baseTorqueNm * acceleratorPedal
    spec.currentTorqueNm = torqueNm
    spec.currentPowerKW = (torqueNm * omegaRadS) / 1000

    if not spec.differentialsAdded then
        local wheelShapes = spec.wheelShapes
        local wheelShapeCount = 0
        for _, wheel in ipairs(spec.driveWheels) do
            local wheelShape = wheel.physics.wheelShape
            if wheelShape ~= nil then
                wheelShapeCount = wheelShapeCount + 1
                wheelShapes[wheelShapeCount] = wheelShape
            end
        end

        if wheelShapeCount >= 2 then
            for i = 2, wheelShapeCount do
                addDifferential(spec.driveNode, wheelShapes[1], true, wheelShapes[i], true, TrailerDriveAxle.DIFFERENTIAL_RATIO, TrailerDriveAxle.DIFFERENTIAL_BIAS)
            end
            spec.differentialsAdded = true
        end
    end

    local clampedOmega = math.max(1, math.abs(boostedOmegaRadS))
    local minRotSpeed = clampedOmega * TrailerDriveAxle.OMEGA_MIN_FACTOR
    local maxRotSpeed = clampedOmega * TrailerDriveAxle.OMEGA_MAX_FACTOR
    local absTorque = math.max(1, math.abs(torqueNm))
    spec.torques[1] = absTorque
    spec.torques[2] = absTorque
    spec.rotationSpeeds[1] = minRotSpeed
    spec.rotationSpeeds[2] = maxRotSpeed

    setMotorProperties(spec.driveNode, minRotSpeed, maxRotSpeed, 1, 0.001, 0.001, 0.001, spec.rotationSpeeds, spec.torques)
    self:setDriveControl(acceleratorPedal, maxSpeed * TrailerDriveAxle.KMH_TO_MS, 2, boostedOmegaRadS * TrailerDriveAxle.OMEGA_MIN_FACTOR, boostedOmegaRadS * TrailerDriveAxle.OMEGA_MAX_FACTOR, boostedOmegaRadS * TrailerDriveAxle.OMEGA_MIN_FACTOR, direction, math.max(1, math.abs(torqueNm)))

    local drivenWheelCount = #spec.driveWheels
    if drivenWheelCount == 0 then
        return
    end

    local signedTorque = torqueNm * (direction >= 0 and 1 or -1)
    local torquePerWheel = signedTorque / drivenWheelCount

    for _, wheel in ipairs(spec.driveWheels) do
        wheel.physics.torque = torquePerWheel
        wheel:updatePhysics(0)
    end
end

---Called on detach
function TrailerDriveAxle:onDetach()
    if not self.isServer then
        return
    end

    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]

    spec.currentTorqueNm = 0
    self:clearDriveDifferentials()
    self:setDriveControl(0, 0, 0, 0, math.huge, 0, 1, 0)
    self:setIsTurnedOn(false)
    self:setTrailerDriveAxleBrakeForce(spec.brakeForce)
end

---Returns whether the trailer drive axle is currently active
-- @return boolean isActive true if the axle is turned on
function TrailerDriveAxle:getTrailerDriveAxleIsActive()
    local isActive = self:getIsTurnedOn()
    return isActive
end

---Returns the current applied torque in Nm
-- @return float torqueNm current torque
function TrailerDriveAxle:getTrailerDriveAxleTorque()
    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    return spec.currentTorqueNm
end

---Sets the brake force on all driven wheels (server-only)
-- @param float|nil brakeForce target brake force; uses spec.brakeForce when nil
function TrailerDriveAxle:setTrailerDriveAxleBrakeForce(brakeForce)
    if not self.isServer then
        return
    end

    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    local targetBrakeForce = brakeForce or spec.brakeForce

    if #spec.driveWheels == 0 then
        return
    end

    for _, wheel in ipairs(spec.driveWheels) do
        wheel.physics.torque = 0
        wheel:updatePhysics(targetBrakeForce)
    end
end

---
function TrailerDriveAxle:clearDriveDifferentials()
    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    if not spec.differentialsAdded or spec.driveNode == nil then
        return
    end

    removeAllDifferentials(spec.driveNode)
    spec.differentialsAdded = false
end

---
function TrailerDriveAxle:setDriveControl(acceleratorPedal, maxSpeed, maxAcceleration, minMotorRotSpeed, maxMotorRotSpeed, maxMotorRotAcceleration, directionSign, maxClutchTorque)
    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    if spec.driveNode == nil then
        return
    end

    local signedRatio = directionSign >= 0 and 1 or -1
    controlVehicle(spec.driveNode, acceleratorPedal, maxSpeed, maxAcceleration, minMotorRotSpeed, maxMotorRotSpeed, maxMotorRotAcceleration, signedRatio, signedRatio, maxClutchTorque, 0)
end

---
function TrailerDriveAxle:getCanBeTurnedOn(superFunc)
    local spec = self[TrailerDriveAxle.SPEC_TABLE_NAME]
    return spec.isValid
end
