package com.airlock.iam.saml2.infrastructure.plugin;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.opensso.jaxb.entityconfig.BaseConfigType;
import com.airlock.iam.saml2.domain.model.idp.Saml2AuthenticationContextClassRef;
import com.airlock.iam.saml2.domain.model.logging.Saml2LogTypes;
import com.airlock.iam.saml2.domain.model.session.FlowFederationSession;
import com.airlock.iam.saml2.domain.model.session.Saml2PropagationAttributes;
import com.airlock.iam.saml2.infrastructure.util.Saml2RequestThreadContext;
import com.google.common.annotations.VisibleForTesting;
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.AuthnContext;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.meta.SAML2MetaException;
import com.sun.identity.saml2.meta.SAML2MetaUtils;
import com.sun.identity.saml2.plugins.IDPAuthnContextInfo;
import com.sun.identity.saml2.plugins.IDPAuthnContextMapper;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.RequestedAuthnContext;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.function.Supplier;

public final class IDPAuthnContextMapperImpl implements IDPAuthnContextMapper {
   private static final ELogger LOG = ELoggerFactory.getLogger(IDPAuthnContextMapperImpl.class);
   private final Supplier<SessionProvider> sessionProviderSupplier;
   private final IDPAuthnContextMapperImpl.AllowedAuthnContextRefsSupplier allowedAuthnContextRefs;

   public IDPAuthnContextMapperImpl() {
      this.sessionProviderSupplier = () -> {
         try {
            return SessionManager.getProvider();
         } catch (SessionException e) {
            throw new FailureException("No Session Provider available. " + e.getMessage(), e);
         }
      };
      this.allowedAuthnContextRefs = this::getAllowedAuthnContextRefs;
   }

   @VisibleForTesting
   IDPAuthnContextMapperImpl(SessionProvider sessionProvider, IDPAuthnContextMapperImpl.AllowedAuthnContextRefsSupplier authnContextRefsSupplier) {
      this.sessionProviderSupplier = () -> sessionProvider;
      this.allowedAuthnContextRefs = authnContextRefsSupplier;
   }

   @Override
   public IDPAuthnContextInfo getIDPAuthnContextInfo(AuthnRequest authnRequest, String idpEntityID, String realm) throws SAML2Exception {
      try {
         this.getSessionProvider().getSession(Saml2RequestThreadContext.getRequest());
      } catch (SessionException e) {
         LOG.info(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "No FlowFederationSession available yet");
         return new IDPAuthnContextInfo(null, new HashSet<>());
      }

      List<String> requestedClassRefs = new ArrayList<>();
      Optional<String> classRefFromAuth = this.getClassRefFromAuth();
      if (classRefFromAuth.isPresent()) {
         LOG.debug("Using requested authn classRef from authentication: " + classRefFromAuth.get());
         requestedClassRefs.add(classRefFromAuth.get());
      } else if (authnRequest != null) {
         LOG.debug("Analyzing AuthnRequest from SP for requested authn classrefs");
         RequestedAuthnContext requestedAuthnContext = authnRequest.getRequestedAuthnContext();
         if (requestedAuthnContext != null) {
            List<String> requestedRefs = requestedAuthnContext.getAuthnContextClassRef();
            if (requestedRefs != null && !requestedRefs.isEmpty()) {
               requestedClassRefs.addAll(requestedRefs);
               LOG.debug("Requested authn classrefs by the SP: " + requestedClassRefs);
            }
         }
      }

      return this.getAuthnContextInfo(requestedClassRefs, idpEntityID, realm);
   }

   @Override
   public AuthnContext getAuthnContextFromAuthLevel(String authLevel, String realm, String idpEntityID) throws SAML2Exception {
      try {
         this.getSessionProvider().getSession(Saml2RequestThreadContext.getRequest());
      } catch (SessionException e) {
         LOG.error(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "No FlowFederationSession available yet");
         throw new SAML2Exception("User has no FlowFederationSession yet");
      }

      List<String> requestedClassRefs = new ArrayList<>();
      this.getClassRefFromAuth().ifPresent(classRef -> {
         LOG.debug("Using configured authn classref from authentication: " + classRef);
         requestedClassRefs.add(classRef);
      });
      IDPAuthnContextInfo res = this.getAuthnContextInfo(requestedClassRefs, idpEntityID, realm);
      return res.getAuthnContext();
   }

   private Optional<String> getClassRefFromAuth() {
      FlowFederationSession ses;
      try {
         ses = FlowFederationSession.toFlowFederationSessionOrThrow(
            this.getSessionProvider().getSession(Saml2RequestThreadContext.getRequest()), () -> new IllegalStateException("Invalid session type")
         );
      } catch (SessionException e) {
         return Optional.empty();
      }

      return ses.getPropagationAttributes().flatMap(Saml2PropagationAttributes::getAuthenticationContext).map(Saml2AuthenticationContextClassRef::value);
   }

   private SessionProvider getSessionProvider() {
      return this.sessionProviderSupplier.get();
   }

   private IDPAuthnContextInfo getAuthnContextInfo(List<String> requestedClassRefs, String idpEntityID, String realm) throws SAML2Exception {
      List<String> idpAuthnContextRefs = this.allowedAuthnContextRefs.getAllowedAuthnContextRefs(idpEntityID, realm);
      if (idpAuthnContextRefs != null && !idpAuthnContextRefs.isEmpty()) {
         String firstClassRef = null;
         Set<String> firstAuthTypeAndValues = new HashSet<>();
         String classRef = null;
         Set<String> authTypeAndValues = new HashSet<>();
         int chosenRequestedClassRefIndex = Integer.MAX_VALUE;

         for (String value : idpAuthnContextRefs) {
            value = value.trim();
            if (LOG.isDebugEnabled()) {
               LOG.debug("Considering configured IDP authn classref: '" + value + "'");
            }

            StringTokenizer st = new StringTokenizer(value, "|");
            if (st.hasMoreTokens()) {
               String currentRef = st.nextToken().trim();
               Set<String> currentOptions = new HashSet<>();

               while (st.hasMoreTokens()) {
                  String authTypeAndValue = st.nextToken().trim();
                  if (!authTypeAndValue.isEmpty()) {
                     currentOptions.add(authTypeAndValue);
                  }
               }

               if (requestedClassRefs.isEmpty()) {
                  classRef = currentRef;
                  authTypeAndValues = currentOptions;
                  break;
               }

               int index = requestedClassRefs.indexOf(currentRef);
               if (index >= 0 && index < chosenRequestedClassRefIndex) {
                  classRef = currentRef;
                  authTypeAndValues = currentOptions;
                  if (index == 0) {
                     break;
                  }

                  chosenRequestedClassRefIndex = index;
               }

               if (firstClassRef == null) {
                  firstClassRef = currentRef;
                  firstAuthTypeAndValues = currentOptions;
               }
            }
         }

         if (classRef == null) {
            LOG.info("None of the requested authn classrefs found, using IDP default: " + firstClassRef);
            classRef = firstClassRef;
            authTypeAndValues = firstAuthTypeAndValues;
         }

         AuthnContext authnContext = AssertionFactory.getInstance().createAuthnContext();
         authnContext.setAuthnContextClassRef(classRef);
         IDPAuthnContextInfo info = new IDPAuthnContextInfo(authnContext, authTypeAndValues);
         if (LOG.isDebugEnabled()) {
            LOG.debug(
               "requested AuthnContextClassRef(s)="
                  + requestedClassRefs
                  + ", returned AuthnContextClassRef="
                  + classRef
                  + ", authTypeAndValues="
                  + authTypeAndValues
            );
         }

         return info;
      } else {
         LOG.error(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "IDP does not specify the attribute 'idpAuthncontextClassrefMapping' in its extended metadata.");
         throw new SAML2Exception("IDP does not specify the attribute 'idpAuthncontextClassrefMapping' in its extended metadata.");
      }
   }

   private List<String> getAllowedAuthnContextRefs(String idpEntityID, String realm) throws SAML2Exception {
      Map<String, List<String>> idpAttrs;
      try {
         BaseConfigType config = SAML2Utils.getSAML2MetaManager().getIDPSSOConfig(realm, idpEntityID);
         idpAttrs = SAML2MetaUtils.getAttributes(config);
      } catch (SAML2MetaException sme) {
         LOG.error(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "get IDPSSOConfig failed for entity " + idpEntityID, sme);
         throw new SAML2Exception(sme);
      }

      return idpAttrs.get("idpAuthncontextClassrefMapping");
   }

   @Override
   public boolean isAuthnContextMatching(List requestedACClassRefs, String acClassRef, String comparison, String realm, String idpEntityID) {
      throw new UnsupportedOperationException("AuthnQuery not implemented.");
   }

   @FunctionalInterface
   interface AllowedAuthnContextRefsSupplier {
      List<String> getAllowedAuthnContextRefs(String var1, String var2) throws SAML2Exception;
   }
}