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

import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.application.service.idp.Saml2SsoContinuationService;
import com.airlock.iam.saml2.domain.model.Saml2MetaManager;
import com.airlock.iam.saml2.domain.model.idp.Saml2FederationSessionIndex;
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.Saml2IdpSsoContinuationId;
import com.airlock.iam.saml2.domain.model.idp.Saml2RequestId;
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.session.Saml2PropagationAttributes;
import com.airlock.iam.saml2.domain.service.Saml2LibraryIdpCacheService;
import com.airlock.iam.saml2.infrastructure.service.idp.Saml2IdpFederationLibrary;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.Optional;

public final class Saml2SsoContinuationServiceImpl implements Saml2SsoContinuationService {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2SsoContinuationServiceImpl.class);
   private final Provider<FlowFederationSessionRepository> repository;
   private final Provider<Saml2IdpFederationLibrary> library;
   private final Provider<Saml2MetaManager> metaAliasFactory;
   private final Provider<TimeProvider> timeProvider;
   private final Provider<Saml2LibraryIdpCacheService> cacheServiceProvider;

   @Inject
   Saml2SsoContinuationServiceImpl(
      Provider<FlowFederationSessionRepository> repository,
      Provider<Saml2IdpFederationLibrary> library,
      Provider<Saml2MetaManager> metaAliasFactory,
      Provider<TimeProvider> timeProvider,
      Provider<Saml2LibraryIdpCacheService> cacheServiceProvider
   ) {
      this.repository = repository;
      this.library = library;
      this.metaAliasFactory = metaAliasFactory;
      this.timeProvider = timeProvider;
      this.cacheServiceProvider = cacheServiceProvider;
   }

   @Override
   public void idpInitiatedSsoFederate(Saml2IdpSsoContinuationId ssoContinuationId) {
      try {
         FlowFederationSession session = this.validateSession(
            ((FlowFederationSessionRepository)this.repository.get())
               .getFederationSession()
               .orElseThrow(() -> new ViolationException("Missing federation session"))
         );
         Saml2PropagationAttributes propagationAttributes = this.validateContinuationAttribute(session.getPropagationAttributes(), ssoContinuationId);
         Saml2IdpMetaAlias idpMetaAlias = ((Saml2MetaManager)this.metaAliasFactory.get()).idpMetaAlias(propagationAttributes.getIdpEntityId());
         Saml2IdpInitiatedSsoContinuationParameters parameters = new Saml2IdpInitiatedSsoContinuationParameters(
            propagationAttributes.getSpEntityId(), idpMetaAlias, propagationAttributes.getRelayState()
         );
         ((Saml2IdpFederationLibrary)this.library.get()).doIdpInitiatedSsoFederate(parameters);
      } finally {
         ((FlowFederationSessionRepository)this.repository.get()).getFederationSession().ifPresent(FlowFederationSession::resetPropagationAttributes);
      }
   }

   @Override
   public void spInitiatedSsoFederate(Saml2IdpSsoContinuationId ssoContinuationId, Saml2RequestId requestId, Saml2IdpMetaAlias idpMetaAlias) {
      try {
         FlowFederationSession session = this.validateSession(
            ((FlowFederationSessionRepository)this.repository.get())
               .getFederationSession()
               .orElseThrow(() -> new ViolationException("Missing federation session"))
         );
         this.validateSpInitiatedParameters(this.validateContinuationAttribute(session.getPropagationAttributes(), ssoContinuationId), requestId, idpMetaAlias);
         ((Saml2LibraryIdpCacheService)this.cacheServiceProvider.get()).removeCachedAuthnContext(requestId);
         ((Saml2IdpFederationLibrary)this.library.get()).doSpInitiatedSsoFederate();
      } finally {
         ((FlowFederationSessionRepository)this.repository.get()).getFederationSession().ifPresent(FlowFederationSession::resetPropagationAttributes);
      }
   }

   private FlowFederationSession validateSession(FlowFederationSession session) {
      if (session.getAuthenticatedUserId().isEmpty()) {
         throw new ViolationException(String.format("Invalid federation session. Authenticated user missing in %s", session));
      }

      if (session.isSingleLogoutInProgress()) {
         throw new ViolationException(String.format("Invalid federation session. Single logout in progress for %s. ", session));
      }

      this.resetSessionIndexIfNotCachedAnymore(session);
      return session;
   }

   private void resetSessionIndexIfNotCachedAnymore(FlowFederationSession session) {
      Optional<Saml2FederationSessionIndex> sessionIndex = session.getSessionIndex();
      if (sessionIndex.isPresent() && !((Saml2LibraryIdpCacheService)this.cacheServiceProvider.get()).isSessionIndexCached(sessionIndex.get())) {
         LOG.debug("Federation Session's SessionIndex not cached anymore.");
         session.clearSessionIndex();
      }
   }

   private Saml2PropagationAttributes validateContinuationAttribute(
      Optional<Saml2PropagationAttributes> propagationAttributes, Saml2IdpSsoContinuationId continuationId
   ) {
      Saml2PropagationAttributes attributes = propagationAttributes.orElseThrow(
         () -> new ViolationException("Missing SAML2 propagation attributes; no SAML2 identity propagation is currently active.")
      );
      if (attributes.isExpired(((TimeProvider)this.timeProvider.get()).getCurrentInstant())) {
         throw new ViolationException("Propagation attributes are no longer valid");
      } else if (!attributes.getSsoContinuationId().equals(continuationId)) {
         throw new ViolationException(
            String.format(
               "Unexpected SSO continuation ID. Expected SSO continuation ID %s but got %s",
               LogUtil.logSafeQuote(attributes.getSsoContinuationId()),
               LogUtil.logSafeQuote(continuationId)
            )
         );
      } else {
         return attributes;
      }
   }

   private void validateSpInitiatedParameters(Saml2PropagationAttributes attributes, Saml2RequestId requestId, Saml2IdpMetaAlias metaAlias) {
      Saml2RequestId expectedRequestId = attributes.getSpInitiatedRequestId()
         .orElseThrow(() -> new ViolationException("No SP-Initiated SSO parameters present."));
      if (!expectedRequestId.equals(requestId)) {
         throw new ViolationException(
            String.format("Unexpected request ID. Expected request ID %s but got %s", LogUtil.logSafeQuote(expectedRequestId), LogUtil.logSafeQuote(requestId))
         );
      }

      if (!((Saml2MetaManager)this.metaAliasFactory.get()).idpMetaAlias(attributes.getIdpEntityId()).equals(metaAlias)) {
         throw new ViolationException(
            String.format(
               "Unexpected meta alias. Meta alias %s does not belong to IdP %s",
               LogUtil.logSafeQuote(metaAlias),
               LogUtil.logSafeQuote(attributes.getIdpEntityId())
            )
         );
      }

      if (((Saml2LibraryIdpCacheService)this.cacheServiceProvider.get()).getAuthnRequestByRequestId(requestId).isEmpty()) {
         throw new ViolationException("Provided request ID does not reference an AuthnRequest in the cache (anymore).");
      }
   }
}