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.log.LogUtil;
import com.airlock.iam.core.misc.token.KnownTokenTypes;
import com.airlock.iam.database.infrastructure.generated.schema.tables.Oauth2SessionAttributeTable;
import com.airlock.iam.database.infrastructure.generated.schema.tables.Oauth2SessionTable;
import com.airlock.iam.database.infrastructure.generated.schema.tables.TokenTable;
import com.airlock.iam.database.infrastructure.generated.schema.tables.records.Oauth2SessionRecord;
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.domain.model.session.OAuth2CustomSessionAttribute;
import com.airlock.iam.oauth2.domain.model.session.OAuth2ExtendedSession;
import com.airlock.iam.oauth2.domain.model.session.OAuth2Session;
import com.airlock.iam.oauth2.domain.model.session.OAuth2SessionId;
import com.airlock.iam.oauth2.domain.model.session.OperationStatistics;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import org.jooq.BatchBindStep;
import org.jooq.Condition;
import org.jooq.DSLContext;
import org.jooq.InsertSetMoreStep;
import org.jooq.Sequence;
import org.jooq.impl.DSL;

public final class OAuth2SessionUpdater {
   private static final ELogger LOG = ELoggerFactory.getLogger(OAuth2SessionUpdater.class);
   private final DSLContext db;
   private final Optional<Sequence<Long>> sequence;
   private final OAuth2SessionFinder sessionFinder;

   public static OAuth2SessionUpdater oAuth2SessionUpdater(DSLContext db, Optional<Sequence<Long>> sequence, boolean usePersistedClaims) {
      return new OAuth2SessionUpdater(db, sequence, OAuth2SessionFinder.oAuth2SessionFinder(db, usePersistedClaims));
   }

   public void insertOAuth2Session(OAuth2ExtendedSession oauth2Session, boolean persistClaims) {
      long oauth2SessionPk = this.insertIntoOAuth2SessionEntity(oauth2Session, persistClaims)
         .returningResult(Oauth2SessionTable.OAUTH2_SESSION.PK)
         .fetchOptional()
         .map(record -> (Long)record.get(Oauth2SessionTable.OAUTH2_SESSION.PK))
         .orElseThrow(
            () -> new FailureException(
               String.format("Error fetching PK from inserted OAuth2 session with ID %s", LogUtil.logSafeQuote(oauth2Session.id().asString()))
            )
         );
      if (!oauth2Session.customAttributes().isEmpty()) {
         this.insertAttributeEntities(oauth2SessionPk, oauth2Session.customAttributes());
      }
   }

   public void updateCustomAttributesOnOAuth2Session(OAuth2ExtendedSession newOAuth2Session, OAuth2ExtendedSession storedOAuth2Session) {
      List<OAuth2CustomSessionAttribute> newAttributes = newOAuth2Session.customAttributes()
         .stream()
         .filter(attr -> storedOAuth2Session.customAttributes().stream().noneMatch(a -> a.getName().equals(attr.getName())))
         .toList();
      List<OAuth2CustomSessionAttribute> updatedAttributes = newOAuth2Session.customAttributes()
         .stream()
         .filter(attr -> storedOAuth2Session.customAttributes().stream().anyMatch(a -> a.getName().equals(attr.getName())))
         .toList();
      List<OAuth2CustomSessionAttribute> deletedAttributes = storedOAuth2Session.customAttributes()
         .stream()
         .filter(attr -> newOAuth2Session.customAttributes().stream().noneMatch(a -> a.getName().equals(attr.getName())))
         .toList();
      long sessionPk = this.getSessionPk(storedOAuth2Session);
      this.deleteAttributes(sessionPk, deletedAttributes);
      this.updateAttributes(sessionPk, updatedAttributes);
      this.insertAttributeEntities(sessionPk, newAttributes);
   }

   private void deleteAttributes(long sessionPk, List<OAuth2CustomSessionAttribute> deletedAttributes) {
      if (!deletedAttributes.isEmpty()) {
         this.db
            .deleteFrom(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE)
            .where(OAuth2SessionRepositoryConditions.matchingOAuth2SessionPk(sessionPk))
            .and(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.NAME.in(deletedAttributes.stream().map(OAuth2CustomSessionAttribute::getName).toList()))
            .execute();
      }
   }

   private void updateAttributes(long sessionPk, List<OAuth2CustomSessionAttribute> updatedAttributes) {
      if (!updatedAttributes.isEmpty()) {
         BatchBindStep batch = this.db
            .batch(
               this.db
                  .update(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE)
                  .set(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.CONTENT, DSL.param(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.CONTENT))
                  .where(
                     Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE
                        .OAUTH2_SESSION_PK
                        .eq(DSL.param(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.OAUTH2_SESSION_PK))
                        .and(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.NAME.eq(DSL.param(Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.NAME)))
                  )
            );

         for (OAuth2CustomSessionAttribute attr : updatedAttributes) {
            batch = batch.bind(new Object[]{attr.getValue(), sessionPk, attr.getName()});
         }

         batch.execute();
      }
   }

   private long getSessionPk(OAuth2ExtendedSession session) {
      return (Long)this.db
         .select(Oauth2SessionTable.OAUTH2_SESSION.PK)
         .from(Oauth2SessionTable.OAUTH2_SESSION)
         .where(Oauth2SessionTable.OAUTH2_SESSION.ID.eq(session.id().asString()))
         .fetchOptional(Oauth2SessionTable.OAUTH2_SESSION.PK)
         .orElseThrow(() -> new ViolationException("Could not find primary key of session"));
   }

   public OperationStatistics deleteOAuth2Sessions(UserId userId) {
      return this.deleteOAuth2Sessions(this.sessionFinder.findAllOAuth2SessionsHavingTokens(userId));
   }

   public OperationStatistics deleteOAuth2Sessions(UserId userId, AuthorizationServerId authorizationServerId) {
      return this.deleteOAuth2Sessions(this.sessionFinder.findAllOAuth2SessionsHavingTokens(userId, authorizationServerId));
   }

   public OperationStatistics deleteOAuth2Session(OAuth2SessionId sessionId, AuthorizationServerId authorizationServerId) {
      Optional<OAuth2Session> session = this.sessionFinder.findOAuth2SessionById(sessionId, authorizationServerId);
      if (session.isEmpty()) {
         LOG.debug("Session %s not found for AS %s".formatted(LogUtil.logSafeQuote(sessionId.asString()), LogUtil.logSafeQuote(authorizationServerId.id())));
         return OperationStatistics.zero();
      } else {
         return this.deleteOAuth2Sessions(sessionId);
      }
   }

   public OperationStatistics deleteOAuth2Sessions(Collection<OAuth2SessionId> sessionIds, AuthorizationServerId authorizationServerId) {
      OAuth2SessionId[] filteredSessionIds = this.filterSessionIdsByAs(sessionIds.stream().map(OAuth2SessionId::getId).toList(), authorizationServerId)
         .stream()
         .map(OAuth2SessionId::oAuth2SessionId)
         .toArray(OAuth2SessionId[]::new);
      return this.deleteOAuth2Sessions(filteredSessionIds);
   }

   public OperationStatistics deleteOAuth2Sessions(OAuth2SessionId... sessionIds) {
      List<String> sessionIdsToDelete = Arrays.stream(sessionIds).map(OAuth2SessionId::getId).collect(Collectors.toList());
      if (sessionIdsToDelete.isEmpty()) {
         return OperationStatistics.zero();
      }

      int deletedTokens = this.deleteOAuth2TokensForSessions(sessionIdsToDelete);
      int deletedSessions = this.db.deleteFrom(Oauth2SessionTable.OAUTH2_SESSION).where(Oauth2SessionTable.OAUTH2_SESSION.ID.in(sessionIdsToDelete)).execute();
      LOG.debug("Deleted %d OAuth2 session(s) and %d OAuth2 token(s)".formatted(deletedSessions, deletedTokens));
      return OperationStatistics.builder().tokenCount(deletedTokens).sessionCount(deletedSessions).build();
   }

   private Collection<String> filterSessionIdsByAs(Collection<String> sessionIds, AuthorizationServerId authorizationServerId) {
      return sessionIds.isEmpty()
         ? Collections.emptyList()
         : this.db
            .selectDistinct(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID)
            .from(TokenTable.TOKEN)
            .where(
               new Condition[]{
                  TokenTable.TOKEN.TYPE.eq(DSL.inline(KnownTokenTypes.OAUTH2.getId())),
                  OAuth2SessionRepositoryFields.TOKEN_SESSION_ID.in(sessionIds),
                  OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID.eq(authorizationServerId.id())
               }
            )
            .fetch(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID);
   }

   private int deleteOAuth2TokensForSessions(List<String> sessionIdsToDelete) {
      return this.db
         .deleteFrom(TokenTable.TOKEN)
         .where(
            new Condition[]{
               TokenTable.TOKEN.TYPE.eq(DSL.inline(KnownTokenTypes.OAUTH2.getId())), OAuth2SessionRepositoryFields.TOKEN_SESSION_ID.in(sessionIdsToDelete)
            }
         )
         .execute();
   }

   private InsertSetMoreStep<Oauth2SessionRecord> insertIntoOAuth2SessionEntity(OAuth2ExtendedSession oauth2Session, boolean persistClaims) {
      InsertSetMoreStep<Oauth2SessionRecord> query = this.db
         .insertInto(Oauth2SessionTable.OAUTH2_SESSION)
         .set(Oauth2SessionTable.OAUTH2_SESSION.ID, oauth2Session.id().asString());
      if (persistClaims) {
         query = query.set(Oauth2SessionTable.OAUTH2_SESSION.CLAIMS, oauth2Session.claims().asJsonString());
      }

      return this.addSequence(query);
   }

   private InsertSetMoreStep<Oauth2SessionRecord> addSequence(InsertSetMoreStep<Oauth2SessionRecord> query) {
      this.sequence.map(seq -> query.set(Oauth2SessionTable.OAUTH2_SESSION.PK, seq.nextval()));
      return query;
   }

   private void insertAttributeEntities(long oauth2SessionPk, Collection<OAuth2CustomSessionAttribute> attributes) {
      if (!attributes.isEmpty()) {
         List<String> names = attributes.stream().map(OAuth2CustomSessionAttribute::getName).toList();
         int result = this.db
            .insertInto(
               Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE,
               Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.OAUTH2_SESSION_PK,
               Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.NAME,
               Oauth2SessionAttributeTable.OAUTH2_SESSION_ATTRIBUTE.CONTENT
            )
            .valuesOfRows(attributes.stream().map(attribute -> DSL.row(oauth2SessionPk, attribute.getName(), attribute.getValue())).toList())
            .execute();
         if (result != attributes.size()) {
            throw new FailureException(
               String.format(
                  "Error inserting %d OAuth2 session attributes with names %s: %d records inserted",
                  names.size(),
                  names.stream().map(LogUtil::logSafeQuote).collect(Collectors.joining(", ")),
                  result
               )
            );
         }
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   private OAuth2SessionUpdater(DSLContext db, Optional<Sequence<Long>> sequence, OAuth2SessionFinder sessionFinder) {
      this.db = db;
      this.sequence = sequence;
      this.sessionFinder = sessionFinder;
   }
}