package com.airlock.iam.authentication.domain.model.airlock2fa;

import com.airlock.iam.airlock2fa.api.domain.model.Airlock2FAAccountId;
import com.airlock.iam.airlock2fa.api.domain.model.Airlock2FADeviceId;
import com.airlock.iam.authentication.api.domain.model.step.AuthenticationFlowStep;
import com.airlock.iam.authentication.api.domain.model.step.AuthenticationMethodVerifyingFlowStep;
import com.airlock.iam.authentication.api.domain.model.step.AuthenticationStepResult;
import com.airlock.iam.authentication.api.domain.model.step.AuthenticationStepResultBuilder;
import com.airlock.iam.authentication.api.domain.model.step.AuthenticationStepResultFactory;
import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.plugin.LogicForConfig;
import com.airlock.iam.base.api.domain.model.user.UserId;
import com.airlock.iam.common.api.domain.model.authmethod.AuthenticationMethod;
import com.airlock.iam.common.api.domain.model.user.UserIdentity;
import com.airlock.iam.common.domain.model.user.PrincipalInfo;
import com.airlock.iam.core.domain.model.authentication.factor.KnownAuthenticationFactorDetail;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FAFactor;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FAFlowBindingToken;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FALogData;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FAUser;
import com.airlock.iam.factor.domain.model.airlock2fa.activation.Airlock2FAActivationCode;
import com.airlock.iam.factor.domain.model.airlock2fa.futurae.FuturaeErrorResponseException;
import com.airlock.iam.factor.domain.service.airlock2fa.activation.Airlock2FAActivationService;
import com.airlock.iam.factor.domain.service.airlock2fa.activation.Airlock2FAFlowBindingService;
import com.airlock.iam.factor.domain.service.airlock2fa.enrollment.Airlock2FAAuthApiEnrollmentService;
import com.airlock.iam.factor.domain.service.airlock2fa.repository.Airlock2FAUserRepository;
import com.airlock.iam.flow.api.domain.model.session.FlowAttributeRepository;
import com.airlock.iam.flow.api.domain.model.step.StepSession;
import com.airlock.iam.flow.shared.application.configuration.airlock2fa.Airlock2faActivationTrustedSessionBindingStepConfig;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.Airlock2FAActivationFlowBindingSettings;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.Airlock2FAActivationFlowBindingStep;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.Airlock2FAActivationFlowBindingStepSession;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.Airlock2FAActivationStatusStep;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.Airlock2FACurrentDeviceIdFlowAttribute;
import com.airlock.iam.flow.shared.domain.service.PollingService;
import com.airlock.iam.flow.shared.domain.service.airlock2fa.Airlock2FAAuthTokenIdService;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.Optional;
import java.util.function.BiFunction;

@LogicForConfig(Airlock2faActivationTrustedSessionBindingStepConfig.class)
public final class Airlock2FAActivationFlowBindingStepImpl implements Airlock2FAActivationFlowBindingStep, AuthenticationMethodVerifyingFlowStep {
   private static final String FLOW_BINDING_TOKEN_ATTRIBUTE = "flowBindingToken";
   private final Provider<Airlock2FAActivationService> activationService;
   private final Provider<Airlock2FAFlowBindingService> flowBindingService;
   private final Provider<Airlock2FAUserRepository> repository;
   private final Provider<PrincipalInfo> user;
   private final Provider<Airlock2FAActivationFlowBindingSettings> settings;
   private final Provider<AuthenticationMethod> authMethod;
   private final Provider<Airlock2FAActivationFlowBindingStepSession> stepSession;
   private final Provider<Airlock2FAAuthTokenIdService> authTokenIdService;
   private final Provider<FlowAttributeRepository> attributes;
   private final PollingService<Airlock2FAActivationFlowBindingStepSession, AuthenticationStepResult> pollingService;

   @Inject
   public Airlock2FAActivationFlowBindingStepImpl(
      Provider<Airlock2FAActivationService> activationService,
      Provider<Airlock2FAFlowBindingService> flowBindingService,
      Provider<Airlock2FAUserRepository> repository,
      Provider<PrincipalInfo> user,
      Provider<Airlock2FAActivationFlowBindingSettings> settings,
      Provider<AuthenticationMethod> authMethod,
      Provider<Airlock2FAActivationFlowBindingStepSession> stepSession,
      Provider<Airlock2FAAuthTokenIdService> authTokenIdService,
      Provider<FlowAttributeRepository> attributes
   ) {
      this.activationService = activationService;
      this.flowBindingService = flowBindingService;
      this.repository = repository;
      this.user = user;
      this.settings = settings;
      this.authMethod = authMethod;
      this.stepSession = stepSession;
      this.authTokenIdService = authTokenIdService;
      this.attributes = attributes;
      this.pollingService = new PollingService<>(
         Airlock2FAActivationFlowBindingStepSession::getPollingCache, this::pollInternal, s -> this.deviceActivationRequired().build()
      );
   }

   @Override
   public AuthenticationStepResult initialize() {
      return this.stepResultFactory()
         .stay(Airlock2FAActivationFlowBindingStep.Airlock2FAActivationFlowBindingNextAction.AIRLOCK_2FA_DEVICE_ACTIVATION_CODE_REQUIRED)
         .build();
   }

   public AuthenticationStepResult getBindingToken(Airlock2FAActivationCode activationCode) {
      return this.checkEnrollmentStatus(activationCode, (status, accountId) -> {
         return switch (status.getResult()) {
            case PENDING -> this.fetchBindingToken(activationCode);
            case EXPIRED -> this.handleExpired(accountId);
            case SUCCESS -> this.deviceActivationFailed();
         };
      });
   }

   private AuthenticationStepResult fetchBindingToken(Airlock2FAActivationCode activationCode) {
      ((Airlock2FAActivationFlowBindingStepSession)this.stepSession.get()).setActivationCode(activationCode);
      Optional<Airlock2FAFlowBindingToken> token = ((Airlock2FAFlowBindingService)this.flowBindingService.get())
         .tryGetFlowBindingToken(activationCode, ((Airlock2FAActivationFlowBindingSettings)this.settings.get()).flowBindingTokenValidity())
         .peek();
      return token.isPresent()
         ? this.deviceActivationRequired().addAdditionalAttribute("flowBindingToken", token.get().token()).build()
         : this.deviceActivationRequired().build();
   }

   public AuthenticationStepResult pollActivationResult(boolean queryOnly) {
      return this.pollingService.poll((Airlock2FAActivationFlowBindingStepSession)this.stepSession.get(), queryOnly);
   }

   private AuthenticationStepResult pollInternal(Airlock2FAActivationFlowBindingStepSession session) {
      return session.getActivationCode().isEmpty()
         ? this.stepResultFactory().failDueToUnexpectedCall("Binding token must be retrieved before polling.").build()
         : this.checkEnrollmentStatus(session.getActivationCode().orElseThrow(), (status, accountId) -> {
            return switch (status.getResult()) {
               case PENDING -> this.deviceActivationRequired().build();
               case EXPIRED -> this.handleExpired(accountId);
               case SUCCESS -> this.handleSuccess(status.getDeviceId().orElseThrow(), accountId);
            };
         });
   }

   private AuthenticationStepResult checkEnrollmentStatus(
      Airlock2FAActivationCode activationCode,
      BiFunction<Airlock2FAAuthApiEnrollmentService.Airlock2FAEnrollmentStatus, Airlock2FAAccountId, AuthenticationStepResult> handleStatus
   ) {
      Optional<Airlock2FAUser> account = ((Airlock2FAUserRepository)this.repository.get()).findBy(this.userId());
      if (account.isEmpty()) {
         return this.noValidToken();
      }

      Airlock2FAAccountId accountId = account.get().getAccountId();
      return ((Airlock2FAActivationService)this.activationService.get())
         .tryCheckEnrollmentStatus(accountId, activationCode)
         .mapResult(status -> handleStatus.apply(status, accountId))
         .orMapFailure(failure -> {
            if (failure.type() == FuturaeErrorResponseException.FuturaeErrorResponseType.BAD_REQUEST) {
               return this.invalidActivationCode();
            } else {
               throw failure;
            }
         });
   }

   private AuthenticationStepResultBuilder deviceActivationRequired() {
      return this.stepResultFactory().stay(Airlock2FAActivationStatusStep.Airlock2FAActivationStepAction.AIRLOCK_2FA_DEVICE_ACTIVATION_REQUIRED);
   }

   private AuthenticationStepResult invalidActivationCode() {
      return this.stepResultFactory()
         .fail(Airlock2FAActivationFlowBindingStep.Airlock2FAActivationFlowBindingStepError.AIRLOCK_2FA_DEVICE_ACTIVATION_CODE_INVALID)
         .build();
   }

   private AuthenticationStepResult handleExpired(Airlock2FAAccountId accountId) {
      ((Airlock2FAActivationService)this.activationService.get()).handleExpired(Airlock2FALogData.airlock2FALogData(accountId));
      return this.deviceActivationFailed();
   }

   private AuthenticationStepResult deviceActivationFailed() {
      return this.stepResultFactory().fail(Airlock2FAActivationStatusStep.Airlock2FAActivationStepErrors.AIRLOCK_2FA_DEVICE_ACTIVATION_FAILED).build();
   }

   private AuthenticationStepResult handleSuccess(Airlock2FADeviceId deviceId, Airlock2FAAccountId accountId) {
      ((FlowAttributeRepository)this.attributes.get())
         .putOrMergeAttribute(Airlock2FACurrentDeviceIdFlowAttribute.airlock2FACurrentDeviceIdFlowAttribute(deviceId));
      ((Airlock2FAAuthTokenIdService)this.authTokenIdService.get()).handleSuccessfulUsage(Airlock2FAFactor.ONE_TOUCH, accountId, deviceId);
      ((Airlock2FAActivationService)this.activationService.get())
         .handleSuccess(accountId, deviceId, Airlock2FALogData.airlock2FALogData(accountId).withDeviceId(deviceId));
      return this.stepResultFactory().success().build();
   }

   private AuthenticationStepResult noValidToken() {
      return this.stepResultFactory().fail(AuthenticationFlowStep.GeneralAuthenticationStepError.NO_VALID_TOKEN).isFailedFactor(false).build();
   }

   @Override
   public Optional<StepSession> createSession() {
      return Optional.of(Airlock2FAActivationFlowBindingStepSession.emptySession());
   }

   public AuthenticationStepResultFactory stepResultFactory() {
      return AuthenticationStepResultFactory.forFactorDetail(KnownAuthenticationFactorDetail.CRAM_AIRLOCK_2FA_ACTIVATION);
   }

   @Override
   public String getDescription() {
      return "Airlock 2FA Activation Trusted Session Binding Step";
   }

   private UserId userId() {
      return ((PrincipalInfo)this.user.get()).getKnownIdentity().map(UserIdentity::id).orElseThrow(() -> new FailureException("User must be known"));
   }

   @Override
   public AuthenticationMethod getAuthenticationMethod() {
      return (AuthenticationMethod)this.authMethod.get();
   }
}