package graphql.execution.instrumentation.dataloader;

import graphql.Assert;
import graphql.Internal;
import graphql.PublicApi;
import graphql.execution.instrumentation.InstrumentationState;
import org.dataloader.DataLoader;
import org.dataloader.DataLoaderRegistry;
import org.slf4j.Logger;

import java.util.concurrent.atomic.AtomicReference;

/**
 * A base class that keeps track of whether aggressive batching can be used
 */
@PublicApi
public class DataLoaderDispatcherInstrumentationState implements InstrumentationState {

    @Internal
    public static final DataLoaderRegistry EMPTY_DATALOADER_REGISTRY = new DataLoaderRegistry() {
        @Override
        public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
            return Assert.assertShouldNeverHappen("You MUST set in your own DataLoaderRegistry to use data loader");
        }
    };

    private final FieldLevelTrackingApproach approach;
    private final AtomicReference<DataLoaderRegistry> dataLoaderRegistry;
    private final InstrumentationState state;
    private volatile boolean aggressivelyBatching = true;
    private volatile boolean hasNoDataLoaders;

    public DataLoaderDispatcherInstrumentationState(Logger log, DataLoaderRegistry dataLoaderRegistry) {
        this.dataLoaderRegistry = new AtomicReference<>(dataLoaderRegistry);
        this.approach = new FieldLevelTrackingApproach(log, this::getDataLoaderRegistry);
        this.state = approach.createState();
        hasNoDataLoaders = checkForNoDataLoader(dataLoaderRegistry);
    }

    private boolean checkForNoDataLoader(DataLoaderRegistry dataLoaderRegistry) {
        //
        // if they have never set a dataloader into the execution input then we can optimize
        // away the tracking code
        //
        return dataLoaderRegistry == EMPTY_DATALOADER_REGISTRY;
    }

    boolean isAggressivelyBatching() {
        return aggressivelyBatching;
    }

    void setAggressivelyBatching(boolean aggressivelyBatching) {
        this.aggressivelyBatching = aggressivelyBatching;
    }

    FieldLevelTrackingApproach getApproach() {
        return approach;
    }

    DataLoaderRegistry getDataLoaderRegistry() {
        return dataLoaderRegistry.get();
    }

    void setDataLoaderRegistry(DataLoaderRegistry newRegistry) {
        dataLoaderRegistry.set(newRegistry);
        hasNoDataLoaders = checkForNoDataLoader(newRegistry);
    }

    boolean hasNoDataLoaders() {
        return hasNoDataLoaders;
    }

    InstrumentationState getState() {
        return state;
    }
}

@@ -338,16 +338,52 @@
 ference;
+%0Aimport java.util.function.Function;
 %0A%0A/**%0A *
@@ -664,32 +664,152 @@
 derRegistry() %7B%0A
+%0A        private static final String ERROR_MESSAGE = %22You MUST set in your own DataLoaderRegistry to use data loader%22;%0A%0A
         @Overrid
@@ -950,72 +950,458 @@
 pen(
-%22You MUST set in your own DataLoaderRegistry to use data loader%22
+ERROR_MESSAGE);%0A        %7D%0A%0A        @Override%0A        public %3CK, V%3E DataLoader%3CK, V%3E computeIfAbsent(final String key,%0A                                                       final Function%3CString, DataLoader%3C?, ?%3E%3E mappingFunction) %7B%0A            return Assert.assertShouldNeverHappen(ERROR_MESSAGE);%0A        %7D%0A%0A        @Override%0A        public DataLoaderRegistry unregister(String key) %7B%0A            return Assert.assertShouldNeverHappen(ERROR_MESSAGE
 );%0A 

Throw an exception when try to register or unregister data loader when no DataLoaderRegistry provided.