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

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.common.api.domain.model.claims.Claim;
import com.airlock.iam.common.api.domain.model.user.PersistentUser;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.oauth2.api.application.configuration.claims.OAuth2UsernameTransformer;
import com.airlock.iam.oauth2.api.domain.model.claims.OAuth2TokenClaimData;
import com.airlock.iam.oauth2.api.domain.model.claims.OAuth2TokenInformation;
import com.airlock.iam.oauth2.domain.model.claims.DistributedClaims;
import com.airlock.iam.oauth2.domain.model.claims.IdTokenUserClaimExtractor;
import com.airlock.iam.oauth2.domain.model.logging.OAuth2LogTypes;
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.token.OAuth2ClientIdTokenExtension;
import com.airlock.iam.oauth2.domain.model.token.OAuth2ClientIdTokenExtensionProvider;
import com.airlock.iam.oauth2.domain.model.validation.OAuth2RequestHandlingException;
import com.airlock.iam.oauth2.domain.service.OAuth2GrantsSessionService;
import com.airlock.iam.oauth2.domain.service.OpenIdConnectCodeFlowIdTokenClaimsService;
import com.airlock.iam.oauth2.domain.service.OpenIdConnectHybridFlowIdTokenClaimsService;
import com.airlock.iam.oauth2.domain.service.persistence.OAuth2UserLoader;
import com.google.common.collect.ImmutableList;
import com.nimbusds.jwt.JWTClaimsSet.Builder;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

public class OpenIdConnectIdTokenClaimsServiceImpl implements OpenIdConnectCodeFlowIdTokenClaimsService, OpenIdConnectHybridFlowIdTokenClaimsService {
   private static final ELogger LOG = ELoggerFactory.getLogger(OpenIdConnectIdTokenClaimsServiceImpl.class);
   private static final Collection<String> ID_TOKEN_FIXED_CLAIMS = ImmutableList.builder().add("auth_time").add("nonce").add("acr").build();
   private final Provider<IdTokenUserClaimExtractor> defaultUserClaimExtractors;
   private final Provider<OpenIdConnectIdTokenDistributedClaimProvider> distributedClaimsProvider;
   private final CommonTokenClaimService commonTokenClaimService;
   private final Provider<OAuth2UsernameTransformer> usernameTransformer;
   private final Provider<OAuth2UserLoader> userLoader;
   private final Provider<OAuth2GrantsSessionService> sessionService;
   private final Provider<OpenIdConnectIdTokenClaimsServiceSettings> settings;
   private final Provider<OAuth2ClientIdTokenExtensionProvider> clientExtensionProvider;

   @Inject
   public OpenIdConnectIdTokenClaimsServiceImpl(
      Provider<IdTokenUserClaimExtractor> defaultUserClaimExtractors,
      Provider<OpenIdConnectIdTokenDistributedClaimProvider> distributedClaimsProvider,
      Provider<OAuth2UsernameTransformer> usernameTransformer,
      Provider<OAuth2UserLoader> userLoader,
      Provider<OAuth2GrantsSessionService> sessionService,
      Provider<OpenIdConnectIdTokenClaimsServiceSettings> settings,
      Provider<OAuth2ClientIdTokenExtensionProvider> clientExtensionProvider
   ) {
      this.defaultUserClaimExtractors = defaultUserClaimExtractors;
      this.distributedClaimsProvider = distributedClaimsProvider;
      this.commonTokenClaimService = new CommonTokenClaimService(ID_TOKEN_FIXED_CLAIMS);
      this.usernameTransformer = usernameTransformer;
      this.userLoader = userLoader;
      this.sessionService = sessionService;
      this.settings = settings;
      this.clientExtensionProvider = clientExtensionProvider;
   }

   @Override
   public List<Claim<?>> evaluateIdTokenClaims(OAuth2TokenInformation idTokenInformation) {
      Builder claims = new Builder();
      PersistentUser user = ((OAuth2UserLoader)this.userLoader.get())
         .getValidUser(idTokenInformation.username())
         .orElseThrow(() -> new OAuth2RequestHandlingException("Token not assigned to valid user: " + LogUtil.quote(idTokenInformation.username())));
      claims.subject(((OAuth2UsernameTransformer)this.usernameTransformer.get()).apply(user));
      Optional<OAuth2ClientIdTokenExtension> clientIdTokenExtension = ((OAuth2ClientIdTokenExtensionProvider)this.clientExtensionProvider.get())
         .extensionPerClientId(idTokenInformation.clientId());
      this.addCustomClaims(claims, user, idTokenInformation, clientIdTokenExtension);
      this.addDistributedClaims(claims, clientIdTokenExtension);
      return claims.build().toJSONObject().entrySet().stream().map(OAuth2AccessTokenClaimsServiceImpl::wrap).toList();
   }

   private void addDistributedClaims(Builder jwtClaimsSet, Optional<OAuth2ClientIdTokenExtension> clientIdTokenExtension) {
      Optional<DistributedClaims> clientSpecificClaims = clientIdTokenExtension.flatMap(OAuth2ClientIdTokenExtension::idTokenDistributedClaims);
      if (clientSpecificClaims.isPresent()) {
         DistributedClaims combinedDistributedClaims = ((OpenIdConnectIdTokenDistributedClaimProvider)this.distributedClaimsProvider.get())
            .distributedClaims()
            .orElse(DistributedClaims.emptyDistributedClaims())
            .combineWith(clientSpecificClaims.get());
         jwtClaimsSet.claim("_claim_names", combinedDistributedClaims.getClaimNames());
         jwtClaimsSet.claim("_claim_sources", combinedDistributedClaims.getClaimSources());
      } else {
         ((OpenIdConnectIdTokenDistributedClaimProvider)this.distributedClaimsProvider.get()).distributedClaims().ifPresent(distributedClaims -> {
            jwtClaimsSet.claim("_claim_names", distributedClaims.getClaimNames());
            jwtClaimsSet.claim("_claim_sources", distributedClaims.getClaimSources());
         });
      }
   }

   @Override
   public List<Claim<?>> loadOrEvaluateIdTokenClaims(OAuth2TokenInformation tokenInformation) {
      return ((OpenIdConnectIdTokenClaimsServiceSettings)this.settings.get()).isUsePersistedClaims()
         ? this.loadIdTokenClaims(OAuth2SessionId.oAuth2SessionId(tokenInformation.getSessionId().getId()))
         : this.evaluateIdTokenClaims(tokenInformation);
   }

   private List<Claim<?>> loadIdTokenClaims(OAuth2SessionId sessionId) {
      return ((OAuth2GrantsSessionService)this.sessionService.get())
         .getSessionClaims(sessionId)
         .map(OAuth2SessionClaims::getIdTokenSnapshot)
         .map(OAuth2SessionClaims::toClaims)
         .orElse(List.of());
   }

   private void addCustomClaims(
      Builder jwtClaimsSet, PersistentUser persistentUser, OAuth2TokenClaimData tokenInfo, Optional<OAuth2ClientIdTokenExtension> clientIdTokenExtension
   ) {
      try {
         List<Claim<?>> defaultClaims = this.commonTokenClaimService
            .extractCustomClaimValues(((IdTokenUserClaimExtractor)this.defaultUserClaimExtractors.get()).userClaimExtractors(), persistentUser, tokenInfo);
         List<Claim<?>> clientSpecificClaims = clientIdTokenExtension.<List<Claim<?>>>map(
               extension -> this.commonTokenClaimService.extractCustomClaimValues(extension.customClaims(), persistentUser, tokenInfo)
            )
            .orElse(Collections.emptyList());
         this.commonTokenClaimService.combineDefaultAndClientSpecificClaims(jwtClaimsSet, defaultClaims, clientSpecificClaims);
      } catch (RuntimeException e) {
         LOG.warn(OAuth2LogTypes.OAUTH2_MESSAGE, "Exception occurred while retrieving ID Token custom claim: " + e.getMessage(), e);
         throw new OAuth2RequestHandlingException(e);
      }
   }
}