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

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.common.domain.model.jwt.signature.HmacAlgorithm;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.model.jwt.signature.JwsAlgorithm;
import com.airlock.iam.common.domain.model.jwt.signature.JwsVerifier;
import com.airlock.iam.common.domain.model.jwt.signature.JwsVerifierImpl;
import com.airlock.iam.common.domain.model.jwt.signature.SignatureAlgorithm;
import com.airlock.iam.common.domain.service.jwt.JwtSettingsValidationError;
import com.airlock.iam.common.domain.service.jwt.jwks.JwksSignatureValidator;
import com.airlock.iam.elog.api.EDefaultLogType;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.security.util.crypto.KeyStoreUtils;
import com.airlock.iam.security.util.crypto.KeyStoreUtils.KeyStoreType;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.File;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;

public class JwtSignatureVerifierImpl implements JwtSignatureVerifier {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwtSignatureVerifierImpl.class);
   private final Optional<JwsAlgorithm> algorithmConstraint;
   private final JwsVerifierFactory verifierFactory;

   public static Optional<JwtSettingsValidationError> validate(HmacAlgorithm hmacAlgorithm, String hmacKey) {
      try {
         assertVerifierSupportsAlgorithm(JwsVerifierImpl.hmacJwsVerifier(Base64.decodeBase64(hmacKey)), hmacAlgorithm);
         return Optional.empty();
      } catch (IllegalArgumentException e) {
         return Optional.of(new JwtSettingsValidationError(e.getMessage(), Optional.of(e)));
      }
   }

   public static Optional<JwtSettingsValidationError> validate(SignatureAlgorithm signatureAlgorithm, PublicKey publicKey) {
      try {
         assertVerifierSupportsAlgorithm(JwsVerifierImpl.forAlgorithm(signatureAlgorithm, publicKey), signatureAlgorithm);
         return Optional.empty();
      } catch (IllegalArgumentException e) {
         return Optional.of(new JwtSettingsValidationError(e.getMessage(), Optional.of(e)));
      }
   }

   public static Optional<JwsVerifier> jwksVerifier(final JsonWebSignatureToken token, final JwksSignatureValidator validator) {
      return Optional.of(new JwsVerifier() {
         @Override
         public boolean verify(JsonWebSignatureToken jwsToken) {
            try {
               return validator.validateSignature(jwsToken.header().keyId(), JwksSignatureValidator.defaultValidator(jwsToken));
            } catch (FailureException e) {
               JwtSignatureVerifierImpl.LOG.warn(EDefaultLogType.GENERIC_LOG_TYPE, "Failed to validate signature: " + e.getMessage());
               return false;
            }
         }

         @Override
         public Collection<JwsAlgorithm> supportedJWSAlgorithms() {
            return List.of(token.header().getValidatedAlgorithm());
         }
      });
   }

   private static void assertVerifierSupportsAlgorithm(JwsVerifier verifier, JwsAlgorithm jwsAlgorithm) {
      if (verifier != null && !verifier.supportedJWSAlgorithms().contains(jwsAlgorithm)) {
         throw new IllegalArgumentException("The selected algorithm (" + jwsAlgorithm.name() + ") is not supported by the given key");
      }
   }

   public JwtSignatureVerifierImpl(HmacAlgorithm hmacAlgorithm, String hmacKey) {
      JwsVerifier hmacVerifier = JwsVerifierImpl.hmacJwsVerifier(Base64.decodeBase64(hmacKey));
      assertVerifierSupportsAlgorithm(hmacVerifier, hmacAlgorithm);
      this.algorithmConstraint = Optional.of(hmacAlgorithm);
      this.verifierFactory = ignored -> Optional.of(hmacVerifier);
   }

   public JwtSignatureVerifierImpl(SignatureAlgorithm signatureAlgorithm, File keyStorePath, String keyStorePassword, String certificateAlias) {
      this.algorithmConstraint = Optional.of(signatureAlgorithm);

      try {
         KeyStore keyStore = KeyStoreUtils.createKeyStoreFromFile(keyStorePath, optional(keyStorePassword), KeyStoreType.AUTO);
         X509Certificate certificate = (X509Certificate)KeyStoreUtils.getUniqueX509Certificate(
               keyStore, optional(certificateAlias), Optional.of(signatureAlgorithm.getCertificateFilter())
            )
            .orElseThrow(() -> new IllegalArgumentException("Could not extract certificate from keystore"));
         JwsVerifier verifier = JwsVerifierImpl.forAlgorithm(signatureAlgorithm, certificate.getPublicKey());
         this.verifierFactory = ignored -> Optional.of(verifier);
         assertVerifierSupportsAlgorithm(verifier, signatureAlgorithm);
      } catch (KeyStoreException e) {
         throw new IllegalArgumentException("Could not load keystore", e);
      }
   }

   public JwtSignatureVerifierImpl(SignatureAlgorithm signatureAlgorithm, X509Certificate certificate) {
      JwsVerifier verifier = JwsVerifierImpl.forAlgorithm(signatureAlgorithm, certificate.getPublicKey());
      assertVerifierSupportsAlgorithm(verifier, signatureAlgorithm);
      this.algorithmConstraint = Optional.of(signatureAlgorithm);
      this.verifierFactory = ignored -> Optional.of(verifier);
   }

   public JwtSignatureVerifierImpl(SignatureAlgorithm signatureAlgorithm, PublicKey publicKey) {
      JwsVerifier verifier = JwsVerifierImpl.forAlgorithm(signatureAlgorithm, publicKey);
      assertVerifierSupportsAlgorithm(verifier, signatureAlgorithm);
      this.algorithmConstraint = Optional.of(signatureAlgorithm);
      this.verifierFactory = ignored -> Optional.of(verifier);
   }

   public JwtSignatureVerifierImpl(JwksSignatureValidator validator) {
      this.algorithmConstraint = Optional.empty();
      this.verifierFactory = token -> jwksVerifier(token, validator);
   }

   private static Optional<String> optional(String val) {
      return StringUtils.isEmpty(val) ? Optional.empty() : Optional.of(val);
   }

   @Override
   public boolean isSignatureValid(JsonWebSignatureToken signedJwt) throws JwtSignatureVerifierException {
      boolean isValid = this.isSameAlgorithmAsConfigured(signedJwt) && this.isValidSignature(signedJwt);
      if (isValid) {
         LOG.debug("Successfully verified the JWT Signature");
      }

      return isValid;
   }

   private boolean isSameAlgorithmAsConfigured(JsonWebSignatureToken signedJWT) {
      boolean isCorrectAlgorithm = this.algorithmConstraint.<Boolean>map(alg -> alg.name().equals(signedJWT.header().algorithmName())).orElse(true);
      if (!isCorrectAlgorithm) {
         LOG.debug(
            "The JWT Algorithm does not match the configured algorithm! Configured: "
               + this.algorithmConstraint.get().name()
               + " Algorithm in JWT Header: "
               + signedJWT.header().algorithmName()
         );
      }

      return isCorrectAlgorithm;
   }

   private boolean isValidSignature(JsonWebSignatureToken signedJwt) throws JwtSignatureVerifierException {
      try {
         Optional<JwsVerifier> verifier = this.verifierFactory.create(signedJwt);
         if (verifier.isPresent()) {
            return verifier.get().verify(signedJwt);
         }

         LOG.warn(EDefaultLogType.GENERIC_LOG_TYPE, "Unable to verify signature: Verification information missing or not supported");
         return false;
      } catch (ViolationException e) {
         throw new JwtSignatureVerifierException("The JWS Object could not be verified", e);
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwtSignatureVerifierImpl.JwtSignatureVerifierImplHmacBuilder builder() {
      return new JwtSignatureVerifierImpl.JwtSignatureVerifierImplHmacBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder keystoreBuilder() {
      return new JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwtSignatureVerifierImpl.JwtSignatureVerifierImplCertificateBuilder certificateBuilder() {
      return new JwtSignatureVerifierImpl.JwtSignatureVerifierImplCertificateBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwtSignatureVerifierImpl.JwtSignatureVerifierImplPublicKeyBuilder publicKeyBuilder() {
      return new JwtSignatureVerifierImpl.JwtSignatureVerifierImplPublicKeyBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static JwtSignatureVerifierImpl.JwtSignatureVerifierImplJwksBuilder jwksBuilder() {
      return new JwtSignatureVerifierImpl.JwtSignatureVerifierImplJwksBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtSignatureVerifierImplCertificateBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private SignatureAlgorithm signatureAlgorithm;
      @SuppressFBWarnings(justification = "generated code")
      private X509Certificate certificate;

      @SuppressFBWarnings(justification = "generated code")
      JwtSignatureVerifierImplCertificateBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplCertificateBuilder signatureAlgorithm(SignatureAlgorithm signatureAlgorithm) {
         this.signatureAlgorithm = signatureAlgorithm;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplCertificateBuilder certificate(X509Certificate certificate) {
         this.certificate = certificate;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl build() {
         return new JwtSignatureVerifierImpl(this.signatureAlgorithm, this.certificate);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtSignatureVerifierImpl.JwtSignatureVerifierImplCertificateBuilder(signatureAlgorithm="
            + this.signatureAlgorithm
            + ", certificate="
            + this.certificate
            + ")";
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtSignatureVerifierImplHmacBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private HmacAlgorithm hmacAlgorithm;
      @SuppressFBWarnings(justification = "generated code")
      private String hmacKey;

      @SuppressFBWarnings(justification = "generated code")
      JwtSignatureVerifierImplHmacBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplHmacBuilder hmacAlgorithm(HmacAlgorithm hmacAlgorithm) {
         this.hmacAlgorithm = hmacAlgorithm;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplHmacBuilder hmacKey(String hmacKey) {
         this.hmacKey = hmacKey;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl build() {
         return new JwtSignatureVerifierImpl(this.hmacAlgorithm, this.hmacKey);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtSignatureVerifierImpl.JwtSignatureVerifierImplHmacBuilder(hmacAlgorithm=" + this.hmacAlgorithm + ", hmacKey=" + this.hmacKey + ")";
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtSignatureVerifierImplJwksBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private JwksSignatureValidator validator;

      @SuppressFBWarnings(justification = "generated code")
      JwtSignatureVerifierImplJwksBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplJwksBuilder validator(JwksSignatureValidator validator) {
         this.validator = validator;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl build() {
         return new JwtSignatureVerifierImpl(this.validator);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtSignatureVerifierImpl.JwtSignatureVerifierImplJwksBuilder(validator=" + this.validator + ")";
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtSignatureVerifierImplKeystoreBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private SignatureAlgorithm signatureAlgorithm;
      @SuppressFBWarnings(justification = "generated code")
      private File keyStorePath;
      @SuppressFBWarnings(justification = "generated code")
      private String keyStorePassword;
      @SuppressFBWarnings(justification = "generated code")
      private String certificateAlias;

      @SuppressFBWarnings(justification = "generated code")
      JwtSignatureVerifierImplKeystoreBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder signatureAlgorithm(SignatureAlgorithm signatureAlgorithm) {
         this.signatureAlgorithm = signatureAlgorithm;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder keyStorePath(File keyStorePath) {
         this.keyStorePath = keyStorePath;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder keyStorePassword(String keyStorePassword) {
         this.keyStorePassword = keyStorePassword;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder certificateAlias(String certificateAlias) {
         this.certificateAlias = certificateAlias;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl build() {
         return new JwtSignatureVerifierImpl(this.signatureAlgorithm, this.keyStorePath, this.keyStorePassword, this.certificateAlias);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtSignatureVerifierImpl.JwtSignatureVerifierImplKeystoreBuilder(signatureAlgorithm="
            + this.signatureAlgorithm
            + ", keyStorePath="
            + this.keyStorePath
            + ", keyStorePassword="
            + this.keyStorePassword
            + ", certificateAlias="
            + this.certificateAlias
            + ")";
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtSignatureVerifierImplPublicKeyBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private SignatureAlgorithm signatureAlgorithm;
      @SuppressFBWarnings(justification = "generated code")
      private PublicKey publicKey;

      @SuppressFBWarnings(justification = "generated code")
      JwtSignatureVerifierImplPublicKeyBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplPublicKeyBuilder signatureAlgorithm(SignatureAlgorithm signatureAlgorithm) {
         this.signatureAlgorithm = signatureAlgorithm;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl.JwtSignatureVerifierImplPublicKeyBuilder publicKey(PublicKey publicKey) {
         this.publicKey = publicKey;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtSignatureVerifierImpl build() {
         return new JwtSignatureVerifierImpl(this.signatureAlgorithm, this.publicKey);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtSignatureVerifierImpl.JwtSignatureVerifierImplPublicKeyBuilder(signatureAlgorithm="
            + this.signatureAlgorithm
            + ", publicKey="
            + this.publicKey
            + ")";
      }
   }
}