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

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.infrastructure.db.DbInstantConverter;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.core.misc.token.KnownTokenTypes;
import com.airlock.iam.core.misc.util.ticket.KeyMultiValue;
import com.airlock.iam.database.infrastructure.generated.schema.tables.Oauth2SessionTable;
import com.airlock.iam.database.infrastructure.generated.schema.tables.TokenAssignmentTable;
import com.airlock.iam.database.infrastructure.generated.schema.tables.TokenTable;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.elog.api.LogTypes;
import com.airlock.iam.oauth2.api.domain.model.AuthorizationServerId;
import com.airlock.iam.oauth2.api.domain.model.OAuth2ClientId;
import com.airlock.iam.oauth2.api.domain.model.OAuth2Scope;
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.OAuth2SessionClaims;
import com.airlock.iam.oauth2.domain.model.session.OAuth2SessionId;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.text.ParseException;
import java.time.Instant;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.commons.lang3.StringUtils;
import org.jooq.Condition;
import org.jooq.DSLContext;
import org.jooq.Record;
import org.jooq.Result;
import org.jooq.impl.DSL;

public final class OAuth2ExtendedSessionFinder {
   private static final ELogger LOG = ELoggerFactory.getLogger(OAuth2ExtendedSessionFinder.class);
   private final DSLContext db;
   private final OAuth2SessionMaterializer materializer;
   private final OAuth2SessionFinder sessionFinder;
   private final AuthorizationServerId authorizationServerId;
   private final TimeProvider timeProvider;

   public static OAuth2ExtendedSessionFinder oAuth2ExtendedSessionFinder(
      DSLContext db, AuthorizationServerId authorizationServerId, boolean usePersistedClaims, TimeProvider timeProvider
   ) {
      return new OAuth2ExtendedSessionFinder(
         db,
         OAuth2SessionMaterializer.oAuth2SessionMaterializer(),
         OAuth2SessionFinder.oAuth2SessionFinder(db, usePersistedClaims),
         authorizationServerId,
         timeProvider
      );
   }

   public OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult findOAuth2ExtendedSessionById(OAuth2SessionId oauth2SessionId) {
      Optional<OAuth2Session> session = this.sessionFinder.findOAuth2SessionById(oauth2SessionId, this.authorizationServerId);
      return session.isEmpty()
         ? OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.notFoundResult()
         : this.findValidTokenAndCreateOAuth2Session(oauth2SessionId, session.get().getClaims(), session.get().getCustomAttributes())
            .map(OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult::validResult)
            .orElseGet(OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult::expiredResult);
   }

   public Collection<OAuth2ExtendedSession> findOAuth2ExtendedSessionsByUser(UserId userId) {
      Result<? extends Record> tokenRecords = this.findAllValidOAuth2TokensOfUser(userId, null);
      return this.findOAuth2ExtendedSessions(tokenRecords);
   }

   public Collection<OAuth2ExtendedSession> findOAuth2ExtendedSessionsByUserAndClient(UserId userId, OAuth2ClientId clientId) {
      Result<? extends Record> tokenRecords = this.findAllValidOAuth2TokensOfUser(userId, clientId);
      return this.findOAuth2ExtendedSessions(tokenRecords);
   }

   private Collection<OAuth2ExtendedSession> findOAuth2ExtendedSessions(Result<? extends Record> tokenRecords) {
      List<String> sessionIds = tokenRecords.stream()
         .map(r -> (String)r.get(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID))
         .distinct()
         .collect(Collectors.toList());
      Map<OAuth2SessionId, OAuth2Session> attributesBySessionId = this.getSessionsWithAttributes(sessionIds);
      Map<OAuth2SessionId, OAuth2ExtendedSession> oAuth2Sessions = tokenRecords.stream()
         .map(
            tokenRecord -> this.mapToOAuth2Session(
               tokenRecord,
               attributesBySessionId.get(OAuth2SessionId.oAuth2SessionId((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID))).getClaims(),
               attributesBySessionId.get(OAuth2SessionId.oAuth2SessionId((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID)))
                  .getCustomAttributes()
            )
         )
         .collect(Collectors.toMap(OAuth2ExtendedSession::id, session -> (OAuth2ExtendedSession)session, (first, second) -> first));
      return oAuth2Sessions.values();
   }

   private Optional<OAuth2ExtendedSession> findValidTokenAndCreateOAuth2Session(
      OAuth2SessionId sessionId, OAuth2SessionClaims claims, Set<OAuth2CustomSessionAttribute> attributes
   ) {
      return this.db
         .select(
            OAuth2SessionRepositoryFields.TOKEN_SESSION_ID,
            OAuth2SessionRepositoryFields.TOKEN_CLIENT_ID,
            OAuth2SessionRepositoryFields.TOKEN_SCOPES,
            OAuth2SessionRepositoryFields.TOKEN_AUTH_TIMESTAMP,
            OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID
         )
         .from(TokenTable.TOKEN)
         .where(
            new Condition[]{
               OAuth2SessionRepositoryFields.TOKEN_SESSION_ID.eq(sessionId.getId()),
               TokenTable.TOKEN.TYPE.eq(DSL.inline(KnownTokenTypes.OAUTH2.getId())),
               TokenTable.TOKEN.ACTIVE.eq((byte)1),
               TokenTable.TOKEN
                  .VALIDITY_RANGE_UPPER
                  .isNull()
                  .or(TokenTable.TOKEN.VALIDITY_RANGE_UPPER.greaterOrEqual(DbInstantConverter.toNanoSecondsSinceEpoch(this.timeProvider.getCurrentInstant()))),
               OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID.eq(this.authorizationServerId.id())
            }
         )
         .limit(1)
         .fetchOptional()
         .map(record -> this.mapToOAuth2Session(record, claims, attributes));
   }

   private Result<? extends Record> findAllValidOAuth2TokensOfUser(UserId userId, OAuth2ClientId clientId) {
      return this.db
         .select(
            OAuth2SessionRepositoryFields.TOKEN_SESSION_ID,
            OAuth2SessionRepositoryFields.TOKEN_CLIENT_ID,
            OAuth2SessionRepositoryFields.TOKEN_SCOPES,
            OAuth2SessionRepositoryFields.TOKEN_AUTH_TIMESTAMP,
            OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID
         )
         .from(TokenAssignmentTable.TOKEN_ASSIGNMENT)
         .join(TokenTable.TOKEN)
         .on(TokenTable.TOKEN.TOKEN_ID.eq(TokenAssignmentTable.TOKEN_ASSIGNMENT.TA_TOKEN_ID))
         .where(
            new Condition[]{
               TokenAssignmentTable.TOKEN_ASSIGNMENT.TA_USER.eq(userId.getValue()),
               TokenTable.TOKEN.TYPE.eq(DSL.inline(KnownTokenTypes.OAUTH2.getId())),
               TokenTable.TOKEN.ACTIVE.eq((byte)1),
               TokenTable.TOKEN
                  .VALIDITY_RANGE_UPPER
                  .isNull()
                  .or(TokenTable.TOKEN.VALIDITY_RANGE_UPPER.greaterOrEqual(DbInstantConverter.toNanoSecondsSinceEpoch(this.timeProvider.getCurrentInstant()))),
               OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID.eq(this.authorizationServerId.id()),
               clientIdCondition(clientId)
            }
         )
         .orderBy(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID)
         .fetch();
   }

   private static Condition clientIdCondition(OAuth2ClientId clientId) {
      return Optional.ofNullable(clientId)
         .map(OAuth2ClientId::id)
         .<Condition>map(OAuth2SessionRepositoryFields.TOKEN_CLIENT_ID::eq)
         .orElse(DSL.trueCondition());
   }

   private Map<OAuth2SessionId, OAuth2Session> getSessionsWithAttributes(List<String> sessionIds) {
      Result<Record> records = this.sessionFinder.fetchOAuth2SessionsWhere(Oauth2SessionTable.OAUTH2_SESSION.ID.in(sessionIds));
      return this.materializer.materializeOAuth2Sessions(records).stream().collect(Collectors.toMap(OAuth2Session::getId, session -> (OAuth2Session)session));
   }

   private OAuth2ExtendedSession mapToOAuth2Session(Record tokenRecord, OAuth2SessionClaims claims, Set<OAuth2CustomSessionAttribute> attributes) {
      return OAuth2ExtendedSession.builder()
         .id(OAuth2SessionId.oAuth2SessionId((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_SESSION_ID)))
         .clientId(OAuth2ClientId.oAuth2ClientId((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_CLIENT_ID)))
         .scopes(this.tokenScopes((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_SCOPES)))
         .authorizationServerId(
            AuthorizationServerId.authorizationServerId((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_AUTHORIZATION_SERVER_ID))
         )
         .authenticationTimestamp(Instant.ofEpochMilli(Long.parseLong((String)tokenRecord.get(OAuth2SessionRepositoryFields.TOKEN_AUTH_TIMESTAMP))))
         .customAttributes(attributes)
         .claims(claims)
         .build();
   }

   private Set<OAuth2Scope> tokenScopes(String scopesString) {
      if (StringUtils.isNotBlank(scopesString)) {
         try {
            return Stream.of(KeyMultiValue.parseFromString(scopesString, "UTF-8").getValues()).map(OAuth2Scope::of).collect(Collectors.toSet());
         } catch (ParseException e) {
            LOG.warn(LogTypes.RESOURCE_NOT_FOUND, "Could not parse scope from generic token. Proceeded with empty scope.");
         }
      }

      return Set.of();
   }

   @SuppressFBWarnings(justification = "generated code")
   private OAuth2ExtendedSessionFinder(
      DSLContext db,
      OAuth2SessionMaterializer materializer,
      OAuth2SessionFinder sessionFinder,
      AuthorizationServerId authorizationServerId,
      TimeProvider timeProvider
   ) {
      this.db = db;
      this.materializer = materializer;
      this.sessionFinder = sessionFinder;
      this.authorizationServerId = authorizationServerId;
      this.timeProvider = timeProvider;
   }

   public static class OAuth2ExtendedSessionResult {
      private final Optional<OAuth2ExtendedSession> session;
      private final OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState resultState;

      public OAuth2ExtendedSession getValidSession() {
         return ViolationException.checkPresent(this.session, "Session not present; state is " + this.resultState);
      }

      public boolean isSessionExpired() {
         return this.resultState == OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState.SESSION_EXPIRED;
      }

      public boolean isValid() {
         return this.resultState == OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState.VALID;
      }

      static OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult validResult(OAuth2ExtendedSession session) {
         return new OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult(
            Optional.of(session), OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState.VALID
         );
      }

      static OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult notFoundResult() {
         return new OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult(
            Optional.empty(), OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState.SESSION_NOT_FOUND
         );
      }

      static OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult expiredResult() {
         return new OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult(
            Optional.empty(), OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState.SESSION_EXPIRED
         );
      }

      @SuppressFBWarnings(justification = "generated code")
      private OAuth2ExtendedSessionResult(
         Optional<OAuth2ExtendedSession> session, OAuth2ExtendedSessionFinder.OAuth2ExtendedSessionResult.ResultState resultState
      ) {
         this.session = session;
         this.resultState = resultState;
      }

      enum ResultState {
         VALID,
         SESSION_NOT_FOUND,
         SESSION_EXPIRED;
      }
   }
}