package com.sun.identity.saml2.common;

import com.sun.identity.plugin.datastore.DataStoreProviderException;
import com.sun.identity.shared.configuration.SystemPropertiesManager;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class AccountUtils {
   private static final String DELIM = "|";
   private static final String NAMEID_INFO_ATTRIBUTE = "com.sun.identity.saml2.nameidinfo.attribute";
   private static final String NAMEID_INFO_KEY_ATTRIBUTE = "com.sun.identity.saml2.nameidinfokey.attribute";

   public static NameIDInfo getAccountFederation(String userID, String hostEntityID, String remoteEntityID) throws SAML2Exception {
      SAML2Utils.DEBUG.message("AccountUtils.getAccountFederation:");
      if (userID == null) {
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullUserID"));
      }

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

      if (remoteEntityID == null) {
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullRemoteEntityID"));
      }

      try {
         Set<String> set = SAML2Utils.getDataStoreProvider().getAttribute(userID, getNameIDInfoAttribute());
         if (set != null && !set.isEmpty()) {
            String filter = hostEntityID + "|" + remoteEntityID + "|";
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message("AccountUtils.getAccountFederation:  filter = " + filter + " userID = " + userID);
            }

            String info = null;

            for (String value : set) {
               if (value.startsWith(filter)) {
                  info = value;
                  break;
               }
            }

            if (info == null) {
               if (SAML2Utils.DEBUG.messageEnabled()) {
                  SAML2Utils.DEBUG.message("AccountUtils.getAccountFederation : user does not have account federation  corresponding to =" + filter);
               }

               return null;
            } else {
               return NameIDInfo.parse(info);
            }
         } else {
            if (SAML2Utils.DEBUG.messageEnabled()) {
               SAML2Utils.DEBUG.message("AccountUtils.getAccountFederation : user does not have any account federations.");
            }

            return null;
         }
      } catch (DataStoreProviderException dse) {
         SAML2Utils.DEBUG.error("AccountUtils.readAccountFederationInfo: DataStoreProviderException", dse);
         throw new SAML2Exception(dse.getMessage());
      }
   }

   public static void setAccountFederation(NameIDInfo info, String userID) throws SAML2Exception {
      SAML2Utils.DEBUG.message("AccountUtils.setAccountFederation:");
      if (info == null) {
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullNameIDInfo"));
      }

      if (userID == null) {
         throw new SAML2Exception(SAML2Utils.BUNDLE.getString("nullUserID"));
      }

      NameIDInfoKey infoKey = new NameIDInfoKey(info.getNameIDValue(), info.getHostEntityID(), info.getRemoteEntityID());
      if (SAML2Utils.DEBUG.messageEnabled()) {
         SAML2Utils.DEBUG
            .message("AccountUtils.setAccountFederation: info to be set:" + info.toValueString() + ",infoKey to be set:" + infoKey.toValueString());
      }

      String filter = info.getHostEntityID() + "|" + info.getRemoteEntityID() + "|";

      try {
         String nameIDInfoAttr = getNameIDInfoAttribute();
         String nameIDInfoKeyAttr = getNameIDInfoKeyAttribute();
         Set<String> set = new HashSet<>();
         set.add(nameIDInfoAttr);
         set.add(nameIDInfoKeyAttr);
         Map<String, Set<String>> map = new HashMap<>();
         Map<String, Set<String>> existMap = SAML2Utils.getDataStoreProvider().getAttributes(userID, set);
         if (existMap != null && !existMap.isEmpty()) {
            Set<String> set1 = existMap.get(nameIDInfoAttr);
            if (set1 != null) {
               set1.removeIf(value -> value.startsWith(filter));
            } else {
               set1 = new HashSet<>();
            }

            set1.add(info.toValueString());
            map.put(nameIDInfoAttr, set1);
            Set<String> set2 = existMap.get(nameIDInfoKeyAttr);
            if (set2 != null) {
               set2.removeIf(value -> value.startsWith(filter));
            } else {
               set2 = new HashSet<>();
            }

            set2.add(infoKey.toValueString());
            map.put(nameIDInfoKeyAttr, set2);
         } else {
            Set<String> set1 = new HashSet<>();
            set1.add(infoKey.toValueString());
            map.put(nameIDInfoKeyAttr, set1);
            Set<String> set2 = new HashSet<>();
            set2.add(info.toValueString());
            map.put(nameIDInfoAttr, set2);
         }

         if (SAML2Utils.DEBUG.messageEnabled()) {
            SAML2Utils.DEBUG.message("AccountUtils.setAccountFederation: set fedinfo " + map + " userID = " + userID);
         }

         SAML2Utils.getDataStoreProvider().setAttributes(userID, map);
      } catch (DataStoreProviderException dse) {
         SAML2Utils.DEBUG.error("SAML2Utils.setAccountFederation: DataStoreProviderException", dse);
         throw new SAML2Exception(dse.getMessage());
      }
   }

   public static String getNameIDInfoAttribute() {
      return SystemPropertiesManager.get("com.sun.identity.saml2.nameidinfo.attribute", "sun-fm-saml2-nameid-info");
   }

   public static String getNameIDInfoKeyAttribute() {
      return SystemPropertiesManager.get("com.sun.identity.saml2.nameidinfokey.attribute", "sun-fm-saml2-nameid-infokey");
   }
}
