package net.minecraft.util;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterators;
import java.util.Collections;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.function.Predicate;
import java.util.Collection;
import java.util.Iterator;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import java.util.List;
import java.util.Map;
import java.util.AbstractCollection;

public class ClassInstanceMultiMap<T> extends AbstractCollection<T> {
    private final Map<Class<?>, List<T>> byClass;
    private final Class<T> baseClass;
    private final List<T> allInstances;
    
    public ClassInstanceMultiMap(final Class<T> class1) {
        this.byClass = Maps.newHashMap();
        this.allInstances = Lists.newArrayList();
        this.baseClass = class1;
        this.byClass.put(class1, this.allInstances);
    }
    
    @Override
    public boolean add(final T object) {
        boolean boolean3 = false;
        for (final Map.Entry<Class<?>, List<T>> entry5 : this.byClass.entrySet()) {
            if (entry5.getKey().isInstance(object)) {
                boolean3 |= entry5.getValue().add(object);
            }
        }
        return boolean3;
    }
    
    @Override
    public boolean remove(final Object object) {
        boolean boolean3 = false;
        for (final Map.Entry<Class<?>, List<T>> entry5 : this.byClass.entrySet()) {
            if (entry5.getKey().isInstance(object)) {
                final List<T> list6 = entry5.getValue();
                boolean3 |= list6.remove(object);
            }
        }
        return boolean3;
    }
    
    @Override
    public boolean contains(final Object object) {
        return this.find(object.getClass()).contains(object);
    }
    
    public <S> Collection<S> find(final Class<S> class1) {
        if (!this.baseClass.isAssignableFrom(class1)) {
            throw new IllegalArgumentException("Don't know how to search for " + class1);
        }
        final List<T> list3 = this.byClass.computeIfAbsent(class1, class1 -> this.allInstances.stream().filter(class1::isInstance).collect(Collectors.toList()));
        return Collections.<S>unmodifiableCollection(list3);
    }
    
    @Override
    public Iterator<T> iterator() {
        if (this.allInstances.isEmpty()) {
            return Collections.<T>emptyIterator();
        }
        return Iterators.unmodifiableIterator(this.allInstances.iterator());
    }
    
    public List<T> getAllInstances() {
        return ImmutableList.copyOf(this.allInstances);
    }
    
    @Override
    public int size() {
        return this.allInstances.size();
    }
}
