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

import com.airlock.iam.common.domain.model.jwt.encryption.DirectEncryptionMethod;
import com.airlock.iam.common.domain.model.jwt.encryption.JweToken;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.service.jwt.DefaultSecureJwtConfig;
import com.airlock.iam.common.domain.service.jwt.JwtSettingsValidationError;
import com.airlock.iam.core.misc.util.ticket.TicketEncodingException;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.nimbusds.jose.EncryptionMethod;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWEAlgorithm;
import com.nimbusds.jose.JWECryptoParts;
import com.nimbusds.jose.JWEHeader;
import com.nimbusds.jose.JWEObject;
import com.nimbusds.jose.JWEProvider;
import com.nimbusds.jose.KeyLengthException;
import com.nimbusds.jose.Payload;
import com.nimbusds.jose.JWEHeader.Builder;
import com.nimbusds.jose.crypto.DirectDecrypter;
import com.nimbusds.jose.crypto.DirectEncrypter;
import com.nimbusds.jose.crypto.impl.AAD;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.text.ParseException;
import java.util.Arrays;
import java.util.Optional;
import org.apache.commons.codec.binary.Base64;

public class JwtTicketDirectAesEncrypterDecrypter implements JwtTicketEncrypter, JwtTicketDecrypter {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwtTicketDirectAesEncrypterDecrypter.class);
   private final DirectEncryptionMethod directEncryptionMethod;
   private final DirectEncrypter directEncrypter;
   private final DirectDecrypter directDecrypter;

   public static JwtTicketDirectAesEncrypterDecrypter createDefaultEncryptorDecryptor(String encryptionKey) {
      return builder().directEncryptionMethod(DefaultSecureJwtConfig.DEFAULT_ENCRYPTION_METHOD).encryptionKey(encryptionKey).build();
   }

   public static Optional<JwtSettingsValidationError> validate(DirectEncryptionMethod directEncryptionMethod, String encryptionKey) {
      try {
         EncryptionMethod encryptionMethod = encryptionMethod(directEncryptionMethod);
         DirectEncrypter directEncrypter = new DirectEncrypter(Base64.decodeBase64(encryptionKey));
         DirectDecrypter directDecrypter = new DirectDecrypter(Base64.decodeBase64(encryptionKey));
         Optional<JwtSettingsValidationError> encrypterResult = assertProviderSupportsAlgorithmEncryptionMethodAndKey(
            directEncrypter, JWEAlgorithm.DIR, encryptionMethod
         );
         return encrypterResult.isPresent()
            ? encrypterResult
            : assertProviderSupportsAlgorithmEncryptionMethodAndKey(directDecrypter, JWEAlgorithm.DIR, encryptionMethod);
      } catch (Exception e) {
         return Optional.of(new JwtSettingsValidationError("Illegal encryption key", Optional.of(e)));
      }
   }

   private static Optional<JwtSettingsValidationError> assertProviderSupportsAlgorithmEncryptionMethodAndKey(
      JWEProvider provider, JWEAlgorithm algorithm, EncryptionMethod encryptionMethod
   ) {
      if (!provider.supportedJWEAlgorithms().contains(algorithm)) {
         return Optional.of(
            new JwtSettingsValidationError("The provided algorithm is not supported for the given key: " + algorithm.getName(), Optional.empty())
         );
      } else {
         return !provider.supportedEncryptionMethods().contains(encryptionMethod)
            ? Optional.of(
               new JwtSettingsValidationError(
                  "The provided encryption method is not supported for the given key: " + encryptionMethod.getName(), Optional.empty()
               )
            )
            : Optional.empty();
      }
   }

   private static JWEHeader jweHeader(JWEAlgorithm algorithm, EncryptionMethod encryptionMethod) {
      return new Builder(algorithm, encryptionMethod).contentType("JWT").build();
   }

   private static EncryptionMethod encryptionMethod(DirectEncryptionMethod encryptionMethod) {
      return DirectEncrypter.SUPPORTED_ENCRYPTION_METHODS
         .stream()
         .filter(enc -> isEqual(enc, encryptionMethod))
         .findAny()
         .orElseThrow(() -> new IllegalArgumentException("Unsupported encryption method: " + encryptionMethod.name()));
   }

   private static boolean isEqual(EncryptionMethod encryptionMethod, DirectEncryptionMethod dirEncryptionMethod) {
      try {
         return DirectEncryptionMethod.of(encryptionMethod).equals(dirEncryptionMethod);
      } catch (IllegalArgumentException e) {
         return false;
      }
   }

   JwtTicketDirectAesEncrypterDecrypter(DirectEncryptionMethod directEncryptionMethod, String encryptionKey) {
      this.directEncryptionMethod = directEncryptionMethod;

      try {
         this.directEncrypter = new DirectEncrypter(Base64.decodeBase64(encryptionKey));
         this.directDecrypter = new DirectDecrypter(Base64.decodeBase64(encryptionKey));
         this.encryptionDecryptionSelfTest();
      } catch (KeyLengthException e) {
         throw new IllegalArgumentException("Could not initialize encrypter/decrypter with given key.", e);
      }
   }

   @Override
   public JsonWebSignatureToken decrypt(JweToken jweToken) throws TicketEncodingException {
      JWEObject jwe = this.parseJweObject(jweToken);
      this.checkAlgorithmAndEncryptionMethod(jwe.getHeader());
      return this.decrypt(jwe);
   }

   private JWEObject parseJweObject(JweToken jweToken) throws TicketEncodingException {
      try {
         return JWEObject.parse(jweToken.getJwe());
      } catch (ParseException e) {
         throw new TicketEncodingException("Unable to parse the decoded JWT as an encrypted JWT. A reason could be that the JWT is not encrypted.", e);
      }
   }

   private void checkAlgorithmAndEncryptionMethod(JWEHeader jweHeader) throws TicketEncodingException {
      if (!JWEAlgorithm.DIR.equals(jweHeader.getAlgorithm())) {
         throw new TicketEncodingException("Unexpected JWT Encryption Algorithm! Expected: DIR, but was: " + jweHeader.getAlgorithm().getName());
      }

      EncryptionMethod encryptionMethod = encryptionMethod(this.directEncryptionMethod);
      if (!encryptionMethod.equals(jweHeader.getEncryptionMethod())) {
         throw new TicketEncodingException(
            "Unexpected JWT Encryption Method! Expected: " + this.directEncryptionMethod.name() + ", but was: " + jweHeader.getEncryptionMethod().getName()
         );
      }
   }

   private JsonWebSignatureToken decrypt(JWEObject jwe) throws TicketEncodingException {
      try {
         jwe.decrypt(this.directDecrypter);
         LOG.debug("Successfully decrypted the JWT ticket.");
         return JsonWebSignatureToken.of(jwe.getPayload().toString());
      } catch (JOSEException e) {
         throw new TicketEncodingException("Unable to decrypt the JWT.", e);
      }
   }

   @Override
   public JweToken encrypt(JsonWebSignatureToken payload) throws TicketEncodingException {
      JWEObject encryptedJwt = new JWEObject(
         new Builder(JWEAlgorithm.DIR, encryptionMethod(this.directEncryptionMethod)).contentType("JWT").build(), new Payload(payload.tokenString())
      );
      this.applyEncryption(encryptedJwt);
      return JweToken.of(encryptedJwt.serialize());
   }

   private void applyEncryption(JWEObject encryptedJwt) throws TicketEncodingException {
      try {
         encryptedJwt.encrypt(this.directEncrypter);
      } catch (JOSEException e) {
         throw new TicketEncodingException("Error encrypting JWT.", e);
      }
   }

   private void encryptionDecryptionSelfTest() {
      try {
         byte[] data = new byte[]{1, 2, 3};
         JWEHeader jweHeader = jweHeader(JWEAlgorithm.DIR, encryptionMethod(this.directEncryptionMethod));
         JWECryptoParts cryptoParts = this.directEncrypter.encrypt(jweHeader, data, AAD.compute(jweHeader));
         JWEHeader header = cryptoParts.getHeader();
         byte[] plaintext = this.directDecrypter
            .decrypt(
               header,
               cryptoParts.getEncryptedKey(),
               cryptoParts.getInitializationVector(),
               cryptoParts.getCipherText(),
               cryptoParts.getAuthenticationTag(),
               AAD.compute(header)
            );
         if (!Arrays.equals(data, plaintext)) {
            throw new IllegalStateException("Encryption-decryption test failed.");
         }
      } catch (JOSEException e) {
         throw new IllegalStateException("Encryption-decryption test failed.", e);
      }
   }

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

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtTicketDirectAesEncrypterDecrypterBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private DirectEncryptionMethod directEncryptionMethod;
      @SuppressFBWarnings(justification = "generated code")
      private String encryptionKey;

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

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketDirectAesEncrypterDecrypter.JwtTicketDirectAesEncrypterDecrypterBuilder directEncryptionMethod(
         DirectEncryptionMethod directEncryptionMethod
      ) {
         this.directEncryptionMethod = directEncryptionMethod;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketDirectAesEncrypterDecrypter.JwtTicketDirectAesEncrypterDecrypterBuilder encryptionKey(String encryptionKey) {
         this.encryptionKey = encryptionKey;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketDirectAesEncrypterDecrypter build() {
         return new JwtTicketDirectAesEncrypterDecrypter(this.directEncryptionMethod, this.encryptionKey);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtTicketDirectAesEncrypterDecrypter.JwtTicketDirectAesEncrypterDecrypterBuilder(directEncryptionMethod="
            + this.directEncryptionMethod
            + ", encryptionKey="
            + this.encryptionKey
            + ")";
      }
   }
}