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.plugin.session.SessionProvider;
import com.sun.identity.saml2.assertion.AssertionFactory;
import com.sun.identity.saml2.assertion.NameID;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.profile.IDPCache;
import com.sun.identity.saml2.profile.IDPSSOUtil;
import com.sun.identity.saml2.profile.IDPSession;
import com.sun.identity.saml2.profile.NameIDandSPpair;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class DefaultIDPAccountMapper extends DefaultAccountMapper implements IDPAccountMapper {
   public DefaultIDPAccountMapper() {
      SAML2SDKUtils.DEBUG.message("DefaultIDPAccountMapper.constructor");
      this.role = "IDPRole";
   }

   @Override
   public NameID getNameID(Object session, String hostEntityID, String remoteEntityID, String realm, String nameIDFormat) throws SAML2Exception {
      String userID = null;

      try {
         SessionProvider sessionProv = SessionManager.getProvider();
         userID = sessionProv.getPrincipalName(session);
      } catch (SessionException se) {
         throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("invalidSSOToken"));
      }

      String nameIDValue = null;
      if (nameIDFormat.equals("urn:oasis:names:tc:SAML:2.0:nameid-format:transient")) {
         String sessionIndex = IDPSSOUtil.getSessionIndex(session);
         if (sessionIndex != null) {
            IDPSession idpSession = IDPCache.idpSessionsByIndices().get(sessionIndex);
            if (idpSession != null) {
               List<NameIDandSPpair> list = idpSession.getNameIDandSPpairs();
               if (list != null && !list.isEmpty()) {
                  for (NameIDandSPpair pair : list) {
                     if (pair.getSPEntityID().equals(remoteEntityID)) {
                        nameIDValue = pair.getNameID().getValue();
                        break;
                     }
                  }
               }
            }
         }

         if (nameIDValue == null) {
            nameIDValue = this.getNameIDValueFromUserProfile(realm, hostEntityID, userID, nameIDFormat);
            if (nameIDValue == null) {
               nameIDValue = SAML2Utils.createNameIdentifier();
            }
         }
      } else {
         nameIDValue = this.getNameIDValueFromUserProfile(realm, hostEntityID, userID, nameIDFormat);
         if (nameIDValue == null) {
            if (!nameIDFormat.equals("urn:oasis:names:tc:SAML:2.0:nameid-format:persistent")) {
               throw new SAML2Exception(SAML2SDKUtils.BUNDLE.getString("unableToGenerateNameIDValue"));
            }

            nameIDValue = SAML2Utils.createNameIdentifier();
         }
      }

      NameID nameID = AssertionFactory.getInstance().createNameID();
      nameID.setValue(nameIDValue);
      nameID.setFormat(nameIDFormat);
      nameID.setNameQualifier(hostEntityID);
      nameID.setSPNameQualifier(remoteEntityID);
      nameID.setSPProvidedID(null);
      return nameID;
   }

   private String getNameIDValueFromUserProfile(String realm, String hostEntityID, String userID, String nameIDFormat) {
      String nameIDValue = null;
      Map<String, String> formatAttrMap = this.getFormatAttributeMap(realm, hostEntityID);
      String attrName = formatAttrMap.get(nameIDFormat);
      if (attrName != null) {
         try {
            Set<String> attrValues = dsProvider.getAttribute(userID, attrName);
            if (attrValues != null && !attrValues.isEmpty()) {
               nameIDValue = attrValues.iterator().next();
            }
         } catch (DataStoreProviderException dspe) {
            if (SAML2SDKUtils.DEBUG.warningEnabled()) {
               SAML2SDKUtils.DEBUG.warning("DefaultIDPAccountMapper.getNameIDValueFromUserProfile:", dspe);
            }
         }
      }

      return nameIDValue;
   }

   private Map<String, String> getFormatAttributeMap(String realm, String hostEntityID) {
      String key = hostEntityID + "|" + realm;
      Map<String, String> formatAttributeMap = IDPCache.formatAttributeHash.get(key);
      if (formatAttributeMap != null) {
         return formatAttributeMap;
      }

      formatAttributeMap = new HashMap<>();
      List<String> values = SAML2Utils.getAllAttributeValueFromSSOConfig(realm, hostEntityID, this.role, "nameIDFormatMap");
      if (values != null && !values.isEmpty()) {
         for (String value : values) {
            int index = value.indexOf(61);
            if (index != -1) {
               String format = value.substring(0, index).trim();
               String attrName = value.substring(index + 1).trim();
               if (!format.isEmpty() && !attrName.isEmpty()) {
                  formatAttributeMap.put(format, attrName);
               }
            }
         }
      }

      IDPCache.formatAttributeHash.put(key, formatAttributeMap);
      return formatAttributeMap;
   }
}