package p.studio.utilities.events;

import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;

public final class TypedEventBus implements EventBusPublisher<SimpleEvent> {
    private final AtomicLong nextSubscriptionId = new AtomicLong(1L);
    private final ConcurrentMap<Class<?>, ConcurrentMap<Long, Consumer<Object>>> subscriptionsByType =
            new ConcurrentHashMap<>();

    public <T extends SimpleEvent> EventSubscription subscribe(Class<T> eventType, Consumer<? super T> consumer) {
        Objects.requireNonNull(eventType, "eventType");
        Objects.requireNonNull(consumer, "consumer");

        final long subscriptionId = nextSubscriptionId.getAndIncrement();
        final Consumer<Object> erasedConsumer = event -> consumer.accept(eventType.cast(event));
        final ConcurrentMap<Long, Consumer<Object>> subscriptions =
                subscriptionsByType.computeIfAbsent(eventType, ignored -> new ConcurrentHashMap<>());
        subscriptions.put(subscriptionId, erasedConsumer);

        final AtomicBoolean active = new AtomicBoolean(true);
        return () -> {
            if (!active.compareAndSet(true, false)) {
                return;
            }

            final ConcurrentMap<Long, Consumer<Object>> currentSubscriptions = subscriptionsByType.get(eventType);
            if (currentSubscriptions == null) {
                return;
            }

            currentSubscriptions.remove(subscriptionId);
            if (currentSubscriptions.isEmpty()) {
                subscriptionsByType.remove(eventType, currentSubscriptions);
            }
        };
    }

    @Override
    public <E extends SimpleEvent> void publish(E event) {
        Objects.requireNonNull(event, "event");

        final ConcurrentMap<Long, Consumer<Object>> subscriptions = subscriptionsByType.get(event.getClass());
        if (subscriptions == null || subscriptions.isEmpty()) {
            return;
        }

        for (Consumer<Object> consumer : subscriptions.values()) {
            consumer.accept(event);
        }
    }
}
