package com.sun.identity.saml2.plugins;

import com.sun.identity.plugin.datastore.DataStoreProviderException;
import com.sun.identity.plugin.session.SessionException;
import com.sun.identity.plugin.session.SessionManager;
import com.sun.identity.saml2.assertion.AssertionFactory;
import com.sun.identity.saml2.assertion.Attribute;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.shared.xml.XMLUtils;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.StringTokenizer;

public class DefaultLibraryIDPAttributeMapper extends DefaultAttributeMapper implements IDPAttributeMapper {
   @Override
   public List<Attribute> getAttributes(Object session, String hostEntityID, String remoteEntityID, String realm) throws SAML2Exception {
      if (hostEntityID == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullHostEntityID"));
      }

      if (realm == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullHostEntityID"));
      }

      if (session == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullSSOToken"));
      }

      try {
         if (!SessionManager.getProvider().isValid(session)) {
            if (SAML2SDKUtils.DEBUG.warningEnabled()) {
               SAML2SDKUtils.DEBUG.warning("DefaultLibraryIDPAttributeMapper.getAttributes: Invalid session");
            }

            return null;
         } else {
            Map<String, String> configMap = this.getConfigAttributeMap(realm, remoteEntityID, "SPRole");
            if (SAML2SDKUtils.DEBUG.messageEnabled()) {
               SAML2SDKUtils.DEBUG.message("DefaultLibraryIDPAttributeMapper.getAttributes: remote SP attribute map = " + configMap);
            }

            if (configMap == null || configMap.isEmpty()) {
               configMap = this.getConfigAttributeMap(realm, hostEntityID, "IDPRole");
               if (configMap == null || configMap.isEmpty()) {
                  if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                     SAML2SDKUtils.DEBUG.message("DefaultLibraryIDPAttributeMapper.getAttributes: Configuration map is not defined.");
                  }

                  return null;
               }

               if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                  SAML2SDKUtils.DEBUG.message("DefaultLibraryIDPAttributeMapper.getAttributes: hosted IDP attribute map=" + configMap);
               }
            }

            List<Attribute> attributes = new ArrayList<>();
            Set<String> localAttributes = new HashSet<>(configMap.values());
            Map<String, Set<String>> valueMap = null;
            if (!this.isDynamicalOrIgnoredProfile(realm)) {
               try {
                  valueMap = dsProvider.getAttributes(SessionManager.getProvider().getPrincipalName(session), localAttributes);
               } catch (DataStoreProviderException dse) {
                  if (SAML2SDKUtils.DEBUG.warningEnabled()) {
                     SAML2SDKUtils.DEBUG.warning("DefaultLibraryIDPAttributeMapper.getAttributes:", dse);
                  }
               }
            }

            for (String samlAttribute : configMap.keySet()) {
               String localAttribute = configMap.get(samlAttribute);
               String nameFormat = null;
               StringTokenizer tokenizer = new StringTokenizer(samlAttribute, "|");
               if (tokenizer.countTokens() > 1) {
                  nameFormat = tokenizer.nextToken();
                  samlAttribute = tokenizer.nextToken();
               }

               String[] localAttributeValues = null;
               if (valueMap != null && !valueMap.isEmpty()) {
                  Set<String> values = valueMap.get(localAttribute);
                  if (values != null && !values.isEmpty()) {
                     localAttributeValues = values.toArray(new String[0]);
                  } else if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                     SAML2SDKUtils.DEBUG
                        .message(
                           "DefaultLibraryIDPAttributeMapper.getAttribute: user profile does not have value for "
                              + localAttribute
                              + " but is going to check ssotoken:"
                        );
                  }
               }

               if (localAttributeValues == null) {
                  localAttributeValues = SessionManager.getProvider().getProperty(session, localAttribute);
               }

               if (localAttributeValues != null && localAttributeValues.length != 0) {
                  attributes.add(this.getSAMLAttribute(samlAttribute, nameFormat, localAttributeValues));
               } else if (SAML2SDKUtils.DEBUG.messageEnabled()) {
                  SAML2SDKUtils.DEBUG.message("DefaultLibraryIDPAttributeMapper.getAttribute: user does not have " + localAttribute);
               }
            }

            return attributes;
         }
      } catch (SessionException se) {
         SAML2SDKUtils.DEBUG.error("DefaultLibraryIDPAttribute.getAttributes: ", se);
         throw new SAML2Exception(se);
      }
   }

   protected Attribute getSAMLAttribute(String name, String nameFormat, String[] values) throws SAML2Exception {
      if (name == null) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("nullInput"));
      }

      AssertionFactory factory = AssertionFactory.getInstance();
      Attribute attribute = factory.createAttribute();
      attribute.setName(name);
      if (nameFormat != null) {
         attribute.setNameFormat(nameFormat);
      }

      if (values != null) {
         List<String> list = new ArrayList<>();

         for (String value : values) {
            list.add(XMLUtils.escapeSpecialCharacters(value));
         }

         attribute.setAttributeValueString(list);
      }

      return attribute;
   }

   protected boolean isDynamicalOrIgnoredProfile(String realm) {
      return true;
   }
}