Herobrine-Rewoven/src/main/java/com/thebrokenrail/herobrine/entity/ai/HerobrineAI.java

package com.thebrokenrail.herobrine.entity.ai;

import com.thebrokenrail.herobrine.config.HardcodedConfig;
import com.thebrokenrail.herobrine.entity.HerobrineEntity;
import com.thebrokenrail.herobrine.entity.ai.stage.FireballStage;
import com.thebrokenrail.herobrine.entity.ai.stage.IdleStage;
import com.thebrokenrail.herobrine.entity.ai.stage.LightningStage;
import com.thebrokenrail.herobrine.entity.ai.stage.TeleportStage;
import net.minecraft.nbt.CompoundTag;

import java.util.ArrayList;
import java.util.List;

public class HerobrineAI {
    private static final AIStage.AIStageFactory<?>[] STAGES;

    private AIStage stage = null;
    private final HerobrineEntity entity;

    public HerobrineAI(HerobrineEntity entity) {
        this.entity = entity;
    }

    public void tick() {
        if (stage == null) {
            nextStage();
        }
        if (stage != null) {
            stage.tick();
        }
    }

    public void setStage(AIStage stage) {
        this.stage = stage;
        entity.handleStageChange();
    }

    public AIStage getStage() {
        return stage;
    }

    public void fromTag(CompoundTag tag) {
        String id = tag.getString("ID");
        AIStage.AIStageFactory<?> selected = null;
        for (AIStage.AIStageFactory<?> stage : STAGES) {
            if (stage.getID().equals(id)) {
                selected = stage;
                break;
            }
        }
        if (selected != null) {
            setStage(selected.create(entity));
            stage.fromTag(tag.getCompound("Data"));
        } else {
            setStage(null);
        }
    }

    public CompoundTag toTag() {
        CompoundTag tag = new CompoundTag();
        if (stage != null) {
            tag.putString("ID", stage.getID());
            tag.put("Data", stage.toTag());
        }
        return tag;
    }

    public void nextStage() {
        List<AIStage.AIStageFactory<?>> options = new ArrayList<>();
        for (AIStage.AIStageFactory<?> stage : STAGES) {
            if (stage.canStart(entity)) {
                options.add(stage);
            }
        }
        if (options.size() > 0) {
            setStage(options.get(entity.getRandom().nextInt(options.size())).create(entity));
        } else {
            setStage(null);
        }
    }

    static {
        STAGES = new AIStage.AIStageFactory<?>[]{
                new AIStage.AIStageFactory<AIStage>() {
                    @Override
                    public boolean canStart(HerobrineEntity entity) {
                        return false;
                    }

                    @Override
                    protected AIStage createInternal(HerobrineEntity entity) {
                        return new TeleportStage(entity);
                    }

                    @Override
                    public String getID() {
                        return "teleport";
                    }
                },
                new AIStage.AIStageFactory<AIStage>() {
                    @Override
                    public boolean canStart(HerobrineEntity entity) {
                        return entity.getRandom().nextDouble() <= HardcodedConfig.HEROBRINE_FIREBALL_CHANCE;
                    }

                    @Override
                    protected AIStage createInternal(HerobrineEntity entity) {
                        return new FireballStage(entity);
                    }

                    @Override
                    public String getID() {
                        return "fireball";
                    }
                },
                new AIStage.AIStageFactory<AIStage>() {
                    @Override
                    public boolean canStart(HerobrineEntity entity) {
                        return true;
                    }

                    @Override
                    protected AIStage createInternal(HerobrineEntity entity) {
                        return new IdleStage(entity);
                    }

                    @Override
                    public String getID() {
                        return "idle";
                    }
                },
                new AIStage.AIStageFactory<AIStage>() {
                    @Override
                    public boolean canStart(HerobrineEntity entity) {
                        return entity.getRandom().nextDouble() <= HardcodedConfig.HEROBRINE_LIGHTNING_CHANCE && (entity.getEntityWorld().isSkyVisible(entity.getBlockPos()) || entity.getEntityWorld().isSkyVisible(entity.getBlockPos().up()));
                    }

                    @Override
                    protected AIStage createInternal(HerobrineEntity entity) {
                        return new LightningStage(entity);
                    }

                    @Override
                    public String getID() {
                        return "lightning";
                    }
                }
        };
    }
}