path: root/src/main/java/org/chrisoft/trashyaddon/commands/CEntitySelector.java

package org.chrisoft.trashyaddon.commands;

import com.google.common.collect.Lists;
import com.mojang.brigadier.exceptions.CommandSyntaxException;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;

import net.minecraft.client.MinecraftClient;
import net.minecraft.client.world.ClientWorld;
import net.minecraft.command.argument.EntityArgumentType;
import net.minecraft.entity.Entity;
import net.minecraft.entity.EntityType;
import net.minecraft.predicate.NumberRange;
import net.minecraft.text.Text;
import net.minecraft.text.Texts;
import net.minecraft.util.TypeFilter;
import net.minecraft.util.function.LazyIterationConsumer;
import net.minecraft.util.math.Box;
import net.minecraft.util.math.Vec3d;
import org.chrisoft.trashyaddon.mixin.ClientWorldMixin;
import org.jetbrains.annotations.Nullable;

public class CEntitySelector {
   public static final int MAX_VALUE = Integer.MAX_VALUE;
   public static final BiConsumer<Vec3d, List<? extends Entity>> ARBITRARY = (pos, entities) -> {
   };
   private static final TypeFilter<Entity, ?> PASSTHROUGH_FILTER = new TypeFilter<Entity, Entity>() {
      public Entity downcast(Entity arg) {
         return arg;
      }

      @Override
      public Class<? extends Entity> getBaseClass() {
         return Entity.class;
      }
   };
   private final int limit;
   private final boolean includesNonPlayers;
   private final boolean localWorldOnly;
   private final Predicate<Entity> basePredicate;
   private final NumberRange.DoubleRange distance;
   private final Function<Vec3d, Vec3d> positionOffset;
   @Nullable
   private final Box box;
   private final BiConsumer<Vec3d, List<? extends Entity>> sorter;
   private final boolean senderOnly;
   @Nullable
   private final String playerName;
   @Nullable
   private final UUID uuid;
   private final TypeFilter<Entity, ?> entityFilter;
   private final boolean usesAt;

   public CEntitySelector(
      int count,
      boolean includesNonPlayers,
      boolean localWorldOnly,
      Predicate<Entity> basePredicate,
      NumberRange.DoubleRange distance,
      Function<Vec3d, Vec3d> positionOffset,
      @Nullable Box box,
      BiConsumer<Vec3d, List<? extends Entity>> sorter,
      boolean senderOnly,
      @Nullable String playerName,
      @Nullable UUID uuid,
      @Nullable EntityType<?> type,
      boolean usesAt
   ) {
      this.limit = count;
      this.includesNonPlayers = includesNonPlayers;
      this.localWorldOnly = localWorldOnly;
      this.basePredicate = basePredicate;
      this.distance = distance;
      this.positionOffset = positionOffset;
      this.box = box;
      this.sorter = sorter;
      this.senderOnly = senderOnly;
      this.playerName = playerName;
      this.uuid = uuid;
      this.entityFilter = (TypeFilter<Entity, ?>)(type == null ? PASSTHROUGH_FILTER : type);
      this.usesAt = usesAt;
   }

   public int getLimit() {
      return this.limit;
   }

   public boolean includesNonPlayers() {
      return this.includesNonPlayers;
   }

   public boolean isSenderOnly() {
      return this.senderOnly;
   }

   public boolean isLocalWorldOnly() {
      return this.localWorldOnly;
   }

   public boolean usesAt() {
      return this.usesAt;
   }

   private int getAppendLimit() {
      return this.sorter == ARBITRARY ? this.limit : Integer.MAX_VALUE;
   }

   private Predicate<Entity> getPositionPredicate(Vec3d pos) {
      Predicate<Entity> predicate = this.basePredicate;
      if (this.box != null) {
         Box lv = this.box.offset(pos);
         predicate = predicate.and(entity -> lv.intersects(entity.getBoundingBox()));
      }

      if (!this.distance.isDummy()) {
         predicate = predicate.and(entity -> this.distance.testSqrt(entity.squaredDistanceTo(pos)));
      }

      return predicate;
   }
   public List<? extends Entity> getClientSideEntityMatches(ClientWorld cw, Vec3d sourcePos) throws CommandSyntaxException {
       Predicate<Entity> p = this.getPositionPredicate(sourcePos);
       List<Entity> ret = new ArrayList<>();
       int limit = this.getAppendLimit();
       if (this.box != null)
           cw.collectEntitiesByType(this.entityFilter, this.box.offset(sourcePos), p, ret, limit);
       else {
           ((ClientWorldMixin)cw).invokeGetEntityLookup().forEach(this.entityFilter, (e) -> {
               if (p.test(e)) {
                   ret.add(e);
                   if (ret.size() >= limit)
                       return LazyIterationConsumer.NextIteration.ABORT;
               }
               return LazyIterationConsumer.NextIteration.CONTINUE;
           });
       }
       return ret;
   }

   private <T extends Entity> List<T> getEntities(Vec3d pos, List<T> entities) {
      if (entities.size() > 1) {
         this.sorter.accept(pos, entities);
      }

      return entities.subList(0, Math.min(this.limit, entities.size()));
   }

   public static Text getNames(List<? extends Entity> entities) {
      return Texts.join(entities, Entity::getDisplayName);
   }
}