ArachnaeSwarm/Source/ArachnaeSwarm/Buildings/Building_ARANutrientDispenser/Patch_DispenserFoodSearch.cs

using HarmonyLib;
using RimWorld;
using Verse;
using System.Linq;

namespace ArachnaeSwarm
{
    [HarmonyPatch(typeof(FoodUtility), nameof(FoodUtility.BestFoodSourceOnMap))]
    public static class Patch_DispenserFoodSearch
    {
        // A state variable to pass our found dispenser from Prefix to Postfix
        private static Building_ARANutrientDispenser foundCustomDispenser = null;

        /// <summary>
        /// Runs BEFORE the original BestFoodSourceOnMap.
        /// Its job is to specifically find OUR custom dispenser.
        /// </summary>
        [HarmonyPrefix]
        public static bool Prefix(Pawn getter, Pawn eater, bool desperate = false, FoodPreferability maxPref = FoodPreferability.MealLavish)
        {
            // Reset state at the beginning of each check
            foundCustomDispenser = null;

            // First, check if our custom dispenser can even be a candidate.
            // We find a representative def for our dispenser to get the meal properties.
            // This is a bit of a hack, but necessary to know the preferability without an instance.
            var representativeDef = ThingDef.Named("ARANutrientDispenser"); // Assuming this is the defName of your dispenser
            var props = representativeDef?.GetModExtension<ARAFoodDispenserProperties>();
            var customMealDef = props?.thingToDispense;

            if (customMealDef == null) {
                // If we can't find our meal def, something is wrong with the XML.
                // It's safer to not interfere.
                return true; 
            }

            // Don't interfere if the pawn is looking for something better than what we can offer.
            if (maxPref < customMealDef.ingestible.preferability)
            {
                return true; // Let the original method run, it's looking for fancy food.
            }
            
            // Find the best available custom dispenser on the map
            var bestDispenser = FindBestCustomDispenser(getter, eater, desperate);

            if (bestDispenser != null)
            {
                // We found one! Store it for the Postfix.
                foundCustomDispenser = bestDispenser;
                ArachnaeLog.Debug($"[ArachnaeSwarm Prefix] Found a potential custom dispenser for {eater.LabelShort}: {bestDispenser.Label}");
            }
            
            // ALWAYS let the original method run. 
            // This allows us to compare its result with ours in the Postfix, ensuring maximum compatibility.
            return true;
        }

        /// <summary>
        /// Runs AFTER the original BestFoodSourceOnMap.
        /// Compares the original method's result with our custom dispenser (if any) and picks the best one.
        /// </summary>
        [HarmonyPostfix]
        public static void Postfix(ref Thing __result, ref ThingDef foodDef, Pawn eater, Pawn getter)
        {
            if (foundCustomDispenser == null)
            {
                // Our prefix didn't find any custom dispensers, so we don't need to do anything.
                return;
            }

            var customDispenserMealDef = foundCustomDispenser.DispensableDef;
            if (customDispenserMealDef == null)
            {
                ArachnaeLog.Debug($"[ArachnaeSwarm Postfix] Custom dispenser {foundCustomDispenser.Label} has a null DispensableDef.");
                return;
            }

            // If the original method found NO food, then our dispenser is the best (and only) choice.
            if (__result == null)
            {
                ArachnaeLog.Debug($"[ArachnaeSwarm Postfix] Original method found no food. Using our custom dispenser: {foundCustomDispenser.Label}");
                __result = foundCustomDispenser;
                foodDef = customDispenserMealDef;
                return;
            }

            // Both we and the original method found food. We must now choose the better one.
            // "FoodOptimality" is the score RimWorld uses to decide. Higher is better.
            float ourScore = FoodUtility.FoodOptimality(eater, foundCustomDispenser, customDispenserMealDef, (getter.Position - foundCustomDispenser.Position).LengthManhattan);
            float theirScore = FoodUtility.FoodOptimality(eater, __result, foodDef, (getter.Position - __result.Position).LengthManhattan);

            ArachnaeLog.Debug($"[ArachnaeSwarm Postfix] Comparing food sources: Our Dispenser (Score: {ourScore:F2}) vs Original Result '{__result.Label}' (Score: {theirScore:F2}).");

            // If our dispenser is a better choice, override the result.
            if (ourScore > theirScore)
            {
                ArachnaeLog.Debug($"[ArachnaeSwarm Postfix] Our dispenser is better. Overriding result.");
                __result = foundCustomDispenser;
                foodDef = customDispenserMealDef;
            }
            else
            {
                ArachnaeLog.Debug($"[ArachnaeSwarm Postfix] Original result is better or equal. Keeping it.");
            }
        }

        /// <summary>
        /// A helper to find the best custom dispenser for a pawn.
        /// </summary>
        private static Building_ARANutrientDispenser FindBestCustomDispenser(Pawn getter, Pawn eater, bool desperate)
        {
            if (getter.Map == null) return null;

            var allCustomDispensers = getter.Map.listerBuildings.AllBuildingsColonistOfClass<Building_ARANutrientDispenser>();
            
            Building_ARANutrientDispenser bestDispenser = null;
            float bestScore = float.MinValue;

            foreach (var dispenser in allCustomDispensers)
            {
                var currentMealDef = dispenser.DispensableDef;
                if (currentMealDef == null) continue;

                // --- Start of refined access checks ---

                // 1. Block hostile pawns.
                if (getter.HostileTo(Faction.OfPlayer))
                {
                    continue;
                }

                // 2. Block wild animals. Tamed player-faction animals are allowed.
                if (getter.RaceProps.Animal && getter.Faction != Faction.OfPlayer)
                {
                    continue;
                }
                
                // 3. Use the game's built-in social properness check for things like prisoner access.
                if (!dispenser.IsSociallyProper(getter))
                {
                    continue;
                }

                // --- End of refined access checks ---

                // Check if the dispenser is usable
                if (!dispenser.CanDispenseNow || dispenser.IsForbidden(getter) || !eater.WillEat(currentMealDef, getter))
                {
                    continue;
                }

                // Check reachability
                if (!getter.CanReach(dispenser, Verse.AI.PathEndMode.InteractionCell, Danger.Some))
                {
                    continue;
                }
                
                // Calculate score
                float score = FoodUtility.FoodOptimality(eater, dispenser, currentMealDef, (getter.Position - dispenser.Position).LengthManhattan);
                if (score > bestScore)
                {
                    bestScore = score;
                    bestDispenser = dispenser;
                }
            }

            return bestDispenser;
        }
    }
}