package com.nimbusds.infinispan.persistence.ldap.backend;


import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

import com.codahale.metrics.Timer;
import com.nimbusds.common.appendable.Appendable;
import com.nimbusds.common.ldap.LDAPConnectionPoolFactory;
import com.nimbusds.common.ldap.LDAPConnectionPoolMetrics;
import com.nimbusds.common.ldap.LDAPHealthCheck;
import com.nimbusds.common.monitor.MonitorRegistries;
import com.nimbusds.infinispan.persistence.ldap.Loggers;
import com.nimbusds.infinispan.persistence.ldap.LDAPStoreConfiguration;
import com.unboundid.asn1.ASN1OctetString;
import com.unboundid.ldap.sdk.*;
import com.unboundid.ldap.sdk.controls.SimplePagedResultsControl;
import net.jcip.annotations.ThreadSafe;
import org.infinispan.persistence.spi.PersistenceException;


/**
 * LDAP connector to the backend directory. Provides retrieval, addition /
 * modify and deletion operations.
 */
@ThreadSafe
public class LDAPConnector {


        /**
         * Filter that matches any LDAP entry.
         */
        public static final Filter MATCH_ANY_FILTER = Filter.createPresenceFilter("objectClass");
        

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


        /**
         * The LDAP connection pool.
         */
        private final LDAPConnectionPool ldapConnPool;


        /**
         * The LDAP request factory.
         */
        private final LDAPModifyRequestFactory ldapModifyRequestFactory;


        /**
         * Indicates support for directory attributes that may include options
         * (such as language tags).
         */
        private final boolean supportAttributeOptions;


        /**
         * The LDAP operation timers.
         */
        private final LDAPTimers ldapTimers;
        
        
        /**
         * The name of the associated Infinispan cache / map.
         */
        private final String cacheName;


        /**
         * Creates a new LDAP connector.
         *
         * @param config                  The LDAP configuration. Must not be
         *                                {@code null}.
         * @param cacheName               The name of the Infinispan cache
         *                                associated with this LDAP connector.
         *                                Must not be {@code null}.
         * @param attributes              The names of the supported directory
         *                                attributes. Must not be empty or
         *                                {@code null}.
         * @param supportAttributeOptions {@code true} if any of the supported
         *                                directory attributes may include
         *                                options (such as language tags), else
         *                                {@code false}.
         */
        public LDAPConnector(final LDAPStoreConfiguration config,
                             final String cacheName,
                             final Set<String> attributes,
                             final boolean supportAttributeOptions) {

                this.config = config;
                this.cacheName = cacheName;
                this.supportAttributeOptions = supportAttributeOptions;

                LDAPConnectionPoolFactory factory = new LDAPConnectionPoolFactory(
                        config.ldapServer,
                        config.customTrustStore,
                        config.customKeyStore,
                        config.ldapUser);

                try {
                        ldapConnPool = factory.createLDAPConnectionPool();
                } catch (Exception e) {
                        throw new PersistenceException("LDAP connection pool creation for " + cacheName + " cache failed: " + e.getMessage(), e);
                }

                ldapConnPool.setConnectionPoolName(cacheName);

                checkBaseDN();

                // Setup factory for composing LDAP requests
                ldapModifyRequestFactory = new LDAPModifyRequestFactory(attributes);

                // Set up DropWizard Metrics

                /// LDAP operation timers
                ldapTimers = new LDAPTimers(cacheName + ".");

                // Set up LDAP connection pool metrics and health check
                final String prefix = cacheName + ".ldapStore";
                MonitorRegistries.register(new LDAPConnectionPoolMetrics(ldapConnPool, prefix));
                MonitorRegistries.register(prefix, new LDAPHealthCheck(ldapConnPool, config.ldapDirectory.baseDN, Loggers.LDAP_LOG));

                Loggers.MAIN_LOG.info("[IL0100] Created new LDAP store connector for " + cacheName + " cache");
        }


        /**
         * Checks if the configured directory base DN for the entries exists.
         * If not a WARN message is logged.
         */
        private void checkBaseDN() {

                try {
                        if (ldapConnPool.getEntry(config.ldapDirectory.baseDN.toString()) == null) {

                                Loggers.MAIN_LOG.warn("[IL0101] The configured LDAP store base DN for {} cache doesn't exist: {}",
                                        cacheName,
                                        config.ldapDirectory.baseDN);
                        }

                } catch (LDAPException e) {

                        Loggers.MAIN_LOG.warn("[IL0102] Couldn't verify the LDAP store base DN for {} cache: {}",
                                cacheName,
                                e.getMessage());
                }
        }


        /**
         * Returns the underlying LDAP connection pool.
         *
         * @return The LDAP connection pool.
         */
        public LDAPConnectionPool getPool() {

                return ldapConnPool;
        }


        /**
         * Retrieves the specified entry from the LDAP directory.
         *
         * @param dn The Distinguished Name (DN) of the entry to retrieve. Must
         *           not be {@code null}.
         *
         * @return The matching directory entry, {@code null} if not found (or
         *         insufficient privileges).
         */
        public ReadOnlyEntry retrieveEntry(final DN dn) {

                Timer.Context timerCtx = ldapTimers.getTimer.time();
                try {
                        return ldapConnPool.getEntry(dn.toString(), SearchRequest.ALL_USER_ATTRIBUTES);
                } catch (LDAPException e) {
                        throw new PersistenceException("LDAP get of " + dn + " failed: " + e.getResultString(), e);
                } finally {
                        timerCtx.stop();
                }
        }


        /**
         * Retrieves all entries from the LDAP directory under the base DN.
         *
         * @param appendable Collects the matching directory entries. Must not
         *                   be {@code null}.
         */
        public void retrieveEntries(final Appendable<ReadOnlyEntry> appendable) {

                SearchRequest searchRequest = new SearchRequest(
                        config.ldapDirectory.baseDN.toString(),
                        SearchScope.ONE,
                        MATCH_ANY_FILTER,
                        SearchRequest.ALL_USER_ATTRIBUTES);

                doSearch(searchRequest, appendable);
        }


        /**
         * Checks that the specified entry exists in the LDAP directory.
         *
         * @param dn The Distinguished Name (DN) of the entry to check. Must
         *           not be {@code null}.
         *
         * @return {@code true} if the entry exists, {@code false} if not
         *         found (or insufficient privileges).
         */
        public boolean entryExists(final DN dn) {

                Timer.Context timerCtx = ldapTimers.getTimer.time();
                try {
                        return ldapConnPool.getEntry(dn.toString(), SearchRequest.NO_ATTRIBUTES) != null;
                } catch (LDAPException e) {
                        throw new PersistenceException("LDAP get of " + dn + " failed: " + e.getResultString(), e);
                } finally {
                        timerCtx.stop();
                }
        }


        /**
         * Adds the specified entry to the LDAP directory.
         *
         * @param entry The entry to add. Must not be {@code null}.
         *
         * @return {@code true} if the entry was added, {@code false} if there
         *         was a stored entry with that Distinguished Name (DN).
         */
        public boolean addEntry(final ReadOnlyEntry entry) {

                LDAPResult ldapResult;

                Timer.Context timerCtx = ldapTimers.addTimer.time();

                try {
                        ldapResult = ldapConnPool.add(entry);

                } catch (LDAPException e) {

                        if (e.getResultCode().equals(ResultCode.ENTRY_ALREADY_EXISTS)) {
                                return false;
                        }

                        throw new PersistenceException("LDAP add for " + entry.getDN()  + " failed: " + e.getResultString(), e);
                } finally {
                        timerCtx.stop();
                }

                ResultCode resultCode = ldapResult.getResultCode();

                if (resultCode.equals(ResultCode.SUCCESS)) {
                        return true;
                }

                if (resultCode.equals(ResultCode.ENTRY_ALREADY_EXISTS)) {
                        return false;
                }

                // Other result code indicates exception
                throw new PersistenceException("LDAP add for " + entry.getDN()  + " failed: " + resultCode.getName());
        }


        /**
         * Replaces the specified entry in the LDAP directory.
         *
         * @param entry The entry to replace. Must not be {@code null}.
         *
         * @return {@code true} if the entry was found, {@code false} if
         *         there was no stored entry with that Distinguished Name (DN).
         */
        public boolean replaceEntry(final ReadOnlyEntry entry) {

                final ModifyRequest modifyRequest;

                if (supportAttributeOptions) {

                        DN dn;

                        try {
                                dn = new DN(entry.getDN());
                        } catch (LDAPException e) {
                                throw new PersistenceException(e.getMessage(), e);
                        }
                        // Fetch previous entry, otherwise attributes with
                        // options cannot be reliably updates
                        ReadOnlyEntry existingEntry = retrieveEntry(dn);

                        if (existingEntry == null) {
                                return false; // No such entry
                        }

                        // Create mod based on diff
                        modifyRequest = ldapModifyRequestFactory.composeModifyRequest(entry, existingEntry);

                        if (modifyRequest == null) {
                                // No diff detected, indicate entry was found though
                                return true;
                        }

                } else {

                        // Create fully-specced mod
                        modifyRequest = ldapModifyRequestFactory.composeModifyRequest(entry);
                }

                LDAPResult ldapResult;
                Timer.Context timerCtx = ldapTimers.modifyTimer.time();

                try {
                        ldapResult = ldapConnPool.modify(modifyRequest);

                } catch (LDAPException e) {

                        if (e.getResultCode().equals(ResultCode.NO_SUCH_OBJECT)) {
                                return false;
                        }

                        throw new PersistenceException("LDAP modify for " + modifyRequest.getDN() + " failed: " + e.getResultString(), e);
                } finally {
                        timerCtx.stop();
                }

                ResultCode resultCode = ldapResult.getResultCode();

                if (resultCode.equals(ResultCode.SUCCESS)) {
                        return true;
                }

                if (resultCode.equals(ResultCode.NO_SUCH_OBJECT)) {
                        return false;
                }

                throw new PersistenceException("LDAP modify " + modifyRequest.getDN() + " failed: " + resultCode.getName());
        }


        /**
         * Deletes the specified entry from the LDAP directory.
         *
         * @param dn The Distinguished Name (DN) of the entry to delete. Must
         *           not be {@code null}.
         *
         * @return {@code true} if the matching directory entry was deleted,
         *         {@code false} if not found (or insufficient privileges).
         */
        public boolean deleteEntry(final DN dn) {

                DeleteRequest deleteRequest = new DeleteRequest(dn);

                LDAPResult result;
                Timer.Context timerCtx = ldapTimers.deleteTimer.time();

                try {
                        result = ldapConnPool.delete(deleteRequest);

                } catch (LDAPException e) {

                        ResultCode resultCode = e.getResultCode();

                        if (resultCode.equals(ResultCode.NO_SUCH_OBJECT))
                                return false;

                        throw new PersistenceException("LDAP delete of " + dn + " failed: " + e.getResultString(), e);
                } finally {
                        timerCtx.stop();
                }

                ResultCode resultCode = result.getResultCode();

                if (resultCode.equals(ResultCode.SUCCESS)) {
                        return true;
                }

                if (resultCode.equals(ResultCode.NO_SUCH_OBJECT)) {
                        return false;
                }

                throw new PersistenceException("LDAP delete of " + dn + " failed: " + resultCode.getName());
        }


        /**
         * Checks if the specified LDAP exception is caused by the LDAP server
         * being unavailable, disconnected or timing out.
         *
         * @param e The LDAP exception. Must not be {@code null}.
         *
         * @return {@code true} if the LDAP exception is caused by the LDAP
         *         server being unavailable, disconnected or timing out, else
         *         {@code false}.
         */
        protected static boolean indicatesConnectionException(final LDAPException e) {

                return indicatesConnectionException(e.getResultCode());
        }


        /**
         * Checks if the specified LDAP result code indicates the LDAP server
         * is unavailable, disconnected or timing out.
         *
         * @param code The LDAP result code. Must not be {@code null}.
         *
         * @return {@code true} if the LDAP result code indicates the LDAP
         *         server is unavailable, disconnected or timing out, else
         *         {@code false}.
         */
        protected static boolean indicatesConnectionException(final ResultCode code) {

                return code.equals(ResultCode.CONNECT_ERROR)
                        || code.equals(ResultCode.SERVER_DOWN)
                        || code.equals(ResultCode.TIMEOUT)
                        || code.equals(ResultCode.UNAVAILABLE);
        }


        /**
         * Parses the specified LDAP search result for a page cookie.
         *
         * @param sr The LDAP search result. Must not be {@code null}.
         *
         * @return The page cookie, {@code null} if not found or undefined.
         */
        private static ASN1OctetString parsePageCookie(final SearchResult sr) {

                Control control = sr.getResponseControl(SimplePagedResultsControl.PAGED_RESULTS_OID);

                if (control instanceof SimplePagedResultsControl) {
                        SimplePagedResultsControl spr = (SimplePagedResultsControl)control;
                        return spr.getCookie();
                } else {
                        return null;
                }
        }


        /**
         * Performs an LDAP search with the specified request.
         *
         * @param searchRequest The LDAP search request. Must not be
         *                      {@code null}.
         * @param appendable    Collects the matching directory entries. Must
         *                      not be {@code null}.
         */
        private void doSearch(final SearchRequest searchRequest, final Appendable<ReadOnlyEntry> appendable) {

                // Check out connection from pool
                final LDAPConnection connection;

                try {
                        connection = ldapConnPool.getConnection();
                } catch (LDAPException e) {
                        throw new PersistenceException(e.getMessage(), e);
                }

                ASN1OctetString cookie = null;

                do {
                        // Set paging cookie if returned from a previous iteration
                        searchRequest.replaceControl(new SimplePagedResultsControl(config.ldapDirectory.pageSize, cookie));

                        final SearchResult searchResult;
                        Timer.Context timerCtx = ldapTimers.searchTimer.time();

                        try {
                                searchResult = connection.search(searchRequest);

                        } catch (LDAPSearchException e) {

                                String msg = "[AS0109] LDAP search " + searchRequest.getFilter() + " failed: " + e.getMessage();

                                if (indicatesConnectionException(e)) {

                                        // Assume connection is unusable
                                        ldapConnPool.releaseDefunctConnection(connection);

                                } else {
                                        // Assume connection is still usable
                                        ldapConnPool.releaseConnection(connection);
                                }

                                throw new PersistenceException(msg, e);
                        } finally {
                                timerCtx.stop();
                        }

                        cookie = parsePageCookie(searchResult);

                        searchResult.getSearchEntries().forEach(appendable::append);

                } while (cookie != null && cookie.getValueLength() > 0);

                ldapConnPool.releaseConnection(connection);
        }


        /**
         * Counts the number of entries under the base DN.
         *
         * @return The entry count.
         */
        public int countEntries() {

                SearchRequest request = new SearchRequest(
                        config.ldapDirectory.baseDN.toString(),
                        SearchScope.ONE,
                        MATCH_ANY_FILTER,
                        SearchRequest.NO_ATTRIBUTES);

                final AtomicInteger count = new AtomicInteger();

                doSearch(request, entry -> count.incrementAndGet());

                return count.intValue();
        }


        /**
         * Deletes all entries under the base DN.
         *
         * @return The number of deleted entries, zero if none found.
         */
        public int deleteEntries() {

                SearchRequest request = new SearchRequest(
                        config.ldapDirectory.baseDN.toString(),
                        SearchScope.ONE,
                        MATCH_ANY_FILTER,
                        SearchRequest.NO_ATTRIBUTES);

                List<String> entryDNs = new LinkedList<>();

                doSearch(request, entry -> entryDNs.add(entry.getDN()));

                int count = 0;

                for(String dn: entryDNs) {

                        try {
                                if (deleteEntry(new DN(dn))) {
                                        ++count;
                                }
                        } catch (LDAPException e) {
                                throw new PersistenceException(e.getMessage(), e);
                        }
                }

                return count;
        }


        /**
         * Shuts down this LDAP connector by releasing any associated resources
         * (LDAP connection pool).
         */
        public void shutdown() {

                ldapConnPool.close();

                if (ldapConnPool.isClosed()) {
                        Loggers.MAIN_LOG.info("[IL0107] Shut down LDAP connector for {} cache", cacheName);
                } else {
                        Loggers.MAIN_LOG.error("[IL0108] Attempted to shut down LDAP connector for {} cache, detected unreleased LDAP connections", cacheName);
                }
        }
}