package com.airlock.iam.flow.shared.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.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.exception.ViolationException;
import com.airlock.iam.base.api.domain.model.plugin.LogicForConfig;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.common.api.domain.model.authmethod.AuthenticationMethod;
import com.airlock.iam.core.domain.model.authentication.factor.KnownAuthenticationFactorDetail;
import com.airlock.iam.core.domain.model.authentication.method.KnownAuthenticationMethods;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FAFactor;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FALogData;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FARegisteredDevice;
import com.airlock.iam.factor.domain.model.airlock2fa.Airlock2FAShortLivedOnlineQrCodeSettings;
import com.airlock.iam.factor.domain.model.airlock2fa.activation.Airlock2FAActivationSession;
import com.airlock.iam.factor.domain.model.airlock2fa.activation.Airlock2FAActivationSettings;
import com.airlock.iam.factor.domain.service.airlock2fa.activation.Airlock2FAActivationService;
import com.airlock.iam.factor.domain.service.airlock2fa.enrollment.Airlock2FAAuthApiEnrollmentService;
import com.airlock.iam.flow.api.domain.model.session.FlowAttributeRepository;
import com.airlock.iam.flow.api.domain.model.step.GenericStepError;
import com.airlock.iam.flow.api.domain.model.step.StepSession;
import com.airlock.iam.flow.shared.application.configuration.airlock2fa.GenericAirlock2FAActivationStepConfig;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.activation.Airlock2FAActivationChallenge;
import com.airlock.iam.flow.shared.domain.model.airlock2fa.activation.Airlock2FAActivationStepSession;
import com.airlock.iam.flow.shared.domain.service.PollingService;
import com.airlock.iam.flow.shared.domain.service.airlock2fa.Airlock2FAAuthTokenIdService;
import com.airlock.iam.flow.types.api.domain.model.session.KnownMultiFlowType;
import com.airlock.iam.flow.types.api.domain.model.session.MultiFlowType;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.time.Duration;
import java.time.Instant;
import java.util.Optional;

@LogicForConfig(GenericAirlock2FAActivationStepConfig.class)
public final class Airlock2FAActivationStepImpl implements Airlock2FAActivationStep, AuthenticationFlowStep {
   private static final String FLOW_BINDING_TOKEN_ATTRIBUTE = "flowBindingToken";
   private static final String ACTIVATION_CODE_SHORT_ATTRIBUTE = "activationCodeShort";
   private final Provider<Airlock2FAActivationService> activationService;
   private final Provider<Airlock2FAAuthTokenIdService> authTokenIdService;
   private final Provider<Airlock2FAActivationStepSession> stepSession;
   private final Provider<Airlock2FAActivationStepSettings> stepSettings;
   private final Provider<FlowAttributeRepository> attributes;
   private final Provider<Airlock2FAActivationStepSkipCheck> skipCheck;
   private final Provider<MultiFlowType> multiFlowType;
   private final PollingService<Airlock2FAActivationStepSession, AuthenticationStepResult> pollingService;
   private final Provider<TimeProvider> timeProvider;

   @Inject
   Airlock2FAActivationStepImpl(
      Provider<Airlock2FAActivationService> activationService,
      Provider<Airlock2FAAuthTokenIdService> authTokenIdService,
      Provider<Airlock2FAActivationStepSession> stepSession,
      Provider<Airlock2FAActivationStepSettings> stepSettings,
      Provider<FlowAttributeRepository> attributes,
      Provider<Airlock2FAActivationStepSkipCheck> skipCheck,
      Provider<MultiFlowType> multiFlowType,
      Provider<TimeProvider> timeProvider
   ) {
      this.activationService = activationService;
      this.authTokenIdService = authTokenIdService;
      this.stepSession = stepSession;
      this.stepSettings = stepSettings;
      this.attributes = attributes;
      this.skipCheck = skipCheck;
      this.multiFlowType = multiFlowType;
      this.pollingService = new PollingService<>(Airlock2FAActivationStepSession::getPollingCache, this::pollInternal, this::pollingRequiredResult);
      this.timeProvider = timeProvider;
   }

   @Override
   public AuthenticationStepResult initialize() {
      Optional<Airlock2FARegisteredDevice> registered = ((Airlock2FAActivationStepSkipCheck)this.skipCheck.get()).shouldSkipActivation();
      if (registered.isPresent()) {
         Airlock2FAAccountId accountId = registered.get().accountId();
         Airlock2FADeviceId deviceId = registered.get().deviceId();
         ((Airlock2FAActivationService)this.activationService.get()).handleSkipped(Airlock2FALogData.airlock2FALogData(accountId).withDeviceId(deviceId));
         this.completeActivation(accountId, deviceId);
         return this.stepResultFactory().success().build();
      }

      Airlock2FAActivationSettings.Airlock2FAActivationSettingsBuilder activationSettingsBuilder = ((Airlock2FAActivationStepSettings)this.stepSettings.get())
         .activationSettings()
         .toBuilder();
      if (((Airlock2FAActivationStepSettings)this.stepSettings.get()).shortLivedQrCodeSettings().isPresent()) {
         activationSettingsBuilder.provideActivationCodeShortEnabled(true);
      }

      Airlock2FAActivationSession activationSession = ((Airlock2FAActivationService)this.activationService.get())
         .startActivation(activationSettingsBuilder.build(), ((Airlock2FAActivationStepSettings)this.stepSettings.get()).flowBindingTokenValidity());
      this.stepSession().setActivationSession(activationSession);
      return this.addAdditionalAttributes(
            this.stepResultFactory().stay(Airlock2FAActivationStatusStep.Airlock2FAActivationStepAction.AIRLOCK_2FA_DEVICE_ACTIVATION_REQUIRED),
            this.stepSession()
         )
         .build();
   }

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

   public AuthenticationStepResult retrieveActivationChallenge() {
      if (((Airlock2FAActivationStepSettings)this.stepSettings.get()).shortLivedQrCodeSettings().isPresent()) {
         return this.stepResultFactory()
            .failDueToUnexpectedCall("Deprecated challenge retrieval endpoint does not support Short-Lived Online QR Codes.")
            .build();
      } else if (this.stepSession().getPollingCache().hasResult()) {
         return this.stepResultFactory().failDueToUnexpectedCall("Call is not allowed after the final polling result is ready.").build();
      } else if (this.stepSession().getActivationSession().isEmpty()) {
         throw new FailureException("No Airlock 2FA device activation ongoing.");
      } else {
         return this.stepResultFactory()
            .stay(Airlock2FAActivationStep.Airlock2FAActivationStepAction.AIRLOCK_2FA_DEVICE_ACTIVATION_POLLING_REQUIRED)
            .resultData(this.asStepResultData(this.stepSession().getActivationSession().get()))
            .build();
      }
   }

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

   private AuthenticationStepResult pollInternal(Airlock2FAActivationStepSession session) {
      if (session.getActivationSession().isEmpty()) {
         return this.stepResultFactory().fail(GenericStepError.UNEXPECTED_CALL).build();
      }

      Airlock2FAActivationSession activationSession = session.getActivationSession().get();
      Airlock2FAAuthApiEnrollmentService.Airlock2FAEnrollmentStatus enrollmentStatus = ((Airlock2FAActivationService)this.activationService.get())
         .checkEnrollmentStatus(activationSession.getAccountId(), activationSession.getActivationCode());
      return this.processEnrollmentStatus(activationSession.getAccountId(), enrollmentStatus);
   }

   private Airlock2FAActivationStepSession stepSession() {
      return (Airlock2FAActivationStepSession)this.stepSession.get();
   }

   private AuthenticationStepResult processEnrollmentStatus(
      Airlock2FAAccountId accountId, Airlock2FAAuthApiEnrollmentService.Airlock2FAEnrollmentStatus enrollmentStatus
   ) {
      return switch (enrollmentStatus.getResult()) {
         case PENDING -> this.pollingRequiredResult(this.stepSession());
         case SUCCESS -> {
            this.completeActivation(accountId, enrollmentStatus.getDeviceId().orElseThrow());
            yield this.stepResultFactory().success().build();
         }
         case EXPIRED -> {
            ((Airlock2FAActivationService)this.activationService.get()).handleExpired(this.logData());
            yield this.stepResultFactory()
               .fail(Airlock2FAActivationStatusStep.Airlock2FAActivationStepErrors.AIRLOCK_2FA_DEVICE_ACTIVATION_FAILED)
               .isFailedFactor(false)
               .build();
         }
      };
   }

   private AuthenticationStepResult pollingRequiredResult(Airlock2FAActivationStepSession session) {
      return this.addAdditionalAttributes(
            this.stepResultFactory().stay(Airlock2FAActivationStep.Airlock2FAActivationStepAction.AIRLOCK_2FA_DEVICE_ACTIVATION_POLLING_REQUIRED), session
         )
         .build();
   }

   private void completeActivation(Airlock2FAAccountId accountId, Airlock2FADeviceId deviceId) {
      ((FlowAttributeRepository)this.attributes.get())
         .putOrMergeAttribute(Airlock2FACurrentDeviceIdFlowAttribute.airlock2FACurrentDeviceIdFlowAttribute(deviceId));
      if (this.multiFlowType.get() != KnownMultiFlowType.USER_SELF_REG) {
         ((Airlock2FAAuthTokenIdService)this.authTokenIdService.get()).handleSuccessfulUsage(Airlock2FAFactor.ONE_TOUCH, accountId, deviceId);
      }

      ((Airlock2FAActivationService)this.activationService.get())
         .handleSuccess(
            accountId,
            deviceId,
            Airlock2FALogData.airlock2FALogData(accountId)
               .withDeviceId(deviceId)
               .withEnrollmentIdIfPresent(this.stepSession().getActivationSession().map(Airlock2FAActivationSession::getEnrollmentId))
         );
   }

   @Override
   public AuthenticationMethod getAuthenticationMethod() {
      return KnownAuthenticationMethods.AIRLOCK_2FA;
   }

   @Override
   public String getDescription() {
      return "Airlock2FA Activation Step";
   }

   private Airlock2FAActivationChallenge asStepResultData(Airlock2FAActivationSession session) {
      return Airlock2FAActivationChallenge.airlock2faActivationChallenge(session.getAppDeviceActivationUrl(), session.getQrCode().asBufferedImage());
   }

   private AuthenticationStepResultBuilder addAdditionalAttributes(AuthenticationStepResultBuilder result, Airlock2FAActivationStepSession session) {
      Airlock2FAActivationSession activationSession = session.getActivationSession().orElseThrow();
      activationSession.getFlowBindingToken().ifPresent(token -> result.addAdditionalAttribute("flowBindingToken", token.token()));
      Optional<Airlock2FAShortLivedOnlineQrCodeSettings> shortLivedQrCodeSettings = ((Airlock2FAActivationStepSettings)this.stepSettings.get())
         .shortLivedQrCodeSettings();
      if (shortLivedQrCodeSettings.isPresent()) {
         Duration enrollmentTimeout = ((Airlock2FAActivationStepSettings)this.stepSettings.get()).enrollmentTimeout();
         if (!enrollmentTimeout.equals(shortLivedQrCodeSettings.get().sessionTimeout())) {
            throw new ViolationException("Short-lived Online QR Code timeout and enrollment timeout do not match.");
         } else {
            return this.addEnrollmentExchangeTokenAttributes(result, activationSession, enrollmentTimeout, shortLivedQrCodeSettings.get().validityOverlap());
         }
      } else {
         activationSession.getActivationCodeShort().ifPresent(code -> result.addAdditionalAttribute("activationCodeShort", code.getValue()));
         return StepResultHelper.addEnrollmentQrCodeChallenge(result, activationSession);
      }
   }

   private AuthenticationStepResultBuilder addEnrollmentExchangeTokenAttributes(
      AuthenticationStepResultBuilder result, Airlock2FAActivationSession activationSession, Duration sessionTimeout, Duration overlapDuration
   ) {
      Instant currentTime = ((TimeProvider)this.timeProvider.get()).getCurrentInstant();
      Optional<Instant> startTime = Optional.of(activationSession.getSessionStartTime());
      return StepResultHelper.hasSessionTimeout(startTime, currentTime, sessionTimeout)
         ? this.stepResultFactory().fail(Airlock2FAActivationStatusStep.Airlock2FAActivationStepErrors.AIRLOCK_2FA_DEVICE_ACTIVATION_FAILED)
         : StepResultHelper.addEnrollmentShortLivedQrCodeChallenge(result, activationSession.getEnrollmentExchangeTokens(), currentTime, overlapDuration);
   }

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

   private Airlock2FALogData logData() {
      return this.stepSession().logData();
   }
}