
WrapperExtensions = {}

function WrapperExtensions.prerequisitesPresent(specializations)
    return SpecializationUtil.hasSpecialization(BaleWrapper, specializations)
       and SpecializationUtil.hasSpecialization(Consumable, specializations)
end

function WrapperExtensions.initSpecialization()
    local schema = Vehicle.xmlSchema
    schema:setXMLSpecializationType("WrapperExtensions")

    -- Drop animation config
    schema:register(XMLValueType.INT,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?)#wrappingAnimationIndex",
        "wrappingAnimation config index (1-based) that enables these drop animations")
    schema:register(XMLValueType.STRING,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#name",
        "Drop animation name")
    schema:register(XMLValueType.FLOAT,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#animSpeed",
        "Drop animation speed", 1)
    schema:register(XMLValueType.L10N_STRING,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#text",
        "Input help text")
    schema:register(XMLValueType.STRING,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#inputAction",
        "Input action name")
    schema:register(XMLValueType.BOOL,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#liftOnDrop",
        "Lift attacher joint on drop", false)
    schema:register(XMLValueType.FLOAT,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#baleWidth",
        "Bale width to match", 1.2)
    schema:register(XMLValueType.FLOAT,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#baleDiameter",
        "Bale diameter to match", 1.5)
    schema:register(XMLValueType.STRING,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#resetAnimName",
        "Reset animation name to play after this drop animation")
    schema:register(XMLValueType.FLOAT,
        "vehicle.wrapperExtensions.dropAnimationConfig.config(?).dropAnimation(?)#resetAnimSpeed",
        "Reset animation speed", 1)

    -- Slot holder config
    g_vehicleConfigurationManager:addConfigurationType("wrapSlotHolder",
        g_i18n:getText("shop_configuration"), "wrapperExtensions", VehicleConfigurationItem)

    schema:register(XMLValueType.L10N_STRING,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)#name",
        "Configuration name")
    schema:register(XMLValueType.FLOAT,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)#price",
        "Configuration price", 0)
    schema:register(XMLValueType.INT,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)#slotCount",
        "Number of storage slots to enable (0 = only consuming slots)", 0)
    schema:register(XMLValueType.STRING,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)#consumableTypeName",
        "Consumable type name to adjust", "BALE_WRAP")
    schema:register(XMLValueType.NODE_INDEX,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?).slot(?)#node",
        "Slot node to link (by index in consumable type)")
    schema:register(XMLValueType.INT,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?).slot(?)#slotIndex",
        "Storage slot index in consumable type (1-based)")
    ObjectChangeUtil.registerObjectChangeXMLPaths(schema,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)")
    schema:register(XMLValueType.INT,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(?)#connectionHoseConfigId",
        "Connection hose configuration index to use (0 = no extra hoses)", 0)

    schema:register(XMLValueType.L10N_STRING,
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations#title",
        "Configuration title")

    schema:setXMLSpecializationType()
end

function WrapperExtensions.registerOverwrittenFunctions(vehicleType)
    SpecializationUtil.registerOverwrittenFunction(vehicleType, "doStateChange", WrapperExtensions.doStateChange)
    SpecializationUtil.registerOverwrittenFunction(vehicleType, "getConnectionHoseConfigIndex", WrapperExtensions.getConnectionHoseConfigIndex)
end

function WrapperExtensions.registerEventListeners(vehicleType)
    SpecializationUtil.registerEventListener(vehicleType, "onLoad", WrapperExtensions)
    SpecializationUtil.registerEventListener(vehicleType, "onPostLoad", WrapperExtensions)
    SpecializationUtil.registerEventListener(vehicleType, "onRegisterActionEvents", WrapperExtensions)
    SpecializationUtil.registerEventListener(vehicleType, "onFillUnitIsFillingStateChanged", WrapperExtensions)
end

function WrapperExtensions.onFillUnitIsFillingStateChanged(self, isFilling)
    if not isFilling then
        local consumableSpec = self.spec_consumable
        if consumableSpec ~= nil and consumableSpec.types ~= nil then
            for _, consumableType in ipairs(consumableSpec.types) do
                -- Only force-update when there are no storage slots (trimmed config),
                -- because the native storage-slot fill mechanism handles updates correctly
                -- when storage slots exist. Forcing it with storage slots present causes
                -- the inner addFillUnitFillLevel(-numConsumingSlots) to drain fillLevel
                -- after every fill-stop.
                if consumableType.numConsumingSlots > 0 and consumableType.numStorageSlots == 0 then
                    self:updateConsumable(consumableType.typeName, 0, nil, true)
                end
            end
        end
    end
end

function WrapperExtensions.onLoad(self, savegame)
    self.spec_wrapperExtensions = {}
    local spec = self.spec_wrapperExtensions

    -- Read drop animation configs
    spec.dropAnimConfigs = {}
    for _, configKey in self.xmlFile:iterator("vehicle.wrapperExtensions.dropAnimationConfig.config") do
        local configEntry = {
            wrappingAnimationIndex = self.xmlFile:getValue(configKey .. "#wrappingAnimationIndex"),
            dropAnimations = {}
        }
        if configEntry.wrappingAnimationIndex ~= nil then
            for _, dropKey in self.xmlFile:iterator(configKey .. ".dropAnimation") do
                local dropAnim = {
                    name = self.xmlFile:getValue(dropKey .. "#name"),
                    animSpeed = self.xmlFile:getValue(dropKey .. "#animSpeed", 1),
                    text = self.xmlFile:getValue(dropKey .. "#text", nil, self.customEnvironment, false),
                    inputAction = self.xmlFile:getValue(dropKey .. "#inputAction"),
                    liftOnDrop = self.xmlFile:getValue(dropKey .. "#liftOnDrop", false),
                    baleWidth = self.xmlFile:getValue(dropKey .. "#baleWidth", 1.2),
                    baleDiameter = self.xmlFile:getValue(dropKey .. "#baleDiameter", 1.5),
                    resetAnimName = self.xmlFile:getValue(dropKey .. "#resetAnimName"),
                    resetAnimSpeed = self.xmlFile:getValue(dropKey .. "#resetAnimSpeed", 1)
                }
                if dropAnim.inputAction ~= nil then
                    dropAnim.inputAction = InputAction[dropAnim.inputAction]
                end
                table.insert(configEntry.dropAnimations, dropAnim)
            end
            table.insert(spec.dropAnimConfigs, configEntry)
        end
    end

    -- Read slot holder configs
    spec.slotHolderConfigs = {}
    for _, configKey in self.xmlFile:iterator(
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration") do
        local entry = {
            slotCount = self.xmlFile:getValue(configKey .. "#slotCount", 0),
            consumableTypeName = self.xmlFile:getValue(configKey .. "#consumableTypeName", "BALE_WRAP"),
            enabledSlotIndices = {},
            objectChanges = {}
        }
        for _, slotKey in self.xmlFile:iterator(configKey .. ".slot") do
            local idx = self.xmlFile:getValue(slotKey .. "#slotIndex")
            if idx ~= nil then
                table.insert(entry.enabledSlotIndices, idx)
            end
        end
        ObjectChangeUtil.loadObjectChangeFromXML(self.xmlFile, configKey,
            entry.objectChanges, self.components, self)
        table.insert(spec.slotHolderConfigs, entry)
    end

    -- Feature 1: Inject drop animations based on wrappingAnimation config
    local wrappingAnimConfigIndex = Utils.getNoNil(self.configurations.wrappingAnimation, 1)
    local baleWrapperSpec = self.spec_baleWrapper

    if baleWrapperSpec ~= nil then
        for _, configEntry in ipairs(spec.dropAnimConfigs) do
            if configEntry.wrappingAnimationIndex == wrappingAnimConfigIndex then
                local wrapper = baleWrapperSpec.roundBaleWrapper or baleWrapperSpec.squareBaleWrapper
                if wrapper ~= nil and wrapper.allowedBaleTypes ~= nil then
                    for _, baleType in ipairs(wrapper.allowedBaleTypes) do
                        if baleType.dropAnimations == nil then
                            baleType.dropAnimations = {}
                        end
                        for _, dropAnim in ipairs(configEntry.dropAnimations) do
                            local widthMatch = math.abs(baleType.width - dropAnim.baleWidth) < 0.01
                            local diamMatch = math.abs(baleType.diameter - dropAnim.baleDiameter) < 0.01
                            if widthMatch and diamMatch then
                                table.insert(baleType.dropAnimations, {
                                    name = dropAnim.name,
                                    animSpeed = dropAnim.animSpeed,
                                    text = dropAnim.text,
                                    inputAction = dropAnim.inputAction,
                                    liftOnDrop = dropAnim.liftOnDrop,
                                    resetAnimName = dropAnim.resetAnimName,
                                    resetAnimSpeed = dropAnim.resetAnimSpeed
                                })
                            end
                        end
                    end
                end
            end
        end
    end
end

function WrapperExtensions.onPostLoad(self, savegame)
    local spec = self.spec_wrapperExtensions

    -- Feature 2: Adjust consumable slots based on wrapSlotHolder config
    local slotHolderConfigIndex = Utils.getNoNil(self.configurations.wrapSlotHolder, 1)
    local slotConfig = spec.slotHolderConfigs[slotHolderConfigIndex]

    if slotConfig ~= nil then
        local consumableSpec = self.spec_consumable
        if consumableSpec ~= nil then
            local consumableType = consumableSpec.typesByName[slotConfig.consumableTypeName]
            if consumableType ~= nil then
                for i = #consumableType.storageSlots, 1, -1 do
                    local isEnabled = false
                    for _, enabledIdx in ipairs(slotConfig.enabledSlotIndices) do
                        if enabledIdx == i then
                            isEnabled = true
                            break
                        end
                    end
                    if not isEnabled and slotConfig.slotCount < #consumableType.storageSlots then
                        local slot = consumableType.storageSlots[i]
                        if slot.node ~= nil then
                            setVisibility(slot.node, false)
                        end
                    end
                end

                local newStorageCount = slotConfig.slotCount
                local totalSlots = newStorageCount + consumableType.numConsumingSlots

                -- Override FillActivatable.getIsActivatable because its default check
                -- "fillLevel < capacity - 1" prevents refilling when capacity is only 1.
                local fillUnitSpec = self.spec_fillUnit
                if fillUnitSpec ~= nil and fillUnitSpec.fillTrigger ~= nil and fillUnitSpec.fillTrigger.activatable ~= nil then
                    local activatable = fillUnitSpec.fillTrigger.activatable
                    local origGetIsActivatable = activatable.getIsActivatable
                    activatable.getIsActivatable = function(actSelf)
                        if actSelf.vehicle:getIsActiveForInput(true) then
                            local fuIdx = actSelf.vehicle:getFirstValidFillUnitToFill(actSelf.fillTypeIndex)
                            if fuIdx ~= nil then
                                local fillLevel = actSelf.vehicle:getFillUnitFillLevel(fuIdx)
                                local capacity = actSelf.vehicle:getFillUnitCapacity(fuIdx)
                                if fillLevel < capacity - 0.001
                                   and actSelf.vehicle:getFillUnitAllowsFillType(fuIdx, actSelf.fillTypeIndex)
                                   and actSelf.vehicle:getFillUnitSupportsToolType(fuIdx, ToolType.TRIGGER) then
                                    local spec = actSelf.vehicle.spec_fillUnit
                                    for _, trigger in ipairs(spec.fillTrigger.triggers) do
                                        if trigger:getIsActivatable(actSelf.vehicle) then
                                            actSelf:updateActivateText(spec.fillTrigger.isFilling)
                                            return true
                                        end
                                    end
                                end
                            end
                        end
                        return false
                    end
                end

                local needsTrim = newStorageCount < consumableType.numStorageSlots

                -- Clear variationIndex on storage slots that will be trimmed away,
                -- so they don't count toward the underlying fillLevel after Consumable.onPostLoad
                -- pre-filled them based on the original numStorageSlots.
                if needsTrim then
                    local enabledMap = {}
                    for _, enabledIdx in ipairs(slotConfig.enabledSlotIndices) do
                        enabledMap[enabledIdx] = true
                    end
                    for i = 1, #consumableType.storageSlots do
                        local keep = false
                        if #slotConfig.enabledSlotIndices > 0 then
                            keep = enabledMap[i] == true
                        else
                            keep = i <= newStorageCount
                        end
                        if not keep and consumableType.storageSlots[i].consumableVariationIndex ~= 0 then
                            self:setConsumableSlotVariationIndex(consumableType.index, i, 0)
                        end
                    end

                    -- Now trim storage slots and update slots arrays
                    local slotsToKeep = {}
                    if #slotConfig.enabledSlotIndices > 0 then
                        for _, enabledIdx in ipairs(slotConfig.enabledSlotIndices) do
                            if enabledIdx <= #consumableType.storageSlots then
                                table.insert(slotsToKeep, consumableType.storageSlots[enabledIdx])
                            end
                        end
                        consumableType.storageSlots = slotsToKeep
                    else
                        local trimmed = {}
                        for i = 1, math.min(newStorageCount, #consumableType.storageSlots) do
                            table.insert(trimmed, consumableType.storageSlots[i])
                        end
                        consumableType.storageSlots = trimmed
                    end
                    consumableType.numStorageSlots = #consumableType.storageSlots

                    consumableType.slots = {}
                    for _, slot in ipairs(consumableType.storageSlots) do
                        table.insert(consumableType.slots, slot)
                    end
                    for _, slot in ipairs(consumableType.consumingSlots) do
                        table.insert(consumableType.slots, slot)
                    end
                    consumableType.numSlots = #consumableType.slots
                end

                -- Only override capacity if we actually changed slot counts
                if needsTrim then
                    self:setFillUnitCapacity(consumableType.fillUnitIndex, totalSlots, true)
                    self:setFillUnitCapacityToDisplay(consumableType.fillUnitIndex, totalSlots)

                    -- Drain any leftover fillLevel beyond the new capacity
                    local currentFillLevel = self:getFillUnitFillLevel(consumableType.fillUnitIndex) or 0
                    if currentFillLevel > totalSlots then
                        local excess = currentFillLevel - totalSlots
                        self:addFillUnitFillLevel(self:getOwnerFarmId(), consumableType.fillUnitIndex, -excess,
                            self:getFillUnitFillType(consumableType.fillUnitIndex), ToolType.UNDEFINED, nil)
                    end

                    -- Refresh fillLevelToDisplay; force=true is needed so updateConsumable runs
                    -- the inner addFillUnitFillLevel(-numConsumingSlots) path that recalculates
                    -- consumingFillLevel correctly for the trimmed setup.
                    self:updateConsumable(consumableType.typeName, 0, nil, true)
                end
            end
        end

        ObjectChangeUtil.setObjectChanges(slotConfig.objectChanges, true, self, self.setMovingToolDirty)
    end
end

function WrapperExtensions.doStateChange(self, superFunc, changeType, ...)
    local baleWrapperSpec = self.spec_baleWrapper

    if changeType == BaleWrapper.CHANGE_WRAPPER_START_DROP_BALE then
        local dropAnimIndex = baleWrapperSpec.dropAnimationIndex
        local baleTypeIndex = baleWrapperSpec.currentBaleTypeIndex
        local wrapper = baleWrapperSpec.currentWrapper

        if wrapper ~= nil and baleTypeIndex ~= nil then
            local baleType = wrapper.allowedBaleTypes[baleTypeIndex]
            if baleType ~= nil and baleType.dropAnimations ~= nil then
                local dropAnim = baleType.dropAnimations[dropAnimIndex]
                if dropAnim ~= nil and dropAnim.resetAnimName ~= nil then
                    -- Override resetAfterDrop BEFORE superFunc runs,
                    -- so the game uses and tracks the correct animation
                    wrapper.animations.resetAfterDrop.animName = dropAnim.resetAnimName
                    wrapper.animations.resetAfterDrop.animSpeed = dropAnim.resetAnimSpeed
                end
            end
        end
    end

    superFunc(self, changeType, ...)
end

function WrapperExtensions.getConnectionHoseConfigIndex(self, superFunc)
    local baseIndex = superFunc(self)
    local configIndex = Utils.getNoNil(self.configurations.wrapSlotHolder, 1)
    local configKey = string.format(
        "vehicle.wrapperExtensions.wrapSlotHolderConfigurations.wrapSlotHolderConfiguration(%d)#connectionHoseConfigId",
        configIndex - 1)
    local lightOffset = self.xmlFile:getValue(configKey, 0)
    return baseIndex + lightOffset
end

function WrapperExtensions.onRegisterActionEvents(self, isActiveForInput)
    if not isActiveForInput then
        return
    end

    local baleWrapperSpec = self.spec_baleWrapper
    if baleWrapperSpec == nil then
        return
    end

    local wrapper = baleWrapperSpec.currentWrapper
    if wrapper == nil then
        return
    end

    local baleTypeIndex = baleWrapperSpec.currentBaleTypeIndex
    if baleTypeIndex == nil or wrapper.allowedBaleTypes == nil then
        return
    end

    local baleType = wrapper.allowedBaleTypes[baleTypeIndex]
    if baleType == nil or baleType.dropAnimations == nil then
        return
    end

    for i = 1, #baleType.dropAnimations do
        local dropAnim = baleType.dropAnimations[i]
        if dropAnim.inputAction ~= nil then
            local existing = baleWrapperSpec.actionEvents[dropAnim.inputAction]
            if existing == nil then
                local _, actionEventId = self:addPoweredActionEvent(
                    baleWrapperSpec.actionEvents,
                    dropAnim.inputAction,
                    self,
                    BaleWrapper.actionEventDrop,
                    false, true, false, false, i)
                if actionEventId ~= nil then
                    g_inputBinding:setActionEventTextPriority(actionEventId, GS_PRIO_HIGH)
                    g_inputBinding:setActionEventText(actionEventId, dropAnim.text)
                    g_inputBinding:setActionEventActive(actionEventId,
                        baleWrapperSpec.baleWrapperState == BaleWrapper.STATE_WRAPPER_FINSIHED)
                end
            end
        end
    end
end
