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.lang.reflect.Modifier;
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 -> {
                try {
                    priorityField.set(goal, ((int) priorityField.get(goal)) + 1);
                    return true;
                } catch(Exception e) {
                    InfinityLib.instance.getLogger().error("Failed to increment priority");
                    InfinityLib.instance.getLogger().printStackTrace(e);
                    return false;
                }
            });
        } 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_220892_d");
        } catch (Exception e) {
            InfinityLib.instance.getLogger().error("Could not retrieve Goals field from GoalSelector.class");
            InfinityLib.instance.getLogger().printStackTrace(e);
            return null;
        }
    }

    //TODO: fix this if reflection does not work
    private static Field initPriorityField() {
        try {
            // Retrieve field
            Field field = ObfuscationReflectionHelper.findField(WrappedGoal.class, "field_220775_b");
            // Remove final modifier
            Field modifiersField = Field.class.getDeclaredField("modifiers");
            modifiersField.setAccessible(true);
            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);
            // 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;
        }
    }
}

@@ -426,43 +426,8 @@
 ld;%0A
-import java.lang.reflect.Modifier;%0A
 impo
@@ -1532,34 +1532,123 @@
     
-try %7B%0A                    
+boolean success = false;%0A                try %7B%0A                    success = UnsafeUtil.getInstance().replaceField(
 prio
@@ -1660,13 +1660,10 @@
 ield
-.set(
+, 
 goal
@@ -1706,41 +1706,8 @@
 1);%0A
-                    return true;%0A
     
@@ -1922,33 +1922,17 @@
         
-    return false;
+%7D
 %0A       
@@ -1932,33 +1932,47 @@
                 
-%7D
+return success;
 %0A            %7D);
@@ -2756,57 +2756,8 @@
  %7D%0A%0A
-    //TODO: fix this if reflection does not work%0A
     
@@ -2964,107 +2964,22 @@
  // 
-Remove final modifier%0A            Field modifiersField = Field.class.getDeclaredField(%22modifiers%22);
+Set accessible
 %0A   
@@ -2983,34 +2983,25 @@
             
-modifiersF
+f
 ield.setAcce
@@ -3018,90 +3018,8 @@
 e);%0A
-            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);%0A
     

Fix EntityHelper Reflection