DiseaseSprayerHealer = {}
DiseaseSprayerHealer.MOD_NAME = g_currentModName or "FS25_CropDiseases_BMP"

function DiseaseSprayerHealer.prerequisitesPresent(specs)
    return SpecializationUtil.hasSpecialization(Sprayer, specs)
end

function DiseaseSprayerHealer.initSpecialization() end

function DiseaseSprayerHealer.registerEventListeners(vehicleType)
    SpecializationUtil.registerEventListener(vehicleType, "onStartWorkAreaProcessing", DiseaseSprayerHealer)
    SpecializationUtil.registerEventListener(vehicleType, "onUpdateTick", DiseaseSprayerHealer)
end

function DiseaseSprayerHealer.registerOverwrittenFunctions(vehicleType)
    SpecializationUtil.registerOverwrittenFunction(vehicleType, "setSprayerAITerrainDetailProhibitedRange", DiseaseSprayerHealer.setSprayerAITerrainDetailProhibitedRange)
end

local function getCurrentSprayerFillType(self)
    local params = self.spec_sprayer and self.spec_sprayer.workAreaParameters or nil
    local activeFillType = params and params.sprayFillType or nil
    if activeFillType ~= nil and activeFillType ~= FillType.UNKNOWN then
        return activeFillType
    end

    if self.getSprayerFillUnitIndex == nil or self.getFillUnitFillType == nil then
        return FillType.UNKNOWN
    end

    local fillUnitIndex = self:getSprayerFillUnitIndex()
    if fillUnitIndex == nil then
        return FillType.UNKNOWN
    end

    return self:getFillUnitFillType(fillUnitIndex) or FillType.UNKNOWN
end

local function clearThreatOnWorkAreas(self, workAreas, category)
    if DiseaseGrid == nil then
        return 0
    end

    local specWorkArea = self.spec_workArea
    workAreas = workAreas or (specWorkArea and specWorkArea.workAreas or self.workAreas)
    if workAreas == nil then
        return 0
    end

    local cleared = 0
    local protected = 0
    local terrainSize = g_currentMission and g_currentMission.terrainSize or 2048
    local gridWidth = DiseaseGrid.width or 1024
    local cellSize = terrainSize / math.max(gridWidth, 1)
    local margin = math.max(0.75, cellSize * 0.75)

    for _, workArea in pairs(workAreas) do
        local isSprayerArea = workArea ~= nil and (WorkAreaType == nil or WorkAreaType.SPRAYER == nil or workArea.type == WorkAreaType.SPRAYER)
        local isActive = workArea ~= nil and (self.getIsWorkAreaActive == nil or self:getIsWorkAreaActive(workArea))
        if workArea ~= nil and isSprayerArea and isActive and workArea.start ~= nil and workArea.width ~= nil and workArea.height ~= nil then
            if category == "pest" then
                cleared = cleared + (DiseaseGrid:clearParallelogramByCategory(workArea.start, workArea.width, workArea.height, margin, "pest") or 0)
                if DiseaseProtectionGrid ~= nil then
                    protected = protected + (DiseaseProtectionGrid:protectParallelogram(workArea.start, workArea.width, workArea.height, margin, "pest") or 0)
                end
            elseif category == "disease" then
                cleared = cleared + (DiseaseGrid:clearParallelogramByCategory(workArea.start, workArea.width, workArea.height, margin, "disease") or 0)
                if DiseaseProtectionGrid ~= nil then
                    protected = protected + (DiseaseProtectionGrid:protectParallelogram(workArea.start, workArea.width, workArea.height, margin, "disease") or 0)
                end
            end
        end
    end

    return cleared, protected
end

local function isUsingInsecticide(self)
    local insecticideFillType = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("INSECTICIDE") or nil
    if insecticideFillType == nil then
        return false
    end

    return getCurrentSprayerFillType(self) == insecticideFillType
end

local function isUsingFungicide(self)
    local fungicideFillType = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("FUNGICIDE") or nil
    if fungicideFillType == nil then
        return false
    end

    return getCurrentSprayerFillType(self) == fungicideFillType
end

local function configureDiseaseAIRequirements(self, fillType)
    if DiseaseGrid == nil or DiseaseGrid.mapId == nil then return false end
    local insecticide = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("INSECTICIDE") or nil
    local fungicide = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("FUNGICIDE") or nil
    local category = fillType == insecticide and "pest" or (fillType == fungicide and "disease" or nil)
    if category == nil then return false end

    if self.clearAITerrainDetailRequiredRange ~= nil then self:clearAITerrainDetailRequiredRange() end
    if self.clearAITerrainDetailProhibitedRange ~= nil then self:clearAITerrainDetailProhibitedRange() end
    if self.clearAIFruitRequirements ~= nil then self:clearAIFruitRequirements() end
    if self.clearAIFruitProhibitions ~= nil then self:clearAIFruitProhibitions() end

    -- DiseaseGrid is 1024x1024 and cannot be registered in FieldCropsQuery,
    -- which requires terrain-detail resolution. The actual area query is
    -- handled by the AIVehicleUtil hook in DiseaseCore.
    if self.addAIGroundTypeRequirements ~= nil and Sprayer ~= nil then
        self:addAIGroundTypeRequirements(Sprayer.AI_REQUIRED_GROUND_TYPES)
    end

    self.spec_diseaseSprayerHealer = self.spec_diseaseSprayerHealer or {}
    self.spec_diseaseSprayerHealer.aiFillType = fillType
    self.spec_diseaseSprayerHealer.aiTreatmentFarmlandId = nil
    self.spec_diseaseSprayerHealer.aiTreatmentCategory = nil
    return true
end

function DiseaseSprayerHealer:setSprayerAITerrainDetailProhibitedRange(superFunc, fillType)
    if configureDiseaseAIRequirements(self, fillType) then return end
    self.spec_diseaseSprayerHealer = self.spec_diseaseSprayerHealer or {}
    self.spec_diseaseSprayerHealer.aiFillType = fillType
    return superFunc(self, fillType)
end

function DiseaseSprayerHealer:onUpdateTick(_, _, _, _)
    if g_server == nil or self.getSprayerFillUnitIndex == nil or self.getFillUnitFillType == nil then return end
    local fillUnitIndex = self:getSprayerFillUnitIndex()
    local fillType = fillUnitIndex ~= nil and self:getFillUnitFillType(fillUnitIndex) or FillType.UNKNOWN
    local state = self.spec_diseaseSprayerHealer
    if state ~= nil and self.getIsAIActive ~= nil and not self:getIsAIActive() then
        state.aiTreatmentFarmlandId = nil
        state.aiTreatmentCategory = nil
    end
    if state == nil or state.aiFillType ~= fillType then
        self:setSprayerAITerrainDetailProhibitedRange(fillType)
    end
end

function DiseaseSprayerHealer:onStartWorkAreaProcessing(dt, workAreas)
    if g_server == nil or DiseaseGrid == nil then
        return
    end

    if self.getIsTurnedOn == nil or not self:getIsTurnedOn() then
        return
    end

    local fillType = getCurrentSprayerFillType(self)
    local insecticideFillType = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("INSECTICIDE") or nil
    local fungicideFillType = g_fillTypeManager and g_fillTypeManager:getFillTypeIndexByName("FUNGICIDE") or nil
    local category = nil
    if insecticideFillType ~= nil and fillType == insecticideFillType then
        category = "pest"
    elseif fungicideFillType ~= nil and fillType == fungicideFillType then
        category = "disease"
    end
    if category == nil then
        return
    end

    local sprayerParams = self.spec_sprayer and self.spec_sprayer.workAreaParameters or nil
    if sprayerParams == nil or (sprayerParams.sprayFillLevel or 0) <= 0 then
        return
    end

    local cleared, protected = clearThreatOnWorkAreas(self, workAreas, category)
    if (cleared > 0 or protected > 0) and DiseaseCore ~= nil and DiseaseCore.markCareProtectionAt ~= nil then
        self.spec_diseaseSprayerHealer = self.spec_diseaseSprayerHealer or {}
        local state = self.spec_diseaseSprayerHealer
        local now = g_currentMission and g_currentMission.time or 0
        if now - (state.lastCareMarkTime or -1000) >= 750 then
            state.lastCareMarkTime = now
            DiseaseCore:markCareProtectionAt(self, category, nil)
        end
    end
end




