package com.airlock.iam.saml2.infrastructure.service.sp.logout;

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.common.application.service.session.SessionTerminationService;
import com.airlock.iam.common.infrastructure.factory.uri.AbsoluteExternalViewUriFactory;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.domain.model.session.FlowFederationSession;
import com.airlock.iam.saml2.domain.model.session.FlowFederationSessionRepository;
import com.airlock.iam.saml2.domain.model.sp.logout.Saml2SpSloContinuationId;
import com.airlock.iam.saml2.domain.service.Saml2FlowFederationSessionTerminationService;
import com.airlock.iam.saml2.infrastructure.service.sp.Saml2SpFederationLibrary;
import com.airlock.iam.saml2.infrastructure.sp.Saml2SpJerseyUrls;
import com.airlock.iam.saml2.infrastructure.sp.logout.Saml2IdpInitiatedSloCache;
import com.sun.identity.saml2.profile.LogoutRequestDataContainer;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import jakarta.ws.rs.core.UriBuilder;
import java.net.URI;

public final class Saml2IdpInitiatedSloServiceImpl {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2IdpInitiatedSloServiceImpl.class);
   private final Provider<AbsoluteExternalViewUriFactory> externalViewUriFactory;
   private final Provider<Saml2SpFederationLibrary> library;
   private final Provider<Saml2IdpInitiatedSloCache> idpInitiatedSloCache;
   private final Provider<URI> uiLogoutUri;
   private final Provider<FlowFederationSessionRepository> federationSessionRepository;
   private final Provider<Saml2FlowFederationSessionTerminationService> federationSessionTerminationService;
   private final Provider<SessionTerminationService> sessionTerminationService;

   @Inject
   Saml2IdpInitiatedSloServiceImpl(
      Provider<AbsoluteExternalViewUriFactory> externalViewUriFactory,
      Provider<Saml2SpFederationLibrary> library,
      Provider<Saml2IdpInitiatedSloCache> idpInitiatedSloCache,
      Provider<URI> uiLogoutUri,
      Provider<FlowFederationSessionRepository> federationSessionRepository,
      Provider<Saml2FlowFederationSessionTerminationService> federationSessionTerminationService,
      Provider<SessionTerminationService> sessionTerminationService
   ) {
      this.externalViewUriFactory = externalViewUriFactory;
      this.library = library;
      this.idpInitiatedSloCache = idpInitiatedSloCache;
      this.uiLogoutUri = uiLogoutUri;
      this.federationSessionRepository = federationSessionRepository;
      this.federationSessionTerminationService = federationSessionTerminationService;
      this.sessionTerminationService = sessionTerminationService;
   }

   public URI processIdpInitiatedSlo() {
      LogoutRequestDataContainer validatedLogoutRequest = ((Saml2SpFederationLibrary)this.library.get()).validateIdpInitiatedSloRequest();
      LOG.info(
         "Received Single-Logout request from IdP %s for SP %s"
            .formatted(LogUtil.logSafeQuote(validatedLogoutRequest.getIdpEntityID()), LogUtil.logSafeQuote(validatedLogoutRequest.getSpEntityID()))
      );
      FlowFederationSession session = ((FlowFederationSessionRepository)this.federationSessionRepository.get())
         .getFederationSession()
         .orElseGet(() -> ((FlowFederationSessionRepository)this.federationSessionRepository.get()).createNewFederationSession());
      session.setIdpInitiatedSingleLogoutInProgressOnSp();
      Saml2SpSloContinuationId id = ((Saml2IdpInitiatedSloCache)this.idpInitiatedSloCache.get()).save(session, validatedLogoutRequest);
      LOG.info("Redirecting to the UI to perform the local logout before processing the SAML 2.0 logout request");
      return this.redirectUri(id);
   }

   private URI redirectUri(Saml2SpSloContinuationId id) {
      return UriBuilder.fromUri(((AbsoluteExternalViewUriFactory)this.externalViewUriFactory.get()).createAbsoluteUri((URI)this.uiLogoutUri.get()))
         .queryParam("Location", new Object[]{this.sloContinueUrl(id)})
         .queryParam("skipDisclaimer", new Object[]{"true"})
         .build(new Object[0]);
   }

   private String sloContinueUrl(Saml2SpSloContinuationId id) {
      return UriBuilder.fromUri(((AbsoluteExternalViewUriFactory)this.externalViewUriFactory.get()).createAbsoluteUri(Saml2SpJerseyUrls.SLO_CONTINUE_URI.uri()))
         .queryParam("id", new Object[]{id.value()})
         .build(new Object[0])
         .toASCIIString();
   }

   public void continueIdpInitiatedSlo(String unverifiedSloId) {
      Saml2IdpInitiatedSloCache.Saml2IdpInitiatedSloCacheEntry cacheEntry = this.retrieveSloCacheEntry(
         Saml2SpSloContinuationId.parseSloContinuationId(unverifiedSloId)
      );

      try {
         this.restoreFederationSession(cacheEntry.federationSession());
         ((Saml2SpFederationLibrary)this.library.get()).continueIdpInitiatedSlo(cacheEntry.logoutData());
         ((SessionTerminationService)this.sessionTerminationService.get()).terminateSessions();
      } finally {
         ((Saml2FlowFederationSessionTerminationService)this.federationSessionTerminationService.get()).terminateSession();
      }
   }

   private void restoreFederationSession(FlowFederationSession federationSession) {
      if (!federationSession.isSingleLogoutInProgress()) {
         throw new FailureException(
            "Unable to do IdP-initiated SLO, either user's FlowFederationSession has not yet been marked for SLO, or SLO has already been performed."
         );
      }

      if (!federationSession.isIdpInitiatedSingleLogoutInProgressOnSp()) {
         throw new FailureException(
            "Request received to do an IdP-initiated SLO (single logout) but another type of SLO has already been started on this session before. It looks like a different SLO process has been interrupted (e.g. a SAML SP failed to redirect the client back to us). Continuing logout."
         );
      }

      ((FlowFederationSessionRepository)this.federationSessionRepository.get()).updateFederationSession(federationSession);
   }

   private Saml2IdpInitiatedSloCache.Saml2IdpInitiatedSloCacheEntry retrieveSloCacheEntry(Saml2SpSloContinuationId id) {
      return ((Saml2IdpInitiatedSloCache)this.idpInitiatedSloCache.get())
         .readOnce(id)
         .orElseThrow(() -> new ViolationException("IdP-Initiated SLO: unknown SLO continuation ID %s".formatted(LogUtil.logSafeQuote(id.toString()))));
   }
}