package com.airlock.iam.oauth2.infrastructure.repository;

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.user.UserId;
import com.airlock.iam.base.misc.exception.Exceptions;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.time.RealTimeProvider;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.core.infrastructure.db.JooqProvider;
import com.airlock.iam.database.infrastructure.DatabaseType;
import com.airlock.iam.database.infrastructure.JooqConfig;
import com.airlock.iam.database.infrastructure.JooqTransaction;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.oauth2.api.domain.model.AuthorizationServerId;
import com.airlock.iam.oauth2.api.domain.model.OAuth2ClientId;
import com.airlock.iam.oauth2.domain.model.session.OAuth2ExtendedSession;
import com.airlock.iam.oauth2.domain.model.session.OAuth2ExtendedSessionRepository;
import com.airlock.iam.oauth2.domain.model.session.OAuth2Session;
import com.airlock.iam.oauth2.domain.model.session.OAuth2SessionClaims;
import com.airlock.iam.oauth2.domain.model.session.OAuth2SessionId;
import com.airlock.iam.oauth2.domain.model.session.OperationStatistics;
import com.google.common.annotations.VisibleForTesting;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.Collection;
import java.util.Optional;
import org.jooq.DSLContext;
import org.jooq.Sequence;
import org.jooq.impl.DSL;
import org.jooq.impl.SQLDataType;

public final class OAuth2ExtendedSessionRepositoryImpl implements OAuth2ExtendedSessionRepository {
   private static final ELogger LOG = ELoggerFactory.getLogger(OAuth2ExtendedSessionRepositoryImpl.class);
   private final Provider<OAuth2SessionRepositorySettings> settings;
   private final Provider<AuthorizationServerId> authorizationServerId;
   private final TimeProvider timeProvider;

   @Inject
   public OAuth2ExtendedSessionRepositoryImpl(Provider<OAuth2SessionRepositorySettings> settings, Provider<AuthorizationServerId> authorizationServerId) {
      this(settings, authorizationServerId, new RealTimeProvider());
   }

   @VisibleForTesting
   OAuth2ExtendedSessionRepositoryImpl(
      Provider<OAuth2SessionRepositorySettings> settings, Provider<AuthorizationServerId> authorizationServerId, TimeProvider timeProvider
   ) {
      this.settings = settings;
      this.authorizationServerId = authorizationServerId;
      this.timeProvider = timeProvider;
   }

   private static boolean isOracleDatabase(JooqConfig jooqConfig) {
      return jooqConfig.dbType().equals(DatabaseType.ORACLE);
   }

   @Override
   public Optional<OAuth2ExtendedSession> findById(OAuth2SessionId oauth2SessionId) {
      return Exceptions.wrapException(
         () -> (Optional<OAuth2ExtendedSession>)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(db -> this.findValidOAuth2ExtendedSession(oauth2SessionId, db)),
         e -> new FailureException(String.format("Error finding OAuth 2.0 session with ID %s", LogUtil.logSafeQuote(oauth2SessionId.asString())), e)
      );
   }

   @Override
   public Optional<OAuth2SessionClaims> getAccessAndIdTokenClaims(OAuth2SessionId oauth2SessionId) {
      return Exceptions.wrapException(
         () -> (Optional<OAuth2SessionClaims>)JooqTransaction.jooqTransaction(this.db()).executeAndReturn(db -> this.findSessionClaims(oauth2SessionId, db)),
         e -> new FailureException(String.format("Error finding OAuth 2.0 session with ID %s", LogUtil.logSafeQuote(oauth2SessionId.asString())), e)
      );
   }

   @Override
   public Collection<OAuth2ExtendedSession> findByUser(UserId userId) {
      return Exceptions.wrapException(
         () -> (Collection<OAuth2ExtendedSession>)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(db -> this.getExtendedSessionFinder(db).findOAuth2ExtendedSessionsByUser(userId)),
         e -> new FailureException(String.format("Error finding OAuth 2.0 sessions for user %s", LogUtil.logSafeQuote(userId)), e)
      );
   }

   @Override
   public Collection<OAuth2ExtendedSession> findByUserAndClient(UserId userId, OAuth2ClientId clientId) {
      return Exceptions.wrapException(
         () -> (Collection<OAuth2ExtendedSession>)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(db -> this.getExtendedSessionFinder(db).findOAuth2ExtendedSessionsByUserAndClient(userId, clientId)),
         e -> new FailureException(
            String.format("Error finding OAuth 2.0 sessions for user %s and client %s", LogUtil.logSafeQuote(userId), LogUtil.logSafeQuote(clientId)), e
         )
      );
   }

   @Override
   public OAuth2ExtendedSession insert(OAuth2ExtendedSession oauth2Session) {
      if (!oauth2Session.authorizationServerId().equals(this.authorizationServerId.get())) {
         throw new ViolationException(
            "Cannot insert a session of authorization server %s using a session repository of different authorization server %s"
               .formatted(oauth2Session.authorizationServerId().id(), ((AuthorizationServerId)this.authorizationServerId.get()).id())
         );
      } else {
         return Exceptions.wrapException(
            () -> (OAuth2ExtendedSession)JooqTransaction.jooqTransaction(this.db()).executeAndReturn(db -> this.insert(oauth2Session, db)),
            e -> new FailureException(String.format("Error inserting OAuth 2.0 session with ID %s", LogUtil.logSafeQuote(oauth2Session.id().asString())), e)
         );
      }
   }

   @Override
   public boolean delete(OAuth2SessionId oauth2SessionId) {
      return Exceptions.wrapException(
         () -> (Boolean)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(
               db -> this.getSessionUpdater(db).deleteOAuth2Session(oauth2SessionId, (AuthorizationServerId)this.authorizationServerId.get()).getSessionCount()
                  == 1
            ),
         e -> wrapOrForwardException(String.format("Error deleting OAuth 2.0 session ID %s", LogUtil.logSafeQuote(oauth2SessionId)), e)
      );
   }

   @Override
   public OperationStatistics delete(Collection<OAuth2SessionId> oauth2SessionIds) {
      return oauth2SessionIds.isEmpty()
         ? OperationStatistics.zero()
         : Exceptions.wrapException(
            () -> (OperationStatistics)JooqTransaction.jooqTransaction(this.db())
               .executeAndReturn(
                  db -> this.getSessionUpdater(db).deleteOAuth2Sessions(oauth2SessionIds, (AuthorizationServerId)this.authorizationServerId.get())
               ),
            e -> wrapOrForwardException(String.format("Error deleting OAuth 2.0 session IDs %s", LogUtil.logSafeQuote(oauth2SessionIds)), e)
         );
   }

   @Override
   public OperationStatistics deleteAllSessionsByUser(UserId userId) {
      return Exceptions.wrapException(
         () -> (OperationStatistics)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(db -> this.getSessionUpdater(db).deleteOAuth2Sessions(userId, (AuthorizationServerId)this.authorizationServerId.get())),
         e -> wrapOrForwardException(String.format("Error deleting OAuth 2.0 sessions for user %s", LogUtil.logSafeQuote(userId.getValue())), e)
      );
   }

   @Override
   public OAuth2ExtendedSession updateCustomAttributes(OAuth2ExtendedSession oauth2Session) {
      return Exceptions.wrapException(
         () -> (OAuth2ExtendedSession)JooqTransaction.jooqTransaction(this.db())
            .executeAndReturn(
               db -> {
                  OAuth2ExtendedSession existingSession = this.findValidOAuth2ExtendedSession(oauth2Session.id(), db)
                     .orElseThrow(() -> new ViolationException("Session to update not found"));
                  this.getSessionUpdater(db).updateCustomAttributesOnOAuth2Session(oauth2Session, existingSession);
                  LOG.debug("Updated custom session attributes");
                  return this.findValidOAuth2ExtendedSession(oauth2Session.id(), db)
                     .orElseThrow(() -> new ViolationException("Session not found after updating"));
               }
            ),
         e -> new FailureException(
            String.format("Error updating OAuth 2.0 session with ID %s: %s", LogUtil.logSafeQuote(oauth2Session.id().asString()), e.getMessage()), e
         )
      );
   }

   private static RuntimeException wrapOrForwardException(String message, Exception originalException) {
      return originalException instanceof ViolationException ? (ViolationException)originalException : new FailureException(message, originalException);
   }

   private OAuth2ExtendedSession insert(OAuth2ExtendedSession oauth2Session, DSLContext db) {
      OAuth2SessionUpdater sessionUpdater = this.getSessionUpdater(db);
      sessionUpdater.insertOAuth2Session(oauth2Session, ((OAuth2SessionRepositorySettings)this.settings.get()).usePersistedClaims());
      return oauth2Session;
   }

   private OAuth2SessionUpdater getSessionUpdater(DSLContext db) {
      return OAuth2SessionUpdater.oAuth2SessionUpdater(db, this.getSequence(), ((OAuth2SessionRepositorySettings)this.settings.get()).usePersistedClaims());
   }

   private OAuth2ExtendedSessionFinder getExtendedSessionFinder(DSLContext db) {
      return OAuth2ExtendedSessionFinder.oAuth2ExtendedSessionFinder(
         db,
         (AuthorizationServerId)this.authorizationServerId.get(),
         ((OAuth2SessionRepositorySettings)this.settings.get()).usePersistedClaims(),
         this.timeProvider
      );
   }

   private Optional<OAuth2ExtendedSession> findValidOAuth2ExtendedSession(OAuth2SessionId oauth2SessionId, DSLContext db) {
      OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult res = this.getExtendedSessionFinder(db).findOAuth2ExtendedSessionById(oauth2SessionId);
      if (res.isValid()) {
         return Optional.of(res.getValidSession());
      }

      if (res.isSessionExpired()) {
         LOG.debug("Cleaning up expired session without active tokens");
         OAuth2SessionUpdater.oAuth2SessionUpdater(db, this.getSequence(), ((OAuth2SessionRepositorySettings)this.settings.get()).usePersistedClaims())
            .deleteOAuth2Sessions(oauth2SessionId);
      }

      return Optional.empty();
   }

   private Optional<OAuth2SessionClaims> findSessionClaims(OAuth2SessionId oauth2SessionId, DSLContext db) {
      return OAuth2SessionFinder.oAuth2SessionFinder(db, ((OAuth2SessionRepositorySettings)this.settings.get()).usePersistedClaims())
         .findOAuth2SessionById(oauth2SessionId, (AuthorizationServerId)this.authorizationServerId.get())
         .map(OAuth2Session::getClaims);
   }

   private Optional<Sequence<Long>> getSequence() {
      return isOracleDatabase(((OAuth2SessionRepositorySettings)this.settings.get()).jooqConfig())
         ? Optional.of(DSL.sequence(DSL.name("oauth2_session_seq"), SQLDataType.BIGINT))
         : Optional.empty();
   }

   private DSLContext db() {
      JooqConfig jooqConfig = ((OAuth2SessionRepositorySettings)this.settings.get()).jooqConfig();
      return JooqProvider.jooqProvider(jooqConfig).db();
   }
}
