package net.minecraft.world.level.entity;

import com.mojang.logging.LogUtils;
import java.util.Collection;
import java.util.stream.Stream;
import net.minecraft.util.AbortableIterationConsumer;
import net.minecraft.util.ClassInstanceMultiMap;
import net.minecraft.util.VisibleForDebug;
import net.minecraft.world.phys.AABB;
import org.slf4j.Logger;

public class EntitySection<T extends EntityAccess> {
   private static final Logger LOGGER = LogUtils.getLogger();
   private final ClassInstanceMultiMap<T> storage;
   private Visibility chunkStatus;

   public EntitySection(Class<T> p_156831_, Visibility p_156832_) {
      this.chunkStatus = p_156832_;
      this.storage = new ClassInstanceMultiMap<>(p_156831_);
   }

   public void add(T p_188347_) {
      this.storage.add(p_188347_);
   }

   public boolean remove(T p_188356_) {
      return this.storage.remove(p_188356_);
   }

   public AbortableIterationConsumer.Continuation getEntities(AABB p_262016_, AbortableIterationConsumer<T> p_261863_) {
      for(T t : this.storage) {
         if (t.getBoundingBox().intersects(p_262016_) && p_261863_.accept(t).shouldAbort()) {
            return AbortableIterationConsumer.Continuation.ABORT;
         }
      }

      return AbortableIterationConsumer.Continuation.CONTINUE;
   }

   public <U extends T> AbortableIterationConsumer.Continuation getEntities(EntityTypeTest<T, U> p_188349_, AABB p_188350_, AbortableIterationConsumer<? super U> p_261535_) {
      Collection<? extends T> collection = this.storage.find(p_188349_.getBaseClass());
      if (collection.isEmpty()) {
         return AbortableIterationConsumer.Continuation.CONTINUE;
      } else {
         for(T t : collection) {
            U u = (U)((EntityAccess)p_188349_.tryCast(t));
            if (u != null && t.getBoundingBox().intersects(p_188350_) && p_261535_.accept(u).shouldAbort()) {
               return AbortableIterationConsumer.Continuation.ABORT;
            }
         }

         return AbortableIterationConsumer.Continuation.CONTINUE;
      }
   }

   public boolean isEmpty() {
      return this.storage.isEmpty();
   }

   public Stream<T> getEntities() {
      return this.storage.stream();
   }

   public Visibility getStatus() {
      return this.chunkStatus;
   }

   public Visibility updateChunkStatus(Visibility p_156839_) {
      Visibility visibility = this.chunkStatus;
      this.chunkStatus = p_156839_;
      return visibility;
   }

   @VisibleForDebug
   public int size() {
      return this.storage.size();
   }
}