package com.airlock.iam.authentication.application.service.saml2.sp;

import com.airlock.iam.authentication.application.service.ApplicationAccessService;
import com.airlock.iam.authentication.application.service.targetapp.TargetApplicationRepository;
import com.airlock.iam.authentication.domain.model.targetapp.TargetApplication;
import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.result.Result;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.flow.api.domain.model.flow.ApplicationId;
import com.airlock.iam.flow.api.domain.model.flow.FlowResult;
import com.airlock.iam.flow.api.domain.model.session.FlowSession;
import com.airlock.iam.flow.api.domain.model.step.GenericStepError;
import com.airlock.iam.flow.api.domain.model.step.StepResultFactory;
import com.airlock.iam.flow.domain.model.flow.OptionalFlow;
import com.airlock.iam.flow.domain.model.result.FlowResultImpl;
import com.airlock.iam.flow.shared.api.domain.model.session.UserSession;
import com.airlock.iam.login.application.configuration.targetapp.ApplicationIdConfig;
import com.airlock.iam.login.domain.model.targetapp.TargetApplicationAccessForwardURI;
import com.airlock.iam.saml2.application.configuration.sp.Saml2FlowSpConfig;
import com.airlock.iam.saml2.domain.model.Saml2RelayState;
import com.airlock.iam.saml2.domain.model.logging.Saml2LogTypes;
import com.airlock.iam.saml2.domain.model.sp.Saml2AuthnRequestId;
import com.airlock.iam.saml2.domain.model.sp.Saml2InResponseTo;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpAssertionAttributes;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpEntityId;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpInitiatedSsoSessionAttribute;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpSsoAssertionData;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpSsoContinuationId;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpSsoMetadata;
import com.airlock.iam.saml2.domain.model.sp.step.Saml2IdpInitiatedSsoFlowAttribute;
import com.airlock.iam.saml2.domain.model.sp.step.Saml2SpSsoSessionAttribute;
import com.airlock.iam.saml2.domain.model.sp.step.Saml2UserIdentifyingStep;
import com.airlock.iam.saml2.domain.service.Saml2FlowSpConfigService;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.net.URI;
import java.util.Optional;
import lombok.NonNull;

public final class Saml2SpSsoAccessService {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2SpSsoAccessService.class);
   private final Provider<StepResultFactory> stepResultFactory;
   private final Provider<OptionalFlow> currentFlow;
   private final Provider<UserSession> userSession;
   private final Provider<ApplicationAccessService> applicationAccessService;
   private final Provider<Saml2FlowSpConfigService> spConfigService;
   private final Provider<TargetApplicationRepository> applicationRepository;

   @Inject
   Saml2SpSsoAccessService(
      @NonNull Provider<StepResultFactory> stepResultFactory,
      @NonNull Provider<OptionalFlow> currentFlow,
      @NonNull Provider<UserSession> userSession,
      @NonNull Provider<ApplicationAccessService> applicationAccessService,
      @NonNull Provider<Saml2FlowSpConfigService> spConfigService,
      @NonNull Provider<TargetApplicationRepository> applicationRepository
   ) {
      if (stepResultFactory == null) {
         throw new NullPointerException("stepResultFactory is marked non-null but is null");
      }

      if (currentFlow == null) {
         throw new NullPointerException("currentFlow is marked non-null but is null");
      }

      if (userSession == null) {
         throw new NullPointerException("userSession is marked non-null but is null");
      }

      if (applicationAccessService == null) {
         throw new NullPointerException("applicationAccessService is marked non-null but is null");
      }

      if (spConfigService == null) {
         throw new NullPointerException("spConfigService is marked non-null but is null");
      }

      if (applicationRepository == null) {
         throw new NullPointerException("applicationRepository is marked non-null but is null");
      }

      this.stepResultFactory = stepResultFactory;
      this.currentFlow = currentFlow;
      this.userSession = userSession;
      this.applicationAccessService = applicationAccessService;
      this.spConfigService = spConfigService;
      this.applicationRepository = applicationRepository;
   }

   public FlowResult access(Saml2SpSsoContinuationId continuationId) {
      boolean removeContinuationId = true;

      try {
         Optional<Saml2SpSsoAssertionData> assertionData = this.getAssertionDataFromUserSession(continuationId);
         if (assertionData.isEmpty()) {
            return this.failDueToUnexpectedCall("No SSO session attribute found in user session for continuationId " + LogUtil.logSafeQuote(continuationId));
         }

         Result<Saml2SpSsoAccessService.SsoInitiator, FlowResult> determineInitiatorResult = this.determineSsoInitiator(assertionData.get());
         if (determineInitiatorResult.isFailure()) {
            if (this.shouldStopRunningFlowAndTryAgain(determineInitiatorResult.getFailure())) {
               removeContinuationId = false;
            }

            return determineInitiatorResult.getFailure();
         } else {
            FlowResult accessResult = switch ((Saml2SpSsoAccessService.SsoInitiator)determineInitiatorResult.get()) {
               case SP_INITIATED_SSO -> this.accessForSpInitiatedSso(assertionData.get().assertionAttributes());
               case IDP_INITIATED_SSO -> this.accessForIdpInitiatedSso(assertionData.get().ssoMetadata(), assertionData.get().assertionAttributes());
            };
            if (this.shouldStopRunningFlowAndTryAgain(determineInitiatorResult, accessResult)) {
               removeContinuationId = false;
            }

            return accessResult;
         }
      } finally {
         if (removeContinuationId) {
            ((UserSession)this.userSession.get()).removeUserSessionAttribute(Saml2SpSsoSessionAttribute.class);
         }
      }
   }

   private Optional<Saml2SpSsoAssertionData> getAssertionDataFromUserSession(Saml2SpSsoContinuationId continuationId) {
      return ((UserSession)this.userSession.get())
         .findUserSessionAttribute(Saml2SpSsoSessionAttribute.class)
         .filter(s -> s.continuationId().equals(continuationId))
         .map(Saml2SpSsoSessionAttribute::assertionData);
   }

   private boolean shouldStopRunningFlowAndTryAgain(FlowResult failure) {
      return failure.isFailed()
         && failure.error().isPresent()
         && failure.error().get().equals(GenericStepError.UNEXPECTED_CALL)
         && failure.nextStep().isPresent();
   }

   private boolean shouldStopRunningFlowAndTryAgain(Result<Saml2SpSsoAccessService.SsoInitiator, FlowResult> determineInitiatorResult, FlowResult accessResult) {
      return determineInitiatorResult.get().equals(Saml2SpSsoAccessService.SsoInitiator.IDP_INITIATED_SSO)
         && this.shouldStopRunningFlowAndTryAgain(accessResult);
   }

   private Result<Saml2SpSsoAccessService.SsoInitiator, FlowResult> determineSsoInitiator(Saml2SpSsoAssertionData ssoAttributes) {
      return ((UserSession)this.userSession.get()).findUserSessionAttribute(Saml2SpInitiatedSsoSessionAttribute.class).isPresent()
         ? this.validateSpInitiatedSso(ssoAttributes)
         : this.validateIdpInitiatedSso(ssoAttributes);
   }

   private Result<Saml2SpSsoAccessService.SsoInitiator, FlowResult> validateSpInitiatedSso(Saml2SpSsoAssertionData ssoAttributes) {
      if (((OptionalFlow)this.currentFlow.get()).isEmpty()) {
         return Result.failure(this.failDueToUnexpectedCall("SP-Initiated SSO: Expected to find a running authentication flow but none was found"));
      } else {
         Saml2SpInitiatedSsoSessionAttribute attribute = ((UserSession)this.userSession.get())
            .findUserSessionAttribute(Saml2SpInitiatedSsoSessionAttribute.class)
            .orElseThrow();
         Saml2AuthnRequestId authnRequestId = attribute.authnRequestId()
            .orElseThrow(() -> new FailureException("AuthnRequest ID not present in user session."));
         ((UserSession)this.userSession.get()).addUserSessionAttribute(attribute.withoutAuthnRequestId());
         Optional<Saml2InResponseTo> inResponseTo = ssoAttributes.ssoMetadata().inResponseTo();
         if (inResponseTo.isEmpty()) {
            LOG.warn(
               Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE,
               "SP-Initiated SSO: Response from IDP does not contain 'InResponseTo' attribute. Aborting local flow and switching to IDP-initiated SSO."
            );
            return Result.failure(FlowResultImpl.withError(((OptionalFlow)this.currentFlow.get()).get().lastResult(), GenericStepError.UNEXPECTED_CALL));
         } else if (!this.hasRequestIdMatchingInResponseTo(authnRequestId, inResponseTo)) {
            logWrongInResponseTo(authnRequestId, inResponseTo);
            return Result.failure(this.failDueToUnexpectedCall("SP-Initiated SSO: SAMLResponse is not a response to the previous authentication request"));
         } else {
            return !attribute.spEntityId().equals(ssoAttributes.assertionAttributes().spEntityId())
               ? Result.failure(
                  this.failDueToUnexpectedCall(
                     "SP-Initiated SSO: expected SAMLResponse for SP entity ID %s but got %s"
                        .formatted(LogUtil.logSafeQuote(attribute.spEntityId()), LogUtil.logSafeQuote(ssoAttributes.assertionAttributes().spEntityId()))
                  )
               )
               : Result.result(Saml2SpSsoAccessService.SsoInitiator.SP_INITIATED_SSO);
         }
      }
   }

   private Result<Saml2SpSsoAccessService.SsoInitiator, FlowResult> validateIdpInitiatedSso(Saml2SpSsoAssertionData ssoAttributes) {
      if (ssoAttributes.ssoMetadata().inResponseTo().isPresent()) {
         return Result.failure(this.failDueToUnexpectedCall("IdP-Initiated SSO: SAMLResponse is a response to an SP-initiated AuthnRequest"));
      }

      Saml2SpEntityId spEntityId = ssoAttributes.ssoMetadata().spEntityId();
      return ((Saml2FlowSpConfigService)this.spConfigService.get()).findFlowSpConfig(spEntityId).isEmpty()
         ? Result.failure(this.failDueToUnexpectedCall("IdP-Initiated SSO: No flow SP was configured for entity ID " + LogUtil.logSafeQuote(spEntityId)))
         : Result.result(Saml2SpSsoAccessService.SsoInitiator.IDP_INITIATED_SSO);
   }

   private boolean hasRequestIdMatchingInResponseTo(Saml2AuthnRequestId authnRequestId, Optional<Saml2InResponseTo> inResponseTo) {
      return inResponseTo.<Boolean>map(id -> authnRequestId.value().equals(id.value())).orElse(false);
   }

   private FlowResult accessForSpInitiatedSso(Saml2SpAssertionAttributes attributes) {
      return ((OptionalFlow)this.currentFlow.get()).get().execute(step -> step.identify(attributes), Saml2UserIdentifyingStep.class);
   }

   private FlowResult accessForIdpInitiatedSso(Saml2SpSsoMetadata metadata, Saml2SpAssertionAttributes attributes) {
      Saml2SpSsoAccessService.TargetApplicationInfo flowApplicationId = this.determineFlowApplicationId(metadata);
      return ((ApplicationAccessService)this.applicationAccessService.get())
         .startNewFlowForApplication(
            flowApplicationId.applicationId(),
            session -> {
               this.addAssertionAttributesToFlowSession(session, attributes);
               this.raiseIdpInitiatedSsoIdentifyRequiredRedFlag(session);
               flowApplicationId.locationFromRelayState()
                  .ifPresent(
                     locationFromRelayState -> this.addRelayStateLocationToFlowSession(
                        session, TargetApplicationAccessForwardURI.targetApplicationAccessForwardURI(locationFromRelayState)
                     )
                  );
            }
         );
   }

   private Saml2SpSsoAccessService.TargetApplicationInfo determineFlowApplicationId(Saml2SpSsoMetadata metadata) {
      Optional<Saml2SpSsoAccessService.TargetApplicationInfo> relayStateApplicationId = this.findApplicationByRelayState(metadata.relayState());
      if (relayStateApplicationId.isPresent()) {
         this.logSelectedRelayStateApplication(relayStateApplicationId.get().applicationId(), metadata.relayState().orElseThrow());
         return relayStateApplicationId.get();
      } else {
         Optional<ApplicationId> spDefaultApplicationId = this.getSpDefaultFlowApplicationId(metadata.spEntityId());
         if (spDefaultApplicationId.isPresent()) {
            this.logSelectedSpDefaultApplication(spDefaultApplicationId.get(), metadata.spEntityId());
            return new Saml2SpSsoAccessService.TargetApplicationInfo(spDefaultApplicationId.get(), metadata.relayState().flatMap(Saml2RelayState::toUri));
         } else {
            ApplicationId defaultApplicationId = ((TargetApplicationRepository)this.applicationRepository.get()).getDefaultApplication().getApplicationId();
            this.logSelectedGlobalDefaultApplication(defaultApplicationId, metadata.spEntityId());
            return new Saml2SpSsoAccessService.TargetApplicationInfo(defaultApplicationId, metadata.relayState().flatMap(Saml2RelayState::toUri));
         }
      }
   }

   private void addAssertionAttributesToFlowSession(FlowSession session, Saml2SpAssertionAttributes attributes) {
      session.flowAttributes().putOrMergeAttribute(new Saml2IdpInitiatedSsoFlowAttribute(attributes));
   }

   private void addRelayStateLocationToFlowSession(FlowSession session, TargetApplicationAccessForwardURI location) {
      session.flowAttributes().putOrMergeAttribute(location);
   }

   private void raiseIdpInitiatedSsoIdentifyRequiredRedFlag(FlowSession session) {
      session.raiseRedFlag(Saml2UserIdentifyingStep.SAML2_IDP_INITIATED_SSO_USER_IDENTIFY_REQUIRED_RED_FLAG);
   }

   private Optional<Saml2SpSsoAccessService.TargetApplicationInfo> findApplicationByRelayState(Optional<Saml2RelayState> relayState) {
      if (relayState.isEmpty()) {
         return Optional.empty();
      } else {
         Optional<URI> location = relayState.get().toUri();
         if (location.isEmpty()) {
            this.logCouldNotConvertRelayStateToUri(relayState.get());
            return Optional.empty();
         } else {
            Optional<ApplicationId> applicationId = ((TargetApplicationRepository)this.applicationRepository.get())
               .getByForwardURI(location.get())
               .map(TargetApplication::getApplicationId);
            if (applicationId.isEmpty()) {
               this.logCouldNotFindApplicationForRelayState(relayState.get());
               return Optional.empty();
            } else {
               return Optional.of(new Saml2SpSsoAccessService.TargetApplicationInfo(applicationId.get(), location));
            }
         }
      }
   }

   private Optional<ApplicationId> getSpDefaultFlowApplicationId(Saml2SpEntityId spEntityId) {
      return ((Saml2FlowSpConfigService)this.spConfigService.get())
         .findFlowSpConfig(spEntityId)
         .flatMap(Saml2FlowSpConfig::getDefaultFlowApplicationId)
         .map(ApplicationIdConfig::getId);
   }

   private static void logWrongInResponseTo(Saml2AuthnRequestId authnRequestId, Optional<Saml2InResponseTo> inResponseTo) {
      LOG.debug(
         "SP-Initiated SSO: Invalid SAMLResponse. The element 'InResponseTo' with value %s does not match the expected authentication request ID %s"
            .formatted(LogUtil.logSafeQuote(inResponseTo.orElse(null)), LogUtil.logSafeQuote(authnRequestId))
      );
   }

   private void logSelectedRelayStateApplication(ApplicationId applicationId, Saml2RelayState relayState) {
      LOG.info(
         "IdP-initiated SSO: Starting new flow for application %s, determined from relay state %s"
            .formatted(LogUtil.logSafeQuote(applicationId), LogUtil.logSafeQuote(relayState))
      );
   }

   private void logSelectedSpDefaultApplication(ApplicationId applicationId, Saml2SpEntityId entityId) {
      LOG.info(
         "IdP-initiated SSO: Starting new flow for application %s, configured as default application for SP %s"
            .formatted(LogUtil.logSafeQuote(applicationId), LogUtil.logSafeQuote(entityId))
      );
   }

   private void logSelectedGlobalDefaultApplication(ApplicationId applicationId, Saml2SpEntityId entityId) {
      LOG.info(
         "IdP-initiated SSO: Starting new flow for global default application %s with SP %s"
            .formatted(LogUtil.logSafeQuote(applicationId), LogUtil.logSafeQuote(entityId))
      );
   }

   private void logCouldNotFindApplicationForRelayState(Saml2RelayState relayState) {
      LOG.debug("IdP-initiated SSO: No target application found matching relay state %s".formatted(LogUtil.logSafeQuote(relayState)));
   }

   private void logCouldNotConvertRelayStateToUri(Saml2RelayState relayState) {
      LOG.debug("IdP-initiated SSO: Provided relay state %s is not a valid URI, discarding provided relay state".formatted(LogUtil.logSafeQuote(relayState)));
   }

   private FlowResult failDueToUnexpectedCall(String reason) {
      return FlowResultImpl.of(((StepResultFactory)this.stepResultFactory.get()).failDueToUnexpectedCall(reason).build());
   }

   @SuppressFBWarnings(justification = "generated code")
   public static Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder builder() {
      return new Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class Saml2SpSsoAccessServiceBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private Provider<StepResultFactory> stepResultFactory;
      @SuppressFBWarnings(justification = "generated code")
      private Provider<OptionalFlow> currentFlow;
      @SuppressFBWarnings(justification = "generated code")
      private Provider<UserSession> userSession;
      @SuppressFBWarnings(justification = "generated code")
      private Provider<ApplicationAccessService> applicationAccessService;
      @SuppressFBWarnings(justification = "generated code")
      private Provider<Saml2FlowSpConfigService> spConfigService;
      @SuppressFBWarnings(justification = "generated code")
      private Provider<TargetApplicationRepository> applicationRepository;

      @SuppressFBWarnings(justification = "generated code")
      Saml2SpSsoAccessServiceBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder stepResultFactory(@NonNull Provider<StepResultFactory> stepResultFactory) {
         if (stepResultFactory == null) {
            throw new NullPointerException("stepResultFactory is marked non-null but is null");
         }

         this.stepResultFactory = stepResultFactory;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder currentFlow(@NonNull Provider<OptionalFlow> currentFlow) {
         if (currentFlow == null) {
            throw new NullPointerException("currentFlow is marked non-null but is null");
         }

         this.currentFlow = currentFlow;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder userSession(@NonNull Provider<UserSession> userSession) {
         if (userSession == null) {
            throw new NullPointerException("userSession is marked non-null but is null");
         }

         this.userSession = userSession;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder applicationAccessService(
         @NonNull Provider<ApplicationAccessService> applicationAccessService
      ) {
         if (applicationAccessService == null) {
            throw new NullPointerException("applicationAccessService is marked non-null but is null");
         }

         this.applicationAccessService = applicationAccessService;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder spConfigService(@NonNull Provider<Saml2FlowSpConfigService> spConfigService) {
         if (spConfigService == null) {
            throw new NullPointerException("spConfigService is marked non-null but is null");
         }

         this.spConfigService = spConfigService;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder applicationRepository(@NonNull Provider<TargetApplicationRepository> applicationRepository) {
         if (applicationRepository == null) {
            throw new NullPointerException("applicationRepository is marked non-null but is null");
         }

         this.applicationRepository = applicationRepository;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public Saml2SpSsoAccessService build() {
         return new Saml2SpSsoAccessService(
            this.stepResultFactory, this.currentFlow, this.userSession, this.applicationAccessService, this.spConfigService, this.applicationRepository
         );
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "Saml2SpSsoAccessService.Saml2SpSsoAccessServiceBuilder(stepResultFactory="
            + this.stepResultFactory
            + ", currentFlow="
            + this.currentFlow
            + ", userSession="
            + this.userSession
            + ", applicationAccessService="
            + this.applicationAccessService
            + ", spConfigService="
            + this.spConfigService
            + ", applicationRepository="
            + this.applicationRepository
            + ")";
      }
   }

   enum SsoInitiator {
      SP_INITIATED_SSO,
      IDP_INITIATED_SSO;
   }

   private record TargetApplicationInfo(ApplicationId applicationId, Optional<URI> locationFromRelayState) {
   }
}