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.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.application.service.Saml2LogoutResumeRedirectService;
import com.airlock.iam.saml2.application.service.sp.Saml2SpInitiatedSloService;
import com.airlock.iam.saml2.domain.model.Saml2MetaManager;
import com.airlock.iam.saml2.domain.model.idp.Saml2IdpEntityId;
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.FlowFederationSessionRepository;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpEntityId;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpMetaAlias;
import com.airlock.iam.saml2.domain.model.sp.logout.Saml2SpInitiatedSloId;
import com.airlock.iam.saml2.infrastructure.service.sp.Saml2SpFederationLibrary;
import com.airlock.iam.saml2.infrastructure.sp.logout.Saml2SpInitiatedSloFederationSessionCache;
import com.sun.identity.saml2.common.AccountUtils;
import com.sun.identity.saml2.common.NameIDInfoKey;
import com.sun.identity.saml2.common.SAML2Exception;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import jakarta.ws.rs.core.UriBuilder;
import java.net.URI;
import java.util.Optional;
import java.util.regex.Pattern;
import org.apache.commons.lang3.Validate;

public class Saml2SpInitiatedSloServiceImpl implements Saml2SpInitiatedSloService {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2SpInitiatedSloServiceImpl.class);
   private static final String SLO_INIT_WARNING_MESSAGE = "Cannot forward control to IdP. Redirecting to UI logout resume URI in order to continue the logout procedure.";
   private final Provider<Pattern> uiLogoutResumeUriPattern;
   private final Provider<Saml2LogoutResumeRedirectService> logoutResumeRedirectService;
   private final Provider<Saml2SpInitiatedSloFederationSessionCache> federationSessionCache;
   private final Provider<FlowFederationSessionRepository> federationSessionRepository;
   private final Provider<Saml2SpFederationLibrary> library;
   private final Provider<Saml2MetaManager> metaManager;

   @Inject
   public Saml2SpInitiatedSloServiceImpl(
      Provider<Pattern> uiLogoutResumeUriPattern,
      Provider<Saml2LogoutResumeRedirectService> logoutResumeRedirectService,
      Provider<Saml2SpInitiatedSloFederationSessionCache> federationSessionCache,
      Provider<FlowFederationSessionRepository> federationSessionRepository,
      Provider<Saml2SpFederationLibrary> library,
      Provider<Saml2MetaManager> metaManager
   ) {
      this.uiLogoutResumeUriPattern = uiLogoutResumeUriPattern;
      this.logoutResumeRedirectService = logoutResumeRedirectService;
      this.federationSessionCache = federationSessionCache;
      this.federationSessionRepository = federationSessionRepository;
      this.library = library;
      this.metaManager = metaManager;
   }

   @Override
   public URI validateAndStoreUiLogoutResumeUri(String uiLogoutResumeUriParameter) {
      URI uiLogoutResumeUri = this.validateUiLogoutResumeUri(uiLogoutResumeUriParameter);
      ((Saml2LogoutResumeRedirectService)this.logoutResumeRedirectService.get()).storeUiLogoutResumeUri(uiLogoutResumeUri);
      return uiLogoutResumeUri;
   }

   @Override
   public URI retrieveAndValidateUiLogoutResumeUri() {
      return this.validateUiLogoutResumeUri(this.retrieveUiLogoutUri().toASCIIString());
   }

   @Override
   public Optional<URI> startSpInitiatedSlo(Saml2SpInitiatedSloId saml2SpInitiatedSloId, URI validatedUiLogoutResumeUri) {
      Optional<FlowFederationSession> flowFederationSession = ((Saml2SpInitiatedSloFederationSessionCache)this.federationSessionCache.get())
         .readOnce(saml2SpInitiatedSloId);
      if (flowFederationSession.isEmpty()) {
         return this.cleanUpSessionAndLogSloInitWarning("Unknown SP-initiated SLO ID received.", validatedUiLogoutResumeUri);
      }

      try {
         ((FlowFederationSessionRepository)this.federationSessionRepository.get()).updateFederationSession(flowFederationSession.get());
      } catch (ViolationException e) {
         return this.cleanUpSessionAndLogSloInitWarning(e.getMessage(), validatedUiLogoutResumeUri);
      }

      try {
         Saml2SpInitiatedSloServiceImpl.EntityIds entityIds = this.resolveEntityIds(flowFederationSession.get());
         Saml2SpMetaAlias metaAlias = ((Saml2MetaManager)this.metaManager.get()).spMetaAlias(entityIds.hostEntityId());
         if (!((Saml2SpFederationLibrary)this.library.get()).initiateLogoutRequest(metaAlias, entityIds.idpEntityId(), validatedUiLogoutResumeUri)) {
            return this.cleanUpSessionAndLogSloInitWarning("SAML library was unable to successfully initiate the SLO.", validatedUiLogoutResumeUri);
         }
      } catch (FailureException e) {
         return this.cleanUpSessionAndLogSloInitWarning(e.getMessage(), validatedUiLogoutResumeUri);
      }

      LOG.info("SAML 2.0 SP-initiated SLO on SP: Done initiating SLO.");
      return Optional.empty();
   }

   private Optional<URI> cleanUpSessionAndLogSloInitWarning(String message, URI validatedUiLogoutResumeUri) {
      ((Saml2LogoutResumeRedirectService)this.logoutResumeRedirectService.get()).cleanupSessions();
      LOG.warn(
         Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE,
         message + " - Cannot forward control to IdP. Redirecting to UI logout resume URI in order to continue the logout procedure."
      );
      return Optional.of(validatedUiLogoutResumeUri);
   }

   @Override
   public void processSpInitiatedLogoutResponse(URI validatedUiLogoutResumeUri) {
      this.ensureHasFederationSessionWithSloInProgress();

      try {
         ((Saml2SpFederationLibrary)this.library.get()).processSpInitiatedLogoutResponse(validatedUiLogoutResumeUri);
         LOG.info("SAML 2.0 SP-initiated SLO on SP: Done processing response from IdP. SP-initiated SLO successfully completed.");
      } catch (FailureException e) {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, e.getMessage());
      } finally {
         ((Saml2LogoutResumeRedirectService)this.logoutResumeRedirectService.get()).cleanupSessions();
      }
   }

   private void ensureHasFederationSessionWithSloInProgress() {
      Optional<FlowFederationSession> session = ((FlowFederationSessionRepository)this.federationSessionRepository.get()).getFederationSession();
      session.ifPresentOrElse(
         s -> Validate.isTrue(s.isSingleLogoutInProgress()),
         () -> ((FlowFederationSessionRepository)this.federationSessionRepository.get())
            .createNewFederationSession()
            .setSpInitiatedSingleLogoutInProgressOnSp()
      );
   }

   private URI retrieveUiLogoutUri() {
      Optional<URI> resumeUriFromFederationSession = ((Saml2LogoutResumeRedirectService)this.logoutResumeRedirectService.get()).getUiLogoutUriFromSession();
      if (resumeUriFromFederationSession.isPresent()) {
         return resumeUriFromFederationSession.get();
      } else {
         Optional<URI> resumeUriFromRelayState = ((Saml2SpFederationLibrary)this.library.get()).relayStateUriFromLogoutResponse();
         if (resumeUriFromRelayState.isPresent()) {
            return resumeUriFromRelayState.get();
         } else {
            throw new ViolationException("SAML2 SP-initiated SLO: Logout resume URI could not be found.");
         }
      }
   }

   private URI validateUiLogoutResumeUri(String uiLogoutResumeUriParameter) {
      boolean valid = ((Pattern)this.uiLogoutResumeUriPattern.get()).matcher(uiLogoutResumeUriParameter).matches();
      if (!valid) {
         String message = String.format(
            "SAML2 SP-initiated SLO: Validation of location parameter failed. Parameter value %s does not match the validation pattern %s.",
            LogUtil.logSafeQuote(uiLogoutResumeUriParameter),
            LogUtil.logSafeQuote(((Pattern)this.uiLogoutResumeUriPattern.get()).pattern())
         );
         throw new ViolationException(message);
      } else {
         URI uiLogoutResumeUri = UriBuilder.fromUri(uiLogoutResumeUriParameter).build(new Object[0]);
         if (!uiLogoutResumeUri.isAbsolute()) {
            throw new ViolationException("SAML2 SP-initiated SLO: Logout resume URI must be absolute.");
         } else {
            return uiLogoutResumeUri;
         }
      }
   }

   private Saml2SpInitiatedSloServiceImpl.EntityIds resolveEntityIds(FlowFederationSession flowFederationSession) {
      Optional<String> infoKeyString = flowFederationSession.getOnlyValue(AccountUtils.getNameIDInfoKeyAttribute());
      Saml2IdpEntityId idpEntityId = null;
      Saml2SpEntityId hostEntityId = null;
      if (infoKeyString.isPresent()) {
         try {
            NameIDInfoKey nameIdInfoKey = NameIDInfoKey.parse(infoKeyString.get());
            idpEntityId = Saml2IdpEntityId.idpEntityId(nameIdInfoKey.getRemoteEntityID());
            hostEntityId = Saml2SpEntityId.spEntityId(nameIdInfoKey.getHostEntityID());
         } catch (SAML2Exception e) {
            throw new FailureException("Error while parsing Info Key String" + e.getMessage(), e);
         }
      }

      if (idpEntityId == null) {
         throw new FailureException("Unable to determine IDP entity ID during SLO.");
      } else {
         return new Saml2SpInitiatedSloServiceImpl.EntityIds(idpEntityId, hostEntityId);
      }
   }

   private record EntityIds(Saml2IdpEntityId idpEntityId, Saml2SpEntityId hostEntityId) {
   }
}