package com.nimbusds.openid.connect.provider.spi.claims.ldap;


import java.io.InputStream;
import java.nio.charset.Charset;
import java.util.*;

import org.apache.commons.io.IOUtils;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import net.minidev.json.JSONObject;

import com.unboundid.ldap.sdk.Entry;
import com.unboundid.ldap.sdk.LDAPConnectionPool;
import com.unboundid.ldap.sdk.SearchResult;

import com.nimbusds.common.ldap.AttributeMapper;
import com.nimbusds.common.ldap.LDAPConnectionPoolFactory;

import com.nimbusds.langtag.LangTag;

import com.nimbusds.oauth2.sdk.id.Subject;
import com.nimbusds.oauth2.sdk.util.JSONObjectUtils;
import com.nimbusds.openid.connect.sdk.claims.UserInfo;

import com.nimbusds.openid.connect.provider.spi.InitContext;
import com.nimbusds.openid.connect.provider.spi.claims.ClaimUtils;
import com.nimbusds.openid.connect.provider.spi.claims.ClaimsSource;


/**
 * LDAP connector for retrieving OpenID Connect UserInfo claims.
 */
public class LDAPClaimsSource implements ClaimsSource {


        /**
         * The configuration file path.
         */
        public static final String CONFIG_FILE_PATH = "/WEB-INF/ldapClaimsSource.properties";


        /**
         * The LDAP claims map file path.
         */
        public static final String MAP_FILE_PATH = "/WEB-INF/ldapClaimsMap.json";


        /**
         * The LDAP connector configuration.
         */
        private Configuration config;


        /**
         * The supported UserInfo claims map. Allows for mapping of complex top
         * level claims, such as "address" to their sub-claims, while simple
         * claims map one-to-one.
         */
        private Map<String,List<String>> claimsMap;


        /**
         * The UserInfo attribute mapper.
         */
        private AttributeMapper attributeMapper;


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


        /**
         * The logger.
         */
        private final Logger log = LogManager.getLogger("MAIN");


        /**
         * Creates a new LDAP claims source. It must be {@link #init
         * initialised} before it can be used.
         */
        public LDAPClaimsSource() { }


        /**
         * Loads the configuration.
         *
         * @param initContext The initialisation context. Must not be
         *                    {@code null}.
         *
         * @return The configuration.
         *
         * @throws Exception If loading failed.
         */
        private static Configuration loadConfiguration(final InitContext initContext)
                throws Exception {

                InputStream inputStream = initContext.getResourceAsStream(CONFIG_FILE_PATH);

                if (inputStream == null) {
                        throw new Exception("Couldn't find LDAP claims source configuration file: " + CONFIG_FILE_PATH);
                }

                Properties props = new Properties();
                props.load(inputStream);

                return new Configuration(props);
        }


        /**
         * Loads the LDAP attribute map.
         *
         * @param initContext The initialisation context. Must not be
         *                    {@code null}.
         *
         * @return The LDAP attribute map.
         *
         * @throws Exception If loading failed.
         */
        private static Map<String,Object> loadLDAPAttributeMap(final InitContext initContext)
                throws Exception {

                InputStream inputStream = initContext.getResourceAsStream(MAP_FILE_PATH);

                if (inputStream == null) {
                        throw new Exception("Couldn't find LDAP claims map file: " + MAP_FILE_PATH);
                }

                try {
                        String jsonText = IOUtils.toString(inputStream, Charset.forName("UTF-8"));
                        return JSONObjectUtils.parseJSONObject(jsonText);

                } catch (Exception e) {

                        throw new Exception("Couldn't load LDAP claims map: " + e.getMessage(), e);
                }
        }


        /**
         * Composes the claims map from the specified LDAP attributes map.
         *
         * @param attrMap The LDAP attribute map. Must not be {@code null}.
         *
         * @return The claims map.
         */
        private static Map<String,List<String>> composeClaimsMap(final Map<String,Object> attrMap) {

                // Derive the supported UserInfo claims map
                Map<String,List<String>> claimsMap = new HashMap<>();

                for (String key: attrMap.keySet()) {

                        String[] parts = key.split("\\.", 2);

                        List<String> subClaims = claimsMap.get(parts[0]);

                        if (subClaims == null) {
                                subClaims = new LinkedList<>();
                        }

                        subClaims.add(key);

                        claimsMap.put(parts[0], subClaims);
                }

                return claimsMap;
        }


        @Override
        public void init(final InitContext initContext)
                throws Exception {

                log.info("Initializing LDAP claims source...");

                config = loadConfiguration(initContext);

                config.log();

                if (! config.enable) {
                        // stop initialisation
                        return;
                }

                // Load the raw LDAP attribute map
                Map<String,Object> ldapAttributeMap = loadLDAPAttributeMap(initContext);
                attributeMapper = new AttributeMapper(ldapAttributeMap);

                if (attributeMapper.getLDAPAttributeName("sub") == null) {
                        throw new Exception("Missing LDAP attribute mapping for \"sub\" claim");
                }

                // Compose the final claims map
                claimsMap = composeClaimsMap(ldapAttributeMap);

                LDAPConnectionPoolFactory factory = new LDAPConnectionPoolFactory(config.server,
                        config.customTrustStore,
                        config.customKeyStore,
                        config.directory.user);

                try {
                        ldapConnPool = factory.createLDAPConnectionPool();

                } catch (Exception e) {

                        // java.security.KeyStoreException
                        // java.security.GeneralSecurityException
                        // com.unboundid.ldap.sdk.LDAPException

                        throw new Exception("Couldn't create LDAP connection pool: " + e.getMessage(), e);
                }

                ldapConnPool.setConnectionPoolName("userinfo-store");
        }


        @Override
        public boolean isEnabled() {

                return config.enable;
        }


        @Override
        public Set<String> supportedClaims() {

                if (! config.enable) {
                        // Empty set
                        return Collections.unmodifiableSet(new HashSet<String>());
                }

                return Collections.unmodifiableSet(claimsMap.keySet());
        }


        /**
         * Resolves the individual requested claims from the specified
         * requested claims and preferred locales.
         *
         * @param claims        The requested claims. May contain optional
         *                      language tags. Must not be {@code null}.
         * @param claimsLocales The preferred locales, {@code null} if not
         *                      specified.
         *
         * @return The resolved individual requested claims.
         */
        protected List<String> resolveRequestedClaims(final Set<String> claims,
                                                      final List<LangTag> claimsLocales) {

                // Use set to ensure no duplicates get into the collection
                Set<String> individualClaims = new HashSet<>();

                for (String claim: claims) {

                        // Check if the claim is supported and if any sub-claims
                        // are associated with it (e.g. for UserInfo address)
                        List<String> claimsList = claimsMap.get(claim);

                        if (claimsList == null) {
                                // claim not supported
                                continue;
                        }

                        individualClaims.addAll(claimsList);
                }

                // Apply the preferred language tags if any
                individualClaims = ClaimUtils.applyLangTags(individualClaims, claimsLocales);

                return new ArrayList<>(individualClaims);
        }


        @Override
        public UserInfo getClaims(final Subject subject,
                                  final Set<String> claims,
                                  final List<LangTag> claimsLocales)
                throws Exception {

                if (! config.enable)
                        return null;

                // Compose search filter from the cofigured template
                String filter = config.directory.filter.apply(subject.getValue());

                // Resolve the individual requested claims
                List<String> claimsToRequest = resolveRequestedClaims(claims, claimsLocales);

                // Map OIDC claim names to LDAP attribute names
                List<String> ldapAttrs = attributeMapper.getLDAPAttributeNames(claimsToRequest);

                // Do LDAP search
                SearchResult searchResult;

                try {
                        searchResult = ldapConnPool.search(
                                config.directory.baseDN.toString(),
                                config.directory.scope,
                                filter,
                                ldapAttrs.toArray(new String[0]));

                } catch (Exception e) {

                        // LDAPException
                        throw new Exception("Couldn't get UserInfo for subject \"" + subject + "\": " + e.getMessage(), e);
                }


                // Get matches count
                final int entryCount = searchResult.getEntryCount();

                if (entryCount == 0) {
                        // Nothing found
                        return null;
                }

                if (entryCount > 1) {
                        // More than one entry found
                        throw new Exception("Found " + entryCount + " entries for subject \"" + subject + "\"");
                }


                // Process user entry
                Entry entry = searchResult.getSearchEntries().get(0);


                Map<String,Object> entryObject = attributeMapper.transform(entry);


                // Remove unrequested attributes that have got into the entry
                // See issue #2
                List<String> unwantedClaims = new ArrayList<>();

                for (String claimName: entryObject.keySet()) {

                        if (! claims.contains(claimName))
                                unwantedClaims.add(claimName);
                }

                for (String claimToRemove: unwantedClaims) {
                        entryObject.remove(claimToRemove);
                }

                // Append mandatory "sub" claim
                entryObject.put("sub", subject.getValue());

                try {
                        return new UserInfo(new JSONObject(entryObject));

                } catch (IllegalArgumentException e) {

                        throw new Exception("Couldn't create UserInfo object: " + e.getMessage(), e);
                }
        }


        @Override
        public void shutdown()
                throws Exception {

                if (ldapConnPool != null) {
                        // Close the LDAP connection pool
                        ldapConnPool.close();
                }
        }
}