package com.nimbusds.infinispan.persistence.ldap;


import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;

import com.nimbusds.infinispan.persistence.common.InfinispanEntry;
import com.nimbusds.infinispan.persistence.common.InfinispanStore;
import com.nimbusds.infinispan.persistence.ldap.backend.LDAPConnector;
import com.unboundid.ldap.sdk.DN;
import com.unboundid.ldap.sdk.ReadOnlyEntry;
import net.jcip.annotations.ThreadSafe;
import org.infinispan.commons.configuration.ConfiguredBy;
import org.infinispan.filter.KeyFilter;
import org.infinispan.marshall.core.MarshalledEntry;
import org.infinispan.marshall.core.MarshalledEntryFactory;
import org.infinispan.persistence.TaskContextImpl;
import org.infinispan.persistence.spi.InitializationContext;
import org.infinispan.persistence.spi.PersistenceException;
import org.kohsuke.MetaInfServices;


/**
 * LDAP store for Infinispan 8.2+ caches and maps.
 */
@ThreadSafe
@MetaInfServices
@ConfiguredBy(LDAPStoreConfiguration.class)
public class LDAPStore<K,V> extends InfinispanStore<K,V> {


        /**
         * The LDAP store configuration.
         */
        private LDAPStoreConfiguration config;


        /**
         * The LDAP backend connector.
         */
        private LDAPConnector ldapConnector;


        /**
         * The LDAP entry transformer (to / from Infinispan entries).
         */
        private LDAPEntryTransformer<K,V> ldapEntryTransformer;


        /**
         * The marshalled Infinispan entry factory.
         */
        private MarshalledEntryFactory<K, V> marshalledEntryFactory;


        /**
         * Purges expired entries found in the LDAP store, as indicated by
         * their persisted metadata (optional, may be ignored / not stored).
         */
        private ExpiredEntryReaper<K,V> reaper;


        /**
         * Loads an LDAP entry transformer with the specified class name.
         *
         * @param className The class name. Must not be {@code null}.
         *
         * @return The LDAP entry transformer.
         */
        @SuppressWarnings( "unchecked" )
        private LDAPEntryTransformer<K,V> loadEntryTransformerClass(final String className) {

                try {
                        Class<LDAPEntryTransformer<K,V>> clazz = (Class<LDAPEntryTransformer<K,V>>)Class.forName(className);
                        return clazz.newInstance();
                } catch (Exception e) {
                        throw new PersistenceException("Couldn't load LDAP entry transformer class: " + e.getMessage(), e);
                }
        }


        @Override
        @SuppressWarnings("unchecked")
        public void init(final InitializationContext ctx) {

                // This method will be invoked by the PersistenceManager during initialization. The InitializationContext
                // contains:
                // - this CacheLoader's configuration
                // - the cache to which this loader is applied. Your loader might want to use the cache's name to construct
                //   cache-specific identifiers
                // - the StreamingMarshaller that needs to be used to marshall/unmarshall the entries
                // - a TimeService which the loader can use to determine expired entries
                // - a ByteBufferFactory which needs to be used to construct ByteBuffers
                // - a MarshalledEntryFactory which needs to be used to construct entries from the data retrieved by the loader

                super.init(ctx);
                
                this.config = ctx.getConfiguration();

                Loggers.MAIN_LOG.info("[IL0201] LDAP store configuration properties for cache {}:", getCacheName());
                config.log();

                Loggers.MAIN_LOG.debug("[IL0202] Loading LDAP entry transformer class {} for cache {}...",
                        config.ldapDirectory.entryTransformer,
                        getCacheName());

                ldapEntryTransformer = loadEntryTransformerClass(config.ldapDirectory.entryTransformer);

                marshalledEntryFactory = (MarshalledEntryFactory<K, V>)ctx.getMarshalledEntryFactory();

                Loggers.MAIN_LOG.info("[IL0203] Initialized LDAP external store for cache {}", getCacheName());
        }


        @Override
        public void start() {

                // This method will be invoked by the PersistenceManager to start the CacheLoader. At this stage configuration
                // is complete and the loader can perform operations such as opening a connection to the external storage,
                // initialize internal data structures, etc.

                if (ldapConnector != null) {
                        throw new IllegalStateException("LDAP store connector already started");
                }

                ldapConnector = new LDAPConnector(
                        config,
                        getCacheName(),
                        ldapEntryTransformer.getModifiableAttributes(),
                        ldapEntryTransformer.includesAttributesWithOptions());

                Loggers.MAIN_LOG.info("[IL0204] Started LDAP external store connector for cache {}", getCacheName());

                reaper = new ExpiredEntryReaper<>(ldapConnector, ldapEntryTransformer);
        }


        @Override
        public void stop() {

                if (ldapConnector != null) {
                        ldapConnector.shutdown();
                        Loggers.MAIN_LOG.info("[IL0205] Stopped LDAP external store connector for cache {}", getCacheName());
                }
                
                super.stop();
        }


        @SuppressWarnings("unchecked")
        private K resolveKey(final Object key) {

                if (key instanceof byte[]) {
                        throw new PersistenceException("Cannot resolve " + getCacheName() + " cache key from byte[], enable compatibility mode");
                }

                return (K)key;
        }


        @Override
        public boolean contains(final Object key) {

                // This method will be invoked by the PersistenceManager to determine if the loader contains the specified key.
                // The implementation should be as fast as possible, e.g. it should strive to transfer the least amount of data possible
                // from the external storage to perform the check. Also, if possible, make sure the field is indexed on the external storage
                // so that its existence can be determined as quickly as possible.
                //
                // Note that keys will be in the cache's native format, which means that if the cache is being used by a remoting protocol
                // such as HotRod or REST and compatibility mode has not been enabled, then they will be encoded in a byte[].

                Loggers.LDAP_LOG.trace("[IL0250] LDAP store: Checking {} cache key {}", getCacheName(), key);

                DN dn = new DN(ldapEntryTransformer.resolveRDN(resolveKey(key)), config.ldapDirectory.baseDN);

                return ldapConnector.entryExists(dn);
        }


        @Override
        @SuppressWarnings("unchecked")
        public MarshalledEntry<K,V> load(final Object key) {

                // Fetches an entry from the storage using the specified key. The CacheLoader should retrieve from the external storage all
                // of the data that is needed to reconstruct the entry in memory, i.e. the value and optionally the metadata. This method
                // needs to return a MarshalledEntry which can be constructed as follows:
                //
                // ctx.getMarshalledEntryFactory().new MarshalledEntry(key, value, metadata);
                //
                // If the entry does not exist or has expired, this method should return null.
                // If an error occurs while retrieving data from the external storage, this method should throw a PersistenceException
                //
                // Note that keys and values will be in the cache's native format, which means that if the cache is being used by a remoting protocol
                // such as HotRod or REST and compatibility mode has not been enabled, then they will be encoded in a byte[].
                // If the loader needs to have knowledge of the key/value data beyond their binary representation, then it needs access to the key's and value's
                // classes and the marshaller used to encode them.

                Loggers.LDAP_LOG.trace("[IL0251] LDAP store: Loading {} cache entry with key {}", getCacheName(), key);

                DN dn = new DN(ldapEntryTransformer.resolveRDN(resolveKey(key)), config.ldapDirectory.baseDN);

                Loggers.LDAP_LOG.trace("[IL0257] LDAP store: Resolved DN {}", dn);

                ReadOnlyEntry ldapEntry = ldapConnector.retrieveEntry(dn);

                if (ldapEntry == null) {
                        // Not found
                        Loggers.LDAP_LOG.trace("[IL0258] LDAP store: Entry not found");
                        return null;
                }


                if (Loggers.LDAP_LOG.isTraceEnabled()) {
                        Loggers.LDAP_LOG.trace("[IL0259] LDAP store: Retrieved entry: {}", ldapEntry.toLDIFString());
                }

                // Transform LDAP entry to Infinispan entry
                InfinispanEntry<K,V> infinispanEntry = ldapEntryTransformer.toInfinispanEntry(new LDAPEntry(ldapEntry));

                return marshalledEntryFactory.newMarshalledEntry(infinispanEntry.getKey(), infinispanEntry.getValue(), infinispanEntry.getMetadata());
        }


        @Override
        public boolean delete(final Object key) {

                // The CacheWriter should remove from the external storage the entry identified by the specified key.
                // Note that keys will be in the cache's native format, which means that if the cache is being used by a remoting protocol
                // such as HotRod or REST and compatibility mode has not been enabled, then they will be encoded in a byte[].

                Loggers.LDAP_LOG.trace("[IL0252] LDAP store: Deleting {} entry with key {}", getCacheName(), key);

                DN dn = new DN(ldapEntryTransformer.resolveRDN(resolveKey(key)), config.ldapDirectory.baseDN);

                return ldapConnector.deleteEntry(dn);
        }


        @Override
        public void write(final MarshalledEntry<? extends K, ? extends V> marshalledEntry) {

                // The CacheWriter should write the specified entry to the external storage.
                //
                // The PersistenceManager uses MarshalledEntry as the default format so that CacheWriters can efficiently store data coming
                // from a remote node, thus avoiding any additional transformation steps.
                //
                // Note that keys and values will be in the cache's native format, which means that if the cache is being used by a remoting protocol
                // such as HotRod or REST and compatibility mode has not been enabled, then they will be encoded in a byte[].

                Loggers.LDAP_LOG.trace("[IL0253] LDAP store: Writing {} entry {}", getCacheName(), marshalledEntry);

                LDAPEntry ldapEntry = ldapEntryTransformer.toLDAPEntry(
                        config.ldapDirectory.baseDN,
                        new InfinispanEntry<>(
                                marshalledEntry.getKey(),
                                marshalledEntry.getValue(),
                                marshalledEntry.getMetadata()));

                // Resolve the LDAP write strategy
                LDAPWriteStrategy writeStrategy = ldapEntry.getWriteStrategy();

                if (writeStrategy != null) {
                        Loggers.LDAP_LOG.trace("[IL0263] LDAP store: Entry transformer suggested {} write strategy", writeStrategy);
                } else {
                        writeStrategy = LDAPWriteStrategy.getDefault();
                        Loggers.LDAP_LOG.trace("[IL0264] LDAP store: Defaulted to {} write strategy", writeStrategy);
                }

                // Entry metadata created timestamp unreliable, cannot be used
                // as hint whether LDAP ADD or LDAP MODIFY should be attempted
                // first
                // InternalMetadataImpl{actual=EmbeddedExpirableMetadata{lifespan=-1, maxIdle=-1, version=null}, created=-1, lastUsed=-1}


                switch (writeStrategy) {

                        case TRY_LDAP_ADD_FIRST:

                                if (ldapConnector.addEntry(ldapEntry.getEntry())) {
                                        Loggers.LDAP_LOG.trace("[IL0256] LDAP store: Added new {} entry with DN {}", getCacheName(), ldapEntry.getEntry().getDN());
                                        return; // success
                                }

                                // Entry already exists, attempt LDAP modify
                                if (ldapConnector.replaceEntry(ldapEntry.getEntry())) {
                                        Loggers.LDAP_LOG.trace("[IL0257] LDAP store: Replaced {} entry with DN {}", getCacheName(), ldapEntry.getEntry().getDN());
                                        return; // success
                                }

                                // Try to recover from concurrent LDAP delete
                                // (entry deleted between first failed add and second modify attempt)
                                if (! ldapConnector.addEntry(ldapEntry.getEntry())) {
                                        // This should be highly unlikely
                                        throw new PersistenceException("Failed recovery from concurrent LDAP delete (" + getCacheName() + " cache): " + ldapEntry.getEntry().getDN());
                                }

                                break;

                        case TRY_LDAP_MODIFY_FIRST:

                                if (ldapConnector.replaceEntry(ldapEntry.getEntry())) {
                                        Loggers.LDAP_LOG.trace("[IL0265] LDAP store: Replaced {} entry with DN {}", getCacheName(), ldapEntry.getEntry().getDN());
                                        return; // success
                                }

                                // Entry doesn't exist, try LDAP add
                                if (ldapConnector.addEntry(ldapEntry.getEntry())) {
                                        Loggers.LDAP_LOG.trace("[IL0266] LDAP store: Added new {} entry with DN {}", getCacheName(), ldapEntry.getEntry().getDN());
                                        return; // success
                                }

                                // Try to recover from concurrent LDAP add
                                // (entry added between first failed replace and second add attempt)
                                if (! ldapConnector.replaceEntry(ldapEntry.getEntry())) {
                                        // This should be highly unlikely
                                        throw new PersistenceException("Failed recovery from concurrent LDAP add (" + getCacheName() + " cache): " + ldapEntry.getEntry().getDN());
                                }

                                break;

                        default:
                                throw new PersistenceException("Unexpected LDAP write strategy: " + writeStrategy);
                }
        }


        @Override
        public void process(final KeyFilter<? super K> keyFilter,
                            final CacheLoaderTask<K, V> cacheLoaderTask,
                            final Executor executor,
                            final boolean fetchValue,
                            final boolean fetchMetadata) {

                Loggers.LDAP_LOG.trace("[IL0262] LDAP store: Processing key filter for {} cache: fetchValue={} fetchMetadata=", getCacheName(), fetchValue, fetchMetadata);

                final TaskContext taskContext = new TaskContextImpl();

                // TODO consider multi-threaded LDAP retrieval?
                executor.execute(() -> ldapConnector.retrieveEntries(ldapEntry -> {

                        // Retrieves entire entry, fetchValue / fetchMetadata params are ignored TODO consider

                        if (taskContext.isStopped()) {
                                // TODO Consider pushing task context to LDAP connector routine
                                return;
                        }

                        InfinispanEntry<K,V> infinispanEntry = ldapEntryTransformer.toInfinispanEntry(new LDAPEntry(ldapEntry));

                        if (keyFilter.accept(infinispanEntry.getKey())) {

                                MarshalledEntry<K,V> marshalledEntry = marshalledEntryFactory.newMarshalledEntry(
                                        infinispanEntry.getKey(),
                                        infinispanEntry.getValue(),
                                        infinispanEntry.getMetadata());

                                try {
                                        cacheLoaderTask.processEntry(marshalledEntry, taskContext);

                                } catch (InterruptedException e) {
                                        throw new PersistenceException(e.getMessage(), e);
                                }
                        }
                }));
        }


        @Override
        public int size() {

                // Infinispan code analysis on 8.2 shows that this method is never called in practise, and
                // is not wired to the data / cache container API

                Loggers.LDAP_LOG.trace("[IL0258] LDAP store: Counting {} entries", getCacheName());

                final int count = ldapConnector.countEntries();

                Loggers.LDAP_LOG.trace("[IL0259] LDAP store: Counted {} {} entries", count, getCacheName());

                return count;
        }


        @Override
        public void clear() {

                Loggers.LDAP_LOG.trace("[IL0260] LDAP store: Clearing {} entries", getCacheName());

                int numDeleted = ldapConnector.deleteEntries();

                Loggers.LDAP_LOG.debug("[IL0254] LDAP store: Cleared {} {} entries", numDeleted, getCacheName());
        }


        @Override
        public void purge(final Executor executor, final PurgeListener<? super K> purgeListener) {

                Loggers.LDAP_LOG.trace("[IL0261] LDAP store: Purging {} entries", getCacheName());

                final AtomicInteger numPurged = new AtomicInteger();

                executor.execute(() -> numPurged.set(reaper.purge(purgeListener)));

                Loggers.LDAP_LOG.debug("[IL0255] LDAP store: Purged {} expired {} entries", numPurged.get(), getCacheName());
        }
}