package com.airlock.iam.oauth2.application.configuration.client;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.common.domain.service.jwt.jwks.JwksSignatureValidator;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.oauth2.domain.model.client.OAuth2ClientSideCredentials;
import com.airlock.iam.oauth2.domain.model.client.communication.OpenIDConnectIdTokenSignatureValidator;
import com.airlock.iam.oauth2.domain.model.token.IdToken;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.nimbusds.jose.Algorithm;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.oauth2.sdk.auth.ClientSecretJWT;
import com.nimbusds.oauth2.sdk.auth.PrivateKeyJWT;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Set;

public final class OpenIDConnectDiscoverySignatureValidator implements OpenIDConnectIdTokenSignatureValidator {
   private static final Set<JWSAlgorithm> SUPPORTED_PRIVATE_KEY_ALGORITHMS = PrivateKeyJWT.supportedJWAs();
   private static final Set<JWSAlgorithm> SUPPORTED_CLIENT_SECRET_ALGORITHMS = ClientSecretJWT.supportedJWAs();
   public static final Set<JWSAlgorithm> COMPATIBLE_ID_TOKEN_SIGNING_ALGORITHMS = ImmutableSet.copyOf(
      Iterables.concat(SUPPORTED_PRIVATE_KEY_ALGORITHMS, SUPPORTED_CLIENT_SECRET_ALGORITHMS)
   );
   private final JwksSignatureValidator validator;
   private final ELogger logger;

   @Override
   public boolean validateIdTokenSignature(IdToken idToken, OAuth2ClientSideCredentials clientCredentials) {
      if (!idToken.isSigned()) {
         this.logger.info("Rejecting ID Token since it was not signed");
         return false;
      }

      Algorithm signatureAlgorithm = Algorithm.parse(idToken.signatureAlgorithmName());
      if (SUPPORTED_CLIENT_SECRET_ALGORITHMS.contains(signatureAlgorithm)) {
         return idToken.validateMACSignature(
            clientCredentials.clientSecret().orElseThrow(() -> new IllegalArgumentException("Client secret missing. Unable to validate signature."))
         );
      }

      if (SUPPORTED_PRIVATE_KEY_ALGORITHMS.contains(signatureAlgorithm)) {
         return this.validatePrivateKeySignature(idToken);
      }

      this.logger.info("Rejecting ID Token since it was signed with unsupported algorithm '" + signatureAlgorithm + "'");
      return false;
   }

   private boolean validatePrivateKeySignature(IdToken idToken) {
      try {
         return idToken.getSignedJWT()
            .map(jwt -> this.validator.validateSignature(jwt.header().keyId(), JwksSignatureValidator.defaultValidator(jwt)))
            .orElse(false);
      } catch (FailureException e) {
         this.logger.info("Unable to fetch JWK from URI: " + e.getMessage());
         return false;
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public OpenIDConnectDiscoverySignatureValidator(JwksSignatureValidator validator, ELogger logger) {
      this.validator = validator;
      this.logger = logger;
   }
}