package net.glowstone.entity;

import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import net.glowstone.chunk.GlowChunk;
import net.glowstone.entity.physics.BoundingBox;
import org.bukkit.Chunk;
import org.bukkit.Location;
import org.bukkit.entity.Entity;

/**
 * A class which manages all of the entities within a world.
 *
 * @author Graham Edgecombe
 */
public class EntityManager implements Iterable<GlowEntity> {

    /**
     * A map of all the entity ids to the corresponding entities.
     */
    private final Map<Integer, GlowEntity> entities = new ConcurrentHashMap<>();

    /**
     * A map of entity types to a set containing all entities of that type.
     */
    private final Map<Class<? extends GlowEntity>, Set<? extends GlowEntity>> groupedEntities = new ConcurrentHashMap<>();

    /**
     * Gets all entities with the specified type.
     *
     * @param type The {@link Class} for the type.
     * @param <T> The type of entity.
     * @return A collection of entities with the specified type.
     */
    @SuppressWarnings("unchecked")
    public <T extends GlowEntity> Collection<T> getAll(Class<T> type) {
        Set<T> set = (Set<T>) groupedEntities.get(type);
        if (set == null) {
            set = new HashSet<>();
            groupedEntities.put(type, set);
        }
        return set;
    }

    /**
     * Gets all entities.
     *
     * @return A collection of entities.
     */
    public Collection<GlowEntity> getAll() {
        return entities.values();
    }

    /**
     * Gets an entity by its id.
     *
     * @param id The id.
     * @return The entity, or {@code null} if it could not be found.
     */
    public GlowEntity getEntity(int id) {
        return entities.get(id);
    }

    /**
     * Registers the entity to this world.
     *
     * @param entity The entity.
     */
    @SuppressWarnings("unchecked")
    void register(GlowEntity entity) {
        if (entity.id == 0) {
            throw new IllegalStateException("Entity has not been assigned an id.");
        }
        entities.put(entity.id, entity);
        ((Collection<GlowEntity>) getAll(entity.getClass())).add(entity);
        ((GlowChunk) entity.location.getChunk()).getRawEntities().add(entity);
    }

    /**
     * Unregister the entity to this world.
     *
     * @param entity The entity.
     */
    void unregister(GlowEntity entity) {
        entities.remove(entity.id);
        getAll(entity.getClass()).remove(entity);
        ((GlowChunk) entity.location.getChunk()).getRawEntities().remove(entity);
    }

    /**
     * Notes that an entity has moved from one location to another for physics and storage purposes.
     *
     * @param entity The entity.
     * @param newLocation The new location.
     */
    void move(GlowEntity entity, Location newLocation) {
        Chunk prevChunk = entity.location.getChunk();
        Chunk newChunk = newLocation.getChunk();
        if (prevChunk != newChunk) {
            ((GlowChunk) prevChunk).getRawEntities().remove(entity);
            ((GlowChunk) newChunk).getRawEntities().add(entity);
        }
    }

    @Override
    public Iterator<GlowEntity> iterator() {
        return entities.values().iterator();
    }

    public List<Entity> getEntitiesInside(BoundingBox searchBox, GlowEntity except) {
        // todo: narrow search based on the box's corners
        return entities.values().stream()
            .filter(entity -> entity != except && entity.intersects(searchBox))
            .collect(Collectors.toCollection(LinkedList::new));
    }
}

@@ -24,16 +24,165 @@
 ntity;%0A%0A
+import static com.google.common.collect.Multimaps.newSetMultimap;%0A%0Aimport com.google.common.collect.Multimap;%0Aimport com.google.common.collect.Sets;%0A
 import j
@@ -219,23 +219,27 @@
 va.util.
-HashSet
+Collections
 ;%0Aimport
@@ -337,30 +337,8 @@
 ap;%0A
-import java.util.Set;%0A
 impo
@@ -1027,16 +1027,21 @@
  final M
+ultim
 ap%3CClass
@@ -1067,22 +1067,8 @@
 ty%3E,
- Set%3C? extends
  Glo
@@ -1075,17 +1075,16 @@
 wEntity%3E
-%3E
  grouped
@@ -1095,35 +1095,183 @@
 ties
- = new ConcurrentHashMap%3C%3E(
+%0A            = newSetMultimap(%0A                    new ConcurrentHashMap%3CClass%3C? extends GlowEntity%3E, Collection%3CGlowEntity%3E%3E(),%0A                    Sets::newConcurrentHashSet
 );%0A%0A
@@ -1277,35 +1277,38 @@
 %0A    /**%0A     * 
-G
+R
 et
+urn
 s all entities w
@@ -1619,80 +1619,51 @@
     
-Set%3CT%3E set = (Set%3CT%3E) groupedEntities.get(type);%0A        if (set == null
+if (GlowEntity.class.isAssignableFrom(type)
 ) %7B%0A
@@ -1678,42 +1678,30 @@
     
-set = new HashSet%3C%3E();%0A           
+return (Collection%3CT%3E)
  gro
@@ -1717,21 +1717,16 @@
 ies.
-pu
+ge
 t(type
-, set
 );%0A 
@@ -1725,33 +1725,44 @@
 type);%0A        %7D
-%0A
+ else %7B%0A    
         return s
@@ -1764,12 +1764,42 @@
 urn 
-set;
+Collections.emptyList();%0A        %7D
 %0A   
@@ -2560,40 +2560,27 @@
     
-((Collection%3CGlowEntity%3E) getAll
+groupedEntities.put
 (ent
@@ -2597,15 +2597,10 @@
 ss()
-)).add(
+, 
 enti
@@ -2876,21 +2876,37 @@
        g
-etAll
+roupedEntities.remove
 (entity.
@@ -2915,25 +2915,18 @@
 tClass()
-).remove(
+, 
 entity);
@@ -3671,24 +3671,386 @@
 r();%0A    %7D%0A%0A
+    /**%0A     * Returns all entities that are inside or partly inside the given bounding box, with optionally%0A     * one exception.%0A     * @param searchBox the bounding box to search inside%0A     * @param except the entity to exclude, or null to include all%0A     * @return the entities contained in or touching %7B@code searchBox%7D, other than %7B@code except%7D%0A     */%0A
     public L

Refactor EntityManager to use a Multimap (#636)