package mezz.jei.events;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

import net.minecraftforge.fml.event.lifecycle.ModLifecycleEvent;
import net.minecraftforge.common.MinecraftForge;
import net.minecraftforge.eventbus.api.Event;
import net.minecraftforge.eventbus.api.EventPriority;
import net.minecraftforge.eventbus.api.IEventBus;

public class EventBusHelper {
    private static class Subscription {
        private final IEventBus eventBus;
        private final Consumer<? extends Event> listener;

        private <T extends Event> Subscription(IEventBus eventBus, Consumer<T> listener) {
            this.eventBus = eventBus;
            this.listener = listener;
        }
    }

    private static final Map<Object, List<Subscription>> subscriptions = new HashMap<>();

    private static IEventBus getInstance() {
        return MinecraftForge.EVENT_BUS;
    }

    /**
     * See {@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)}
     */
    public static <T extends Event> void addListener(Object owner, Class<T> eventType, Consumer<T> listener) {
        addListener(owner, getInstance(), eventType, listener);
    }

    /**
     * @param owner The owner of this listener. When the owner is unregistered, all listeners with the same owner will
     *              also be unregistered.
     */
    public static <T extends Event> void addListener(Object owner, IEventBus eventBus, Class<T> eventType, Consumer<T> listener) {
        subscriptions.computeIfAbsent(owner, e -> new ArrayList<>()).add(new Subscription(eventBus, listener));
        eventBus.addListener(EventPriority.NORMAL, false, eventType, listener);
    }

    /**
     * See {@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)}
     */
    public static <T extends ModLifecycleEvent> void addLifecycleListener(Object owner, IEventBus eventBus, Class<T> eventType, Consumer<T> listener) {
        addListener(owner, eventBus, eventType, listener);
    }

    /**
     * See {@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)}
     */
    public static <T extends Event> void removeListener(Object owner, Consumer<T> listener) {
        removeListener(owner, getInstance(), listener);
    }

    /**
     * See {@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)}
     */
    public static <T extends Event> void removeListener(Object owner, IEventBus eventBus, Consumer<T> listener) {
        subscriptions.get(owner).remove(listener);
        eventBus.unregister(listener);
    }

    public static void register(Object owner) {
        IEventBus eventBus = getInstance();
        subscriptions.putIfAbsent(owner, new ArrayList<>());
        eventBus.register(owner);
    }

    /**
     * Unregisters an object from the event bus.
     * All listeners owned by this object will also be unregistered.
     * See {@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)}
     */
    public static void unregister(Object owner) {
        IEventBus eventBus = getInstance();
        for (Subscription sub : subscriptions.get(owner)) {
            removeListener(owner, sub.eventBus, sub.listener);
        }
        eventBus.unregister(owner);
        subscriptions.remove(owner);
    }

    public static void post(Event event) {
        IEventBus eventBus = getInstance();
        eventBus.post(event);
    }
}

@@ -114,24 +114,50 @@
 a.util.Map;%0A
+import java.util.Objects;%0A
 import java.
@@ -789,24 +789,114 @@
 r;%0A        %7D
+%0A%0A        private void unregister() %7B%0A            eventBus.unregister(listener);%0A        %7D
 %0A    %7D%0A%0A    
@@ -2382,262 +2382,171 @@
     
-removeListener(owner, getInstance(), listener);%0A    %7D%0A
+List%3CSubscription%3E subs = subscriptions.get(owner);%0A        if (subs != null) %7B
 %0A    
-/**%0A
      
-* See %7B@link EventBusHelper#addListener(Object, IEventBus, Class, Consumer)%7D
+   subs.removeIf(sub -%3E %7B
 %0A     
-*/%0A
+  
     
-public static %3CT extends Event%3E void removeListener(Object owner, IEventBus eventBus, Consumer%3CT%3E
+     if (Objects.equals(sub.listener,
  lis
@@ -2543,32 +2543,33 @@
 tener, listener)
+)
  %7B%0A        subsc
@@ -2567,48 +2567,35 @@
     
+     
-subscriptions.get(owner).remove(l
+       sub.unreg
 iste
-ne
 r
+(
 );%0A 
@@ -2605,38 +2605,106 @@
     
-eventBus.unregister(listener);
+            return true;%0A                %7D%0A                return false;%0A            %7D);%0A        %7D
 %0A   
@@ -3223,13 +3223,13 @@
     
-for (
+List%3C
 Subs
@@ -3240,14 +3240,16 @@
 tion
+%3E
  sub
- :
+s =
  sub
@@ -3263,84 +3263,94 @@
 ons.
-get
+remove
 (owner)
-) %7B%0A            removeListener(owner, sub.eventBus, sub.listener);
+;%0A        if (subs != null) %7B%0A            for (Subscription sub : subs) %7B
 %0A   
@@ -3354,26 +3354,24 @@
         
-%7D%0A
         
 eventBus
@@ -3358,32 +3358,27 @@
             
-eventBus
+sub
 .unregister(
@@ -3373,29 +3373,24 @@
 .unregister(
-owner
 );%0A        s
@@ -3388,36 +3388,59 @@
         
-subscriptions.remove
+    %7D%0A        %7D%0A        eventBus.unregister
 (owner);

Fix erroneous event listener removal (#2157)