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

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.domain.model.Saml2Binding;
import com.airlock.iam.saml2.domain.model.Saml2RelayState;
import com.airlock.iam.saml2.domain.model.idp.Saml2IdpEntityId;
import com.airlock.iam.saml2.domain.model.idp.Saml2IdpInitiatedSsoContinuationParameters;
import com.airlock.iam.saml2.domain.model.idp.Saml2IdpMetaAlias;
import com.airlock.iam.saml2.domain.model.idp.Saml2RequestParameter;
import com.airlock.iam.saml2.domain.model.idp.Saml2ResponseParameter;
import com.airlock.iam.saml2.domain.model.logging.Saml2LogTypes;
import com.airlock.iam.saml2.domain.service.Saml2FlowFederationSessionTerminationService;
import com.airlock.iam.saml2.infrastructure.request.RebuiltInitialSpInitiatedSloRequest;
import com.airlock.iam.saml2.infrastructure.request.Saml2Request;
import com.airlock.iam.saml2.infrastructure.request.Saml2Response;
import com.airlock.iam.saml2.infrastructure.service.idp.Saml2IdpFederationLibrary;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.profile.IDPArtifactResolution;
import com.sun.identity.saml2.profile.IDPSSOFederate;
import com.sun.identity.saml2.profile.IDPSSOUtil;
import com.sun.identity.saml2.profile.IDPSingleLogout;
import com.sun.identity.saml2.profile.LogoutUtil;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.InputStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.Callable;

public final class Saml2IdpFederationLibraryImpl implements Saml2IdpFederationLibrary {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2IdpFederationLibraryImpl.class);
   private final Provider<Saml2Request> request;
   private final Provider<Saml2Response> response;
   private final Provider<Saml2FlowFederationSessionTerminationService> sessionTerminationService;

   @Inject
   Saml2IdpFederationLibraryImpl(
      Provider<Saml2Request> request, Provider<Saml2Response> response, Provider<Saml2FlowFederationSessionTerminationService> sessionTerminationService
   ) {
      this.request = request;
      this.response = response;
      this.sessionTerminationService = sessionTerminationService;
   }

   @Override
   public void doIdpInitiatedSsoFederate(Saml2IdpInitiatedSsoContinuationParameters parameters) {
      LOG.info(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Using SAML 2.0 to redirect user to SP with " + parameters);

      try {
         IDPSSOUtil.doSSOFederate(
            (HttpServletRequest)this.request.get(),
            (HttpServletResponse)this.response.get(),
            null,
            parameters.getSpEntityId().value(),
            parameters.getIdpMetaAlias().value(),
            null,
            parameters.getRelayState().map(Saml2RelayState::value).orElse(null)
         );
      } catch (SAML2Exception e) {
         throw new FailureException("Error while initiating SAML2 SSO with " + parameters, e);
      }
   }

   @Override
   public void doSpInitiatedSsoFederate() {
      IDPSSOFederate.doSSOFederate((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get());
   }

   @Override
   public Saml2Binding doIdpInitiatedSlo(Saml2IdpMetaAlias metaAlias) {
      LOG.info(
         Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE,
         String.format("Using SAML 2.0 to do IdP initiated single logout for meta alias %s", LogUtil.logSafeQuote(metaAlias.value()))
      );
      return this.callLibrary(() -> {
         Saml2IdpEntityId idpEntityId = Saml2IdpEntityId.idpEntityId(SAML2Utils.getSAML2MetaManager().getEntityByMetaAlias(metaAlias.value()));
         Saml2Binding binding = this.getSloBinding(metaAlias, idpEntityId);
         Map<String, String> paramsMap = new HashMap<>();
         paramsMap.put("metaAlias", metaAlias.value());
         paramsMap.put("role", "IDPRole");
         paramsMap.put("binding", binding.value());
         IDPSingleLogout.initiateLogoutRequest((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get(), binding.value(), paramsMap);
         return binding;
      }, "Error while doing IdP initiated single logout");
   }

   @Override
   public boolean processSingleLogoutResponse(Saml2ResponseParameter logoutResponse, Optional<Saml2RelayState> relayState) {
      return this.callLibrary(
         () -> IDPSingleLogout.processLogoutResponse(
            (HttpServletRequest)this.request.get(),
            (HttpServletResponse)this.response.get(),
            logoutResponse.value(),
            relayState.map(Saml2RelayState::value).orElse(null)
         ),
         "Error while processing SAML 2.0 Logout Response"
      );
   }

   @Override
   public void resumeSpInitiatedSlo(
      RebuiltInitialSpInitiatedSloRequest originalSpInitiatedSloRequest, Saml2RequestParameter logoutRequest, Optional<Saml2RelayState> relayState
   ) {
      this.callLibrary(
         () -> {
            IDPSingleLogout.processLogoutRequest(
               originalSpInitiatedSloRequest,
               (HttpServletResponse)this.response.get(),
               logoutRequest.value(),
               relayState.map(Saml2RelayState::value).orElse(null)
            );
            return null;
         },
         "Error while processing SP-initiated single logout request"
      );
   }

   @Override
   public void resolveArtifact(InputStream is) {
      IDPArtifactResolution.doArtifactResolution((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get(), () -> is);
      LOG.info("Artifact Resolution request finished");
   }

   private <V> V callLibrary(Callable<V> library, String errorMessage) {
      try {
         return library.call();
      } catch (SAML2Exception e) {
         ((Saml2FlowFederationSessionTerminationService)this.sessionTerminationService.get()).terminateSession();
         throw new FailureException(errorMessage + ": " + e.getMessage(), e);
      } catch (Exception e) {
         throw new FailureException(errorMessage + ": " + e.getMessage(), e);
      }
   }

   private Saml2Binding getSloBinding(Saml2IdpMetaAlias metaAlias, Saml2IdpEntityId idpEntityId) throws SAML2Exception {
      String binding = LogoutUtil.getSLOBindingInfo((HttpServletRequest)this.request.get(), metaAlias.value(), "SPRole", idpEntityId.value());
      return Arrays.stream(Saml2Binding.values())
         .filter(e -> e.value().equals(binding))
         .findFirst()
         .orElseThrow(() -> new ViolationException(String.format("Unsupported type %s.", binding)));
   }
}