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

import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.plugin.api.config.ConfigurationContext;
import com.airlock.iam.saml2.application.configuration.SamlConfig;
import com.airlock.iam.saml2.application.configuration.sp.Saml2FlowSpConfig;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpEntityId;
import com.airlock.iam.saml2.infrastructure.service.Saml2FlowSpConfigServiceImpl;
import com.airlock.iam.saml2.infrastructure.util.Saml2LoginappConfigProvider;
import com.airlock.iam.saml2.infrastructure.util.Saml2RequestThreadContext;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.plugins.SAML2ServiceProviderAdapter;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.Extensions;
import com.sun.identity.saml2.protocol.LogoutRequest;
import com.sun.identity.saml2.protocol.LogoutResponse;
import com.sun.identity.saml2.protocol.Response;
import com.sun.identity.saml2.protocol.impl.ExtensionsImpl;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.apache.commons.lang3.StringUtils;

public final class Saml2ServiceProviderAdapter extends SAML2ServiceProviderAdapter {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2ServiceProviderAdapter.class);

   @Override
   public void initialize(Map<String, String> initParams) {
   }

   @Override
   public void postSingleLogoutSuccess(
      String spEntityId,
      String realm,
      HttpServletRequest request,
      HttpServletResponse response,
      String userID,
      LogoutRequest logoutRequest,
      LogoutResponse logoutResponse,
      String binding
   ) {
   }

   @Override
   public boolean postSingleSignOnFailure(
      String spEntityId,
      String realm,
      HttpServletRequest request,
      HttpServletResponse response,
      AuthnRequest authnRequest,
      Response ssoResponse,
      String profile,
      int failureCode
   ) {
      return false;
   }

   @Override
   public boolean postSingleSignOnSuccess(
      String spEntityId,
      String realm,
      HttpServletRequest request,
      HttpServletResponse response,
      Object session,
      AuthnRequest authnRequest,
      Response ssoResponse,
      String profile,
      boolean isFederation
   ) {
      return false;
   }

   @Override
   public void preSingleLogoutProcess(
      String spEntityId,
      String realm,
      HttpServletRequest request,
      HttpServletResponse response,
      String userID,
      LogoutRequest logoutRequest,
      LogoutResponse logoutResponse,
      String binding
   ) {
   }

   @Override
   public void preSingleSignOnProcess(
      String spEntityId,
      String realm,
      HttpServletRequest request,
      HttpServletResponse response,
      AuthnRequest authnRequest,
      Response ssoResponse,
      String profile
   ) {
   }

   @Override
   public void preSingleSignOnRequest(
      String spEntityId, String idpEntityID, String realm, HttpServletRequest request, HttpServletResponse response, AuthnRequest authnRequest
   ) throws SAML2Exception {
      ConfigurationContext context = Saml2RequestThreadContext.getConfigurationContext();
      addCustomExtensionsToAuthnRequest(spEntityId, authnRequest, context);
   }

   private static void addCustomExtensionsToAuthnRequest(String spEntityId, AuthnRequest authnRequest, ConfigurationContext context) throws SAML2Exception {
      Optional<SamlConfig> samlConfig = Saml2LoginappConfigProvider.getSamlConfig(context);
      Optional<String> configuredExtensions = new Saml2FlowSpConfigServiceImpl(() -> samlConfig)
         .findFlowSpConfig(Saml2SpEntityId.spEntityId(spEntityId))
         .flatMap(Saml2FlowSpConfig::getCustomAuthnRequestExtensions)
         .filter(StringUtils::isNotBlank);
      if (configuredExtensions.isPresent()) {
         LOG.debug("adding custom extensions to AuthnRequest");
         Extensions ext = authnRequest.getExtensions();
         if (ext == null) {
            ext = new ExtensionsImpl();
         }

         List<String> extValues = ext.getAny();
         if (extValues == null) {
            extValues = new ArrayList<>();
         }

         extValues.add(configuredExtensions.get());
         ext.setAny(extValues);
         authnRequest.setExtensions(ext);
      }
   }
}