package com.infinityraider.infinitylib.capability;

import com.infinityraider.infinitylib.InfinityLib;
import com.infinityraider.infinitylib.utility.ISerializable;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraftforge.common.capabilities.CapabilityManager;
import net.minecraftforge.common.capabilities.ICapabilityProvider;
import net.minecraftforge.event.AttachCapabilitiesEvent;
import net.minecraftforge.event.entity.player.PlayerEvent;
import net.minecraftforge.eventbus.api.SubscribeEvent;

import java.util.*;

public class CapabilityHandler {
    private static final CapabilityHandler INSTANCE = new CapabilityHandler();

    public static CapabilityHandler getInstance() {
        return INSTANCE;
    }

    private final Map<Class<? extends ICapabilityProvider>, List<ICapabilityImplementation<ICapabilityProvider, ? extends ISerializable>>> map;

    private CapabilityHandler() {
        this.map = new HashMap<>();
        InfinityLib.instance.registerEventHandler(this);
    }

    @SuppressWarnings("unchecked")
    public void registerCapability(ICapabilityImplementation implementation) {
        if(!this.map.containsKey(implementation.getCarrierClass())) {
            this.map.put(implementation.getCarrierClass(), new ArrayList<>());
        }
        this.map.get(implementation.getCarrierClass()).add(implementation);
        CapabilityManager.INSTANCE.register(implementation.getCapabilityClass(), new CapabilityStorage(), () -> null);
    }

    protected <T extends ICapabilityProvider, C extends ISerializable> void addCapability(
            AttachCapabilitiesEvent event, ICapabilityImplementation<T , C> implementation, T carrier) {
        C value = implementation.createNewValue(carrier);
        event.addCapability(implementation.getCapabilityKey(), new CapabilityProvider<>(implementation.getCapability(), value));
    }

    @SuppressWarnings("unchecked")
    protected  <T extends ICapabilityProvider> void addCapabilities(AttachCapabilitiesEvent<T> event) {
        Object object = event.getObject();
        T carrier = (T) object;
        Class<T> clazz = (Class<T>) carrier.getClass();
        this.map.entrySet().stream().filter(entry -> entry.getKey().isAssignableFrom(clazz)).forEach(
                entry -> entry.getValue().stream()
                        .filter(impl -> impl.getCarrierClass().isAssignableFrom(clazz))
                        .filter(impl -> impl.shouldApplyCapability(carrier))
                        .forEach(impl -> this.addCapability(event, impl, carrier)));
    }

    @SubscribeEvent
    @SuppressWarnings("unused")
    public void addEntityCapabilitiesRaw(AttachCapabilitiesEvent event) {
        if(event.getObject() instanceof ICapabilityProvider) {
            this.addCapabilities(event);
        }
    }

    @SubscribeEvent
    @SuppressWarnings("unused")
    public void onPlayerClone(PlayerEvent.Clone event) {
        PlayerEntity oldPlayer = event.getOriginal();
        PlayerEntity newPlayer = event.getPlayer();
        if(oldPlayer != null && newPlayer != null && !oldPlayer.getEntityWorld().isRemote) {
            List<ICapabilityImplementation<ICapabilityProvider, ? extends ISerializable>> list = this.map.get(oldPlayer.getClass());
            if(list == null) {
                return;
            }
            list.stream()
                    .filter(impl -> PlayerEntity.class.isAssignableFrom(impl.getCarrierClass()))
                    .forEach(impl -> {
                        oldPlayer.getCapability(impl.getCapability(), null).ifPresent(oldProps -> {
                            newPlayer.getCapability(impl.getCapability(), null).ifPresent(newProps -> {
                                newProps.readFromNBT(oldProps.writeToNBT());
                            });
                        });
                    });
        }
    }
}

@@ -45,16 +45,55 @@
 ility;%0A%0A
+import com.google.common.collect.Sets;%0A
 import c
@@ -789,53 +789,10 @@
 nal 
-Map%3CClass%3C? extends ICapabilityProvider%3E, Lis
+Se
 t%3CIC
@@ -865,13 +865,34 @@
 le%3E%3E
-%3E map
+ capabilityImplementations
 ;%0A%0A 
@@ -941,27 +941,61 @@
 his.
-map = new HashMap%3C%3E
+capabilityImplementations = Sets.newConcurrentHashSet
 ();%0A
@@ -1157,185 +1157,72 @@
 tion
- implementation) %7B%0A        if(!this.map.containsKey(implementation.getCarrierClass())) %7B%0A            this.map.put(implementation.getCarrierClass(), new ArrayList%3C%3E());%0A        %7D
+%3CICapabilityProvider, ? extends ISerializable%3E implementation) %7B
 %0A   
@@ -1235,17 +1235,19 @@
 his.
-map.get(i
+capabilityI
 mple
@@ -1255,35 +1255,17 @@
 entation
-.getCarrierClass())
+s
 .add(imp
@@ -1380,16 +1380,18 @@
 yStorage
+%3C%3E
 (), () -
@@ -2083,138 +2083,33 @@
 his.
-map.entrySet().stream().filter(entry -%3E entry.getKey().isAssignableFrom(clazz)).forEach(%0A                entry -%3E entry.getValue()
+capabilityImplementations
 .str
@@ -2126,32 +2126,24 @@
             
-        
 .filter(impl
@@ -2194,24 +2194,16 @@
 clazz))%0A
-        
         
@@ -2275,32 +2275,24 @@
             
-        
 .forEach(imp
@@ -2337,17 +2337,16 @@
 arrier))
-)
 ;%0A    %7D%0A
@@ -2917,15 +2917,14 @@
     
-List%3CIC
+this.c
 apab
@@ -2946,184 +2946,9 @@
 tion
-%3CICapabilityProvider, ? extends ISerializable%3E%3E list = this.map.get(oldPlayer.getClass());%0A            if(list == null) %7B%0A                return;%0A            %7D%0A            list
+s
 .str

Make registering of Capabilities threadsafe