package nl.dykam.dev.reutil.data;

import org.bukkit.Chunk;
import org.bukkit.block.Block;
import org.bukkit.entity.Player;

import java.util.*;

class MapComponentStorage implements ComponentStorage {
    // Quick&Dirty.
    Map<Object, ComponentMap> map = new HashMap<>();

    @Override
    @SuppressWarnings("unchecked")
    public <T extends Component> T get(Object object, Class<T> type) {
        return (T) getComponentMap(object).get(type);
    }

    @SuppressWarnings("unchecked")
    private ComponentMap getComponentMap(Object object) {
        if(!map.containsKey(object))
            map.put(object, new ComponentMap());
        return map.get(object);
    }

    @Override
    public <T extends Component> T get(Player player, Class<T> type) {
        return get((Object)player, type);
    }

    @Override
    public <T extends Component> T get(Chunk chunk, Class<T> type) {
        return get((Object)chunk, type);
    }

    @Override
    public <T extends Component> T get(Block block, Class<T> type) {
        return get((Object)block, type);
    }

    @Override
    public <T extends Component> void set(Object object, Class<T> type, T component) {
        ComponentMap componentMap = getComponentMap(object);
        componentMap.put(type, component);
    }

    @Override
    public <T extends Component> void set(Player player, Class<T> type, T component) {
        set((Object)player, type, component);
    }

    @Override
    public <T extends Component> void set(Chunk chunk, Class<T> type, T component) {
        set((Object)chunk, type, component);
    }

    @Override
    public <T extends Component> void set(Block block, Class<T> type, T component) {
        set((Object)block, type, component);
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T extends Component> T remove(Object object, Class<T> type) {
        ComponentMap componentMap = getComponentMap(object);
        return (T) componentMap.remove(type);
    }

    @Override
    public <T extends Component> T remove(Player player, Class<T> type) {
        return remove((Object)player, type);
    }

    @Override
    public <T extends Component> T remove(Chunk chunk, Class<T> type) {
        return remove((Object)chunk, type);
    }

    @Override
    public <T extends Component> T remove(Block block, Class<T> type) {
        return remove((Object)block, type);
    }

    @Override
    public Iterator<Component> iterator() {
        return new Iterator<Component>() {
            private Iterator<ComponentMap> mapIterator;
            private Iterator<Component> componentIterator;

            {
                mapIterator = map.values().iterator();
                componentIterator = mapIterator.next().values().iterator();
            }

            @Override
            public boolean hasNext() {
                return componentIterator.hasNext() || mapIterator.hasNext();
            }

            @Override
            public Component next() {
                if(componentIterator.hasNext())
                    return componentIterator.next();
                if(!mapIterator.hasNext())
                    throw new NoSuchElementException();

                componentIterator = mapIterator.next().values().iterator();
                return next();
            }

            @Override
            public void remove() {

            }
        };
    }

    /**
     * Help the type inference
     */
    private static class ComponentMap extends HashMap<Class<? extends Component>, Component> {

    }
}

@@ -28,16 +28,42 @@
 .data;%0A%0A
+import org.bukkit.Bukkit;%0A
 import o
@@ -2699,84 +2699,8 @@
 ();%0A
-                componentIterator = mapIterator.next().values().iterator();%0A
     
@@ -2791,46 +2791,122 @@
     
-return componentIterator.hasNext() %7C%7C 
+if(componentIterator != null && componentIterator.hasNext())%0A                    return true;%0A                if(!
 mapI
@@ -2918,25 +2918,25 @@
 or.hasNext()
-;
+)
 %0A           
@@ -2940,108 +2940,129 @@
     
-%7D%0A%0A
+ 
        
-     @Override%0A            public Component next() %7B%0A                if(componentIterator.
+return false;%0A                componentIterator = mapIterator.next().values().iterator();%0A                return 
 hasN
@@ -3058,34 +3058,35 @@
 return hasNext()
-)%0A
+;%0A 
                 
@@ -3071,32 +3071,36 @@
 t();%0A           
+ %7D%0A%0A
           return
@@ -3097,40 +3097,57 @@
     
-return componentIterator.
+  @Override%0A            public Component 
 next()
-;
+ %7B
 %0A   
@@ -3159,36 +3159,24 @@
         if(!
-mapIterator.
 hasNext())%0A 
@@ -3226,25 +3226,24 @@
 xception();%0A
-%0A
             
@@ -3238,32 +3238,39 @@
                 
+return 
 componentIterato
@@ -3274,50 +3274,88 @@
 ator
- = mapIterator.next().values().iterator();
+.next();%0A            %7D%0A%0A            @Override%0A            public void remove() %7B
 %0A   
@@ -3371,22 +3371,22 @@
     
-return n
+if(!hasN
 ext()
-;
+)
 %0A   
@@ -3385,36 +3385,32 @@
 t())%0A           
- %7D%0A%0A
             @Ove
@@ -3406,53 +3406,86 @@
     
-   @Override%0A            public void 
+throw new NoSuchElementException();%0A                componentIterator.
 remove()
  %7B%0A%0A
@@ -3480,19 +3480,17 @@
 remove()
- %7B%0A
+;
 %0A       

Fix MapComponentStorage iteration