bazaarflux/src/main/java/me/lethunderhawk/custom/entity/CustomEntityRegistry.java

package me.lethunderhawk.custom.entity;

import me.lethunderhawk.fluxapi.FluxService;
import me.lethunderhawk.main.BazaarFlux;
import org.bukkit.entity.Entity;
import org.bukkit.scheduler.BukkitRunnable;
import org.jetbrains.annotations.NotNull;

import java.util.*;

public class CustomEntityRegistry {
    private static Map<Entity, CustomEntity> entities = new HashMap<>();
    private static Map<Entity, Integer> indicators = new HashMap<>();

    static {
        new BukkitRunnable() {
            Set<Entity> armorStands = indicators.keySet();
            List<Entity> removals = new ArrayList<>();
            @Override
            public void run() {
                for(Entity stand : armorStands) {
                    int ticksLeft = indicators.get(stand);
                    if(ticksLeft == 0) {
                        stand.remove();
                        removals.add(stand);
                        continue;
                    }
                    ticksLeft--;
                    indicators.put(stand, ticksLeft);
                }
                armorStands.removeAll(removals);
            }
        }.runTaskTimer(FluxService.get(BazaarFlux.class), 0L, 1L);
    }

    public static CustomEntity remove(@NotNull Entity entity) {
        return entities.remove(entity);
    }

    public static void add(@NotNull Entity entity, @NotNull CustomEntity customEntity) {
        entities.put(entity, customEntity);
    }

    public static boolean contains(@NotNull Entity entity) {
        return entities.containsKey(entity);
    }

    public static CustomEntity get(Entity entity) {
        return entities.get(entity);
    }

    public static Map<Entity, Integer> getIndicators() {
        return indicators;
    }
}