package com.airlock.iam.common.domain.model.jwt.signature;

import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.common.domain.service.jwt.signature.JwtSignatureUtils;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSVerifier;
import com.nimbusds.jose.crypto.ECDSAVerifier;
import com.nimbusds.jose.crypto.MACVerifier;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.SignedJWT;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.security.PublicKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.util.Collection;
import java.util.Optional;
import java.util.stream.Collectors;

public final class JwsVerifierImpl implements JwsVerifier {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwsVerifierImpl.class);
   private final JWSVerifier verifier;

   public static JwsVerifier ecdsaJwsVerifier(ECPublicKey publicKey) {
      JwtSignatureUtils.addBouncyCastleSecurityProviderIfAbsent();

      try {
         return new JwsVerifierImpl(new ECDSAVerifier(publicKey));
      } catch (JOSEException e) {
         throw new ViolationException("Could not create verifier from public key", e);
      }
   }

   public static JwsVerifier rsaJwsVerifier(RSAPublicKey publicKey) {
      return new JwsVerifierImpl(new RSASSAVerifier(publicKey));
   }

   public static JwsVerifier hmacJwsVerifier(byte[] secret) {
      JwtSignatureUtils.addBouncyCastleSecurityProviderIfAbsent();

      try {
         return new JwsVerifierImpl(new MACVerifier(secret));
      } catch (JOSEException e) {
         throw new IllegalArgumentException("Illegal hmac key", e);
      }
   }

   public static JwsVerifier forPublicKey(PublicKey publicKey) {
      String algorithm = publicKey.getAlgorithm();
      if (algorithm.contains("RSA")) {
         return rsaJwsVerifier((RSAPublicKey)publicKey);
      } else if (algorithm.contains("EC")) {
         return ecdsaJwsVerifier((ECPublicKey)publicKey);
      } else {
         throw new IllegalArgumentException("The public key associated with the token is of unknown type: " + LogUtil.quote(algorithm));
      }
   }

   public static JwsVerifier forAlgorithm(SignatureAlgorithm jwsAlgorithm, PublicKey publicKey) {
      return switch (jwsAlgorithm.family()) {
         case RSA -> {
            JwsVerifier var6;
            try {
               var6 = rsaJwsVerifier((RSAPublicKey)publicKey);
            } catch (ClassCastException e) {
               throw new IllegalArgumentException("Illegal RSA public key", e);
            }

            yield var6;
         }
         case EC -> {
            JwsVerifier var2;
            try {
               var2 = ecdsaJwsVerifier((ECPublicKey)publicKey);
            } catch (ClassCastException e) {
               throw new IllegalArgumentException("Illegal EC public key", e);
            }

            yield var2;
         }
         case HMAC -> throw new IllegalArgumentException("The provided algorithm is no signature algorithm: " + jwsAlgorithm);
      };
   }

   @Override
   public boolean verify(JsonWebSignatureToken jwsToken) {
      SignedJWT signedJWT;
      try {
         signedJWT = SignedJWT.parse(jwsToken.tokenString());
      } catch (ParseException e) {
         throw new ViolationException("Invalid signed JWT", e);
      }

      try {
         return this.verifier.verify(signedJWT.getHeader(), signedJWT.getSigningInput(), signedJWT.getSignature());
      } catch (JOSEException e) {
         LOG.info("Failed to verify signature: " + e.getMessage());
         return false;
      }
   }

   @Override
   public Collection<JwsAlgorithm> supportedJWSAlgorithms() {
      return this.verifier.supportedJWSAlgorithms().stream().map(JwsVerifierImpl::getMatchingAlgorithm).flatMap(Optional::stream).collect(Collectors.toSet());
   }

   private static Optional<JwsAlgorithm> getMatchingAlgorithm(JWSAlgorithm algorithm) {
      return JwsAlgorithm.JWS_ALGORITHMS.stream().filter(jwsAlgorithm -> algorithm.getName().equals(jwsAlgorithm.name())).findFirst();
   }

   @SuppressFBWarnings(justification = "generated code")
   private JwsVerifierImpl(JWSVerifier verifier) {
      this.verifier = verifier;
   }
}