package com.infinityraider.infinitylib.utility;

import com.infinityraider.infinitylib.InfinityLib;
import net.minecraft.world.entity.Mob;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.ai.goal.GoalSelector;
import net.minecraft.world.entity.ai.goal.WrappedGoal;
import net.minecraftforge.fml.util.ObfuscationReflectionHelper;

import javax.annotation.Nullable;
import java.lang.reflect.Field;
import java.util.Set;

public class EntityHelper {
    private static final Field goalsField;
    private static final Field priorityField;

    public static boolean injectGoal(Mob entity, Goal goal, int priority) {
        if(incrementPriorities(entity.goalSelector, priority)) {
            entity.goalSelector.addGoal(priority, goal);
            return true;
        }
        return false;
    }

    @SuppressWarnings("unchecked")
    private static boolean incrementPriorities(GoalSelector selector, int priority) {
        if(goalsField == null) {
            InfinityLib.instance.getLogger().error("Could not increment priorities, goals field not retrieved");
            return false;
        }
        if(priorityField == null) {
            InfinityLib.instance.getLogger().error("Could not increment priorities, priority field not retrieved");
            return false;
        }
        try {
            Set<WrappedGoal> goals = (Set<WrappedGoal>) goalsField.get(selector);
            return goals.stream().filter(goal -> goal.getPriority() >= priority).allMatch(goal -> {
                boolean success = false;
                try {
                    success = UnsafeUtil.getInstance().replaceField(priorityField, goal, ((int) priorityField.get(goal)) + 1);
                } catch(Exception e) {
                    InfinityLib.instance.getLogger().error("Failed to increment priority");
                    InfinityLib.instance.getLogger().printStackTrace(e);
                }
                return success;
            });
        } catch (Exception e) {
            InfinityLib.instance.getLogger().error("Encountered error while getting the goals from the GoalSelector");
            InfinityLib.instance.getLogger().printStackTrace(e);
            return false;
        }
    }

    static {
        goalsField = initGoalsField();
        priorityField = initPriorityField();
    }

    @Nullable
    private static Field initGoalsField() {
        try {
            return ObfuscationReflectionHelper.findField(GoalSelector.class, "field_6461");
        } catch (Exception e) {
            InfinityLib.instance.getLogger().error("Could not retrieve Goals field from GoalSelector.class");
            InfinityLib.instance.getLogger().printStackTrace(e);
            return null;
        }
    }

    private static Field initPriorityField() {
        try {
            // Retrieve field
            Field field = ObfuscationReflectionHelper.findField(WrappedGoal.class, "field_18417");
            // Set accessible
            field.setAccessible(true);
            // Return the field
            return field;
        } catch (Exception e) {
            InfinityLib.instance.getLogger().error("Could not retrieve Priority field from PrioritizedGoal.class");
            InfinityLib.instance.getLogger().printStackTrace(e);
            return null;
        }
    }
}

@@ -294,72 +294,8 @@
 oal;
-%0Aimport net.minecraftforge.fml.util.ObfuscationReflectionHelper;
 %0A%0Aim
@@ -358,16 +358,41 @@
 .Field;%0A
+import java.util.Arrays;%0A
 import j
@@ -2366,138 +2366,262 @@
     
-try %7B%0A            return ObfuscationReflectionHelper.findField(GoalSelector.class, %22field_6461%22);%0A        %7D catch (Exception e) %7B%0A
+return Arrays.stream(GoalSelector.class.getDeclaredFields())%0A                .peek(field -%3E field.setAccessible(true))%0A                .filter(field -%3E field.getType() == Set.class)%0A                .findFirst()%0A                .orElseGet(() -%3E %7B%0A        
     
@@ -2738,65 +2738,8 @@
     
-    InfinityLib.instance.getLogger().printStackTrace(e);%0A
     
@@ -2763,33 +2763,43 @@
 ll;%0A        
-%7D
+        %7D);
 %0A    %7D%0A%0A
     private 
@@ -2782,24 +2782,38 @@
  %7D);%0A    %7D%0A%0A
+    @Nullable%0A
     private 
@@ -2859,184 +2859,98 @@
     
-try %7B%0A            // Retrieve field%0A            Field field = ObfuscationReflectionHelper.findField(WrappedGoal.class, %22field_18417%22);%0A            // Set accessible%0A           
+return Arrays.stream(WrappedGoal.class.getDeclaredFields())%0A                .peek(field -%3E
  fie
@@ -2971,18 +2971,19 @@
 le(true)
-;%0A
+)%0A 
         
@@ -2990,86 +2990,130 @@
     
-// Return the field%0A            return field;%0A        %7D catch (Exception e) %7B%0A
+   .filter(field -%3E field.getType() == int.class)%0A                .findFirst()%0A                .orElseGet(() -%3E %7B%0A        
     
@@ -3236,65 +3236,8 @@
     
-    InfinityLib.instance.getLogger().printStackTrace(e);%0A
     
@@ -3261,22 +3261,32 @@
 ll;%0A        
-%7D
+        %7D);
 %0A    %7D%0A%7D%0A

Fix Entity AI injection