do


local Game = import('/lua/game.lua')
local FactoryFixes = import(Game.CBFP_Path..'/lua/FactoryFixes.lua').FactoryFixes

# The altered factory unit class would be ideal except that it doesn't work. The code in this file gets appended at 
# the end to the existing file from stock FA. Because the air, ground and naval factory classes are generated before 
# this script is even executed the altered factory class won't be used. I can ofcourse re-generate the factory 
# classes but that will affect already loaded mods that change this code aswell. So the best sollution to the problem
# is to apply the bug fix that was originally meant to go in the factory unit class to each dedicated factory class.


#-------------------------------------------------------------
#  FACTORY  UNITS
#-------------------------------------------------------------
FactoryUnit = FactoryFixes(FactoryUnit)

#-------------------------------------------------------------
#  AIR FACTORY UNITS
#-------------------------------------------------------------
AirFactoryUnit = FactoryFixes(AirFactoryUnit)

#-------------------------------------------------------------
#  LAND FACTORY UNITS
#-------------------------------------------------------------
LandFactoryUnit = FactoryFixes(LandFactoryUnit)

#-------------------------------------------------------------
#  SEA FACTORY UNITS
#-------------------------------------------------------------
SeaFactoryUnit = FactoryFixes(SeaFactoryUnit)

#-------------------------------------------------------------
#  MASS COLLECTION UNITS
#-------------------------------------------------------------
CBFP_oldMassCollectionUnit = MassCollectionUnit
MassCollectionUnit = Class(CBFP_oldMassCollectionUnit) {

    WatchUpgradeConsumption = function(self, massConsumption, massProduction)

        # Fix for weird mex behaviour when upgrading with depleted resource stock or while paused [100]
        # Replaced Gowerly's fix with this which is very much inspired by his code. My code looks much better and 
        # seems to work a little better aswell.
        
        local aiBrain = self:GetAIBrain()

        local CalcEnergyFraction = function()
            local fraction = 1
            if aiBrain:GetEconomyStored( 'ENERGY' ) < self:GetConsumptionPerSecondEnergy() then
                fraction = math.min( 1, aiBrain:GetEconomyIncome('ENERGY') / aiBrain:GetEconomyRequested('ENERGY') )
            end
            return fraction
        end

        local CalcMassFraction = function()
            local fraction = 1
            if aiBrain:GetEconomyStored( 'MASS' ) < self:GetConsumptionPerSecondMass() then
                fraction = math.min( 1, aiBrain:GetEconomyIncome('MASS') / aiBrain:GetEconomyRequested('MASS') )
            end
            return fraction
        end

        while not self:IsDead() do

            if self:IsPaused() then
                # paused mex upgrade (another bug here that caused paused upgrades to continue use resources)
                self:SetConsumptionPerSecondMass( 0 )
                self:SetProductionPerSecondMass( massProduction * CalcEnergyFraction() )

            elseif aiBrain:GetEconomyStored( 'MASS' ) < 1 then
                # mex upgrade while out of mass (this is where the engine code has a bug)
                self:SetConsumptionPerSecondMass( massConsumption )
                self:SetProductionPerSecondMass( massProduction / CalcMassFraction() )
                # to use Gowerly's words; the above division cancels the engine bug like matter and anti-matter.
                # the engine seems to do the exact opposite of this division.

            else
                # mex upgrade while enough mass (don't care about energy, that works fine)
                self:SetConsumptionPerSecondMass( massConsumption )
                self:SetProductionPerSecondMass( massProduction * CalcEnergyFraction() )

            end

            WaitTicks(1)
        end
    end,     
}

end