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

import com.airlock.iam.common.api.domain.model.claims.Claim;
import com.airlock.iam.common.domain.model.claims.StringClaim;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.oauth2.domain.model.client.OAuth2Client;
import com.airlock.iam.oauth2.domain.model.logger.IntegrationLoggerFactory;
import com.airlock.iam.oauth2.domain.model.logging.OAuth2LogTypes;
import com.airlock.iam.oauth2.domain.model.token.IdToken;
import com.airlock.iam.oauth2.domain.model.token.IdTokenBuilder;
import com.airlock.iam.oauth2.domain.model.token.OpenIDConnectIDTokenSigner;
import com.airlock.iam.oauth2.domain.model.token.OpenIdConnectIdTokenData;
import com.airlock.iam.oauth2.domain.model.validation.OAuth2RequestHandlingException;
import com.airlock.iam.oauth2.domain.service.codegrant.OAuth2AuthorizationCode;
import com.airlock.iam.oauth2.domain.service.token.OpenIdConnectSignedIdTokenFactory;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.oauth2.sdk.AuthorizationCode;
import com.nimbusds.oauth2.sdk.token.TypelessAccessToken;
import com.nimbusds.openid.connect.sdk.claims.AccessTokenHash;
import com.nimbusds.openid.connect.sdk.claims.CodeHash;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

public final class OpenIdConnectSignedIdTokenFactoryImpl implements OpenIdConnectSignedIdTokenFactory {
   private final Provider<OpenIDConnectIDTokenSigner> idTokenSigner;
   private final Provider<IntegrationLoggerFactory> loggerFactory;
   private final Provider<OpenIdConnectSignedIdTokenFactorySettings> settings;

   @Inject
   public OpenIdConnectSignedIdTokenFactoryImpl(
      Provider<IntegrationLoggerFactory> loggerFactory,
      Provider<OpenIDConnectIDTokenSigner> idTokenSigner,
      Provider<OpenIdConnectSignedIdTokenFactorySettings> settings
   ) {
      this.idTokenSigner = idTokenSigner;
      this.loggerFactory = loggerFactory;
      this.settings = settings;
   }

   @Override
   public String getSignedIdToken(OpenIdConnectIdTokenData idToken, OAuth2Client client) {
      IdTokenBuilder idTokenBody = this.getIDTokenBody(idToken);

      try {
         return ((OpenIDConnectIDTokenSigner)this.idTokenSigner.get()).signed(idTokenBody, client).tokenString();
      } catch (IdTokenBuilder.IdTokenException e) {
         throw new OAuth2RequestHandlingException("Failed to sign OpenID Connect ID Token: " + e.getMessage(), e);
      }
   }

   @Override
   public String getSignedIdTokenWithHybridFlowHashes(
      OpenIdConnectIdTokenData idToken, OAuth2Client client, OAuth2AuthorizationCode authorizationCode, Optional<String> accessTokenString
   ) {
      try {
         JWSAlgorithm algorithm = this.getAlgorithm();
         List<Claim<?>> hybridFlowHashClaims = new ArrayList<>();
         hybridFlowHashClaims.add(new StringClaim("c_hash", computeCodeHash(authorizationCode, algorithm).getValue()));
         computeAccessTokenHash(accessTokenString, algorithm)
            .ifPresent(accessTokenHash -> hybridFlowHashClaims.add(new StringClaim("at_hash", accessTokenHash.getValue())));
         IdTokenBuilder idTokenBody = this.getIDTokenBody(idToken, hybridFlowHashClaims);
         return ((OpenIDConnectIDTokenSigner)this.idTokenSigner.get()).signed(idTokenBody, client).tokenString();
      } catch (IdTokenBuilder.IdTokenException e) {
         throw new OAuth2RequestHandlingException("Failed to sign OpenID Connect ID Token: " + e.getMessage(), e);
      }
   }

   private static Optional<AccessTokenHash> computeAccessTokenHash(Optional<String> accessTokenString, JWSAlgorithm algorithm) {
      if (accessTokenString.isEmpty()) {
         return Optional.empty();
      } else {
         AccessTokenHash accessTokenHash = AccessTokenHash.compute(new TypelessAccessToken(accessTokenString.get()), algorithm, null);
         if (accessTokenHash == null) {
            throw new OAuth2RequestHandlingException(
               String.format("Failed to compute access token hash. JWS algorithm is not supported: '%s'", algorithm.getName())
            );
         } else {
            return Optional.of(accessTokenHash);
         }
      }
   }

   private static CodeHash computeCodeHash(OAuth2AuthorizationCode authorizationCode, JWSAlgorithm algorithm) {
      CodeHash codeHash = CodeHash.compute(new AuthorizationCode(authorizationCode.getAuthorizationCode()), algorithm, null);
      if (codeHash == null) {
         throw new OAuth2RequestHandlingException(String.format("Failed to compute code hash. JWS algorithm is not supported: '%s'", algorithm.getName()));
      } else {
         return codeHash;
      }
   }

   private JWSAlgorithm getAlgorithm() {
      return JWSAlgorithm.parse(
         ((OpenIDConnectIDTokenSigner)this.idTokenSigner.get())
            .algorithm()
            .orElseThrow(() -> new OAuth2RequestHandlingException("Missing ID token algorithm to generate c_hash and at_hash."))
      );
   }

   private IdTokenBuilder getIDTokenBody(OpenIdConnectIdTokenData idToken) {
      return this.getIDTokenBody(idToken, List.of());
   }

   private IdTokenBuilder getIDTokenBody(OpenIdConnectIdTokenData idToken, List<Claim<?>> additionalClaims) {
      if (!((OpenIdConnectSignedIdTokenFactorySettings)this.settings.get()).claimIssuer().isValid()) {
         this.getLogger().warn(OAuth2LogTypes.OAUTH2_MESSAGE, "Configured token issuer is not conform with the OpenID Connect specification.");
      }

      try {
         return IdTokenBuilder.buildFromTokenData(
            ((IntegrationLoggerFactory)this.loggerFactory.get()).create(IdToken.class),
            ((OpenIdConnectSignedIdTokenFactorySettings)this.settings.get()).claimIssuer().getValue(),
            idToken,
            additionalClaims.stream().collect(Collectors.toMap(Claim::name, claim -> claim.value().toJSONSerializable()))
         );
      } catch (IdTokenBuilder.IdTokenException e) {
         throw new OAuth2RequestHandlingException("Generated ID Token does not contain all required claims: " + e.getMessage(), e);
      }
   }

   private ELogger getLogger() {
      return ((IntegrationLoggerFactory)this.loggerFactory.get()).create(OpenIdConnectSignedIdTokenFactoryImpl.class);
   }
}