package me.earthme.luminol.mobaireducer;

import com.google.common.collect.MapMaker;
import me.earthme.luminol.config.modules.optimizations.MobAiReducerConfig;
import net.minecraft.tags.ItemTags;
import net.minecraft.tags.TagKey;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.PathfinderMob;
import net.minecraft.world.entity.ai.goal.BreedGoal;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.ai.goal.TemptGoal;
import net.minecraft.world.entity.ai.goal.WrappedGoal;
import net.minecraft.world.entity.ai.targeting.TargetingConditions;
import net.minecraft.world.entity.animal.*;
import net.minecraft.world.entity.animal.armadillo.Armadillo;
import net.minecraft.world.entity.animal.axolotl.Axolotl;
import net.minecraft.world.entity.animal.camel.Camel;
import net.minecraft.world.entity.animal.frog.Frog;
import net.minecraft.world.entity.animal.goat.Goat;
import net.minecraft.world.entity.animal.horse.Horse;
import net.minecraft.world.entity.animal.horse.Llama;
import net.minecraft.world.entity.animal.sheep.Sheep;
import net.minecraft.world.entity.animal.sniffer.Sniffer;
import net.minecraft.world.entity.animal.wolf.Wolf;
import net.minecraft.world.entity.monster.Strider;
import net.minecraft.world.entity.monster.hoglin.Hoglin;
import net.minecraft.world.entity.monster.piglin.Piglin;
import net.minecraft.world.item.Item;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public final class MobAiReducerHandler {

    private static final Logger LOGGER = LoggerFactory.getLogger("NekoCore/MobAiReducer");
    private static final Map<PathfinderMob, Boolean> OPTIMIZED = new MapMaker().weakKeys().concurrencyLevel(4).makeMap();
    private static final Map<Class<? extends Entity>, TargetingConditions> TEMPT_CONDITIONS = new ConcurrentHashMap<>();
    private static TargetingConditions breedConditions;

    private MobAiReducerHandler() {}

    public static void register() {
        if (!MobAiReducerConfig.enabled) return;

        breedConditions = TargetingConditions.forNonCombat().ignoreLineOfSight().range(MobAiReducerConfig.breedRange);

        registerTempt(Horse.class,    ItemTags.HORSE_FOOD);
        registerTempt(Cow.class,      ItemTags.COW_FOOD);
        registerTempt(Sheep.class,    ItemTags.SHEEP_FOOD);
        registerTempt(Fox.class,      ItemTags.FOX_FOOD);
        registerTempt(Pig.class,      ItemTags.PIG_FOOD);
        registerTempt(Chicken.class,  ItemTags.CHICKEN_FOOD);
        registerTempt(Parrot.class,   ItemTags.PARROT_FOOD);
        registerTempt(Frog.class,     ItemTags.FROG_FOOD);
        registerTempt(Axolotl.class,  ItemTags.AXOLOTL_FOOD);
        registerTempt(Goat.class,     ItemTags.GOAT_FOOD);
        registerTempt(Bee.class,      ItemTags.BEE_FOOD);
        registerTempt(Wolf.class,     ItemTags.WOLF_FOOD);
        registerTempt(Turtle.class,   ItemTags.TURTLE_FOOD);
        registerTempt(Strider.class,  ItemTags.STRIDER_FOOD);
        registerTempt(Rabbit.class,   ItemTags.RABBIT_FOOD);
        registerTempt(Piglin.class,   ItemTags.PIGLIN_FOOD);
        registerTempt(Panda.class,    ItemTags.PANDA_FOOD);
        registerTempt(Ocelot.class,   ItemTags.OCELOT_FOOD);
        registerTempt(Llama.class,    ItemTags.LLAMA_FOOD);
        registerTempt(Hoglin.class,   ItemTags.HOGLIN_FOOD);
        registerTempt(Camel.class,    ItemTags.CAMEL_FOOD);
        registerTempt(Armadillo.class,ItemTags.ARMADILLO_FOOD);
        registerTempt(Sniffer.class,  ItemTags.SNIFFER_FOOD);

        LOGGER.info("MobAiReducer registered for {} entity types", TEMPT_CONDITIONS.size());
    }

    private static void registerTempt(Class<? extends Entity> cls, TagKey<Item> foodTag) {
        boolean bothHands = MobAiReducerConfig.temptTriggerBothHands;
        TEMPT_CONDITIONS.computeIfAbsent(cls, k -> TargetingConditions.forNonCombat().ignoreLineOfSight())
            .range(MobAiReducerConfig.temptRange)
            .selector((entity, level) ->
                entity.getMainHandItem().is(foodTag) || (bothHands && entity.getOffhandItem().is(foodTag))
            );
    }

    public static void optimize(PathfinderMob mob) {
        if (OPTIMIZED.containsKey(mob)) return;

        if (!MobAiReducerConfig.enabled) return;

        mob.collides = !MobAiReducerConfig.disableCollides;
        mob.setSilent(MobAiReducerConfig.silent);
        OPTIMIZED.put(mob, Boolean.TRUE);

        boolean isAnimal = mob instanceof Animal;
        TargetingConditions temptCond = MobAiReducerConfig.temptEnabled ? TEMPT_CONDITIONS.get(mob.getClass()) : null;

        Set<WrappedGoal> goals = mob.goalSelector.getAvailableGoals();
        synchronized (goals) {
            Set<WrappedGoal> toAdd = new HashSet<>(), toRemove = new HashSet<>();
            for (WrappedGoal wg : goals) {
                Goal goal = wg.getGoal();
                Class<?> goalClass = goal.getClass();

                if (MobAiReducerConfig.keepDedicated && !goalClass.getName().contains("ai.goal")) continue;

                if (isAnimal && MobAiReducerConfig.breedEnabled && goalClass == BreedGoal.class) {
                    wg.stop();
                    toRemove.add(wg);
                    toAdd.add(new WrappedGoal(wg.getPriority(), new OptimizedBreedGoal(new MobAiReducerConfig(), (Animal) mob, breedConditions)));
                    continue;
                }
                if (MobAiReducerConfig.temptEnabled && goalClass == TemptGoal.class && temptCond != null) {
                    wg.stop();
                    toRemove.add(wg);
                    toAdd.add(new WrappedGoal(wg.getPriority(), new OptimizedTemptGoal(new MobAiReducerConfig(), mob, temptCond)));
                    continue;
                }

                String simpleName = goalClass.getSimpleName();
                boolean inList = MobAiReducerConfig.aiGoalList.stream().anyMatch(simpleName::contains);
                if (inList == MobAiReducerConfig.aiListMode) {
                    wg.stop();
                    toRemove.add(wg);
                }
            }
            goals.removeAll(toRemove);
            goals.addAll(toAdd);
        }
    }

    public static void purge() {
        OPTIMIZED.keySet().removeIf(mob -> !mob.isAlive() || !mob.valid);
    }
}
