package com.airlock.iam.common.domain.service.e2ee;

import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.base.api.domain.model.plugin.LogicForConfig;
import com.airlock.iam.base.misc.util.Base64Util;
import com.airlock.iam.common.api.domain.model.e2ee.E2eEncryptionSession;
import com.airlock.iam.common.api.domain.model.password.Password;
import com.airlock.iam.common.application.configuration.e2ee.JwePasswordDecryptionConfig;
import com.airlock.iam.common.domain.model.e2ee.JweE2eEncryptionSession;
import com.airlock.iam.common.domain.model.password.PasswordImpl;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.security.util.crypto.nonce.SecureRandomNonceGenerator;
import com.fasterxml.jackson.core.JsonPointer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWEDecrypter;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jwt.EncryptedJWT;
import jakarta.inject.Inject;
import java.text.ParseException;
import java.util.Optional;

@LogicForConfig(JwePasswordDecryptionConfig.class)
public class JwePasswordDecryptionService implements PasswordDecryptionService {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwePasswordDecryptionService.class);
   private final JWK publicKey;
   private final JWEDecrypter decrypter;
   private final JsonPointer passwordJsonPointer;

   @Inject
   public JwePasswordDecryptionService(JWEDecrypter decrypter, JWK publicKey, JsonPointer passwordJsonPointer) {
      this.publicKey = publicKey;
      this.passwordJsonPointer = passwordJsonPointer;
      this.decrypter = decrypter;
   }

   @Override
   public Optional<Password> decryptPassword(Password encryptedPassword, E2eEncryptionSession session) {
      String encryptedPasswordJWT = encryptedPassword.getValue();
      JweE2eEncryptionSession jweE2eEncryptionSession = ViolationException.checkType(session, JweE2eEncryptionSession.class);
      String expectedNonce = jweE2eEncryptionSession.getPublicEncryptionInformation().getNonce();

      try {
         EncryptedJWT passwordJWT = EncryptedJWT.parse(encryptedPasswordJWT);
         Optional<String> nonceVerificationResult = this.parseAndVerifyNonce(passwordJWT, expectedNonce);
         if (nonceVerificationResult.isEmpty()) {
            return Optional.empty();
         }

         passwordJWT.decrypt(this.decrypter);
         Optional<String> passwordParseResult = parsePasswordFromPayload(passwordJWT, this.passwordJsonPointer);
         if (passwordParseResult.isEmpty()) {
            return Optional.empty();
         }

         LOG.debug("Successfully decrypted the password JWT and verified the contained nonce.");
         return Optional.of(PasswordImpl.passwordOf(passwordParseResult.get()));
      } catch (JOSEException e) {
         LOG.debug("Decrypting the password JWT failed: " + e.getMessage());
         return Optional.empty();
      } catch (ParseException e) {
         LOG.debug(String.format("Parsing the encrypted JWT (%s) failed with error: %s", encryptedPasswordJWT, e.getMessage()));
         return Optional.empty();
      }
   }

   public JweE2eEncryptionSession createE2eEncryptionSession(E2eEncryptionSession.E2eEncryptionUseCase useCase) {
      return new JweE2eEncryptionSession(generateFreshNonce(), this.publicKey);
   }

   @Override
   public void handleFailedPasswordCheck(E2eEncryptionSession e2eEncryptionSession) {
      JweE2eEncryptionSession jweE2eEncryptionSession = ViolationException.checkType(e2eEncryptionSession, JweE2eEncryptionSession.class);
      jweE2eEncryptionSession.resetNonce(generateFreshNonce());
   }

   private static String generateFreshNonce() {
      return Base64Util.encodeUrlSafeToString(new SecureRandomNonceGenerator().generateNonce());
   }

   private Optional<String> parseAndVerifyNonce(EncryptedJWT jwe, String expectedNonce) {
      Optional<String> nonceParseResult = parseNonceFromHeader(jwe);
      if (nonceParseResult.isEmpty()) {
         return Optional.empty();
      } else if (!nonceParseResult.get().equals(expectedNonce)) {
         LOG.info(
            String.format(
               "Password decryption problem: E2EE nonce returned by client (%s) does not match nonce provided by server (%s).",
               nonceParseResult.get(),
               expectedNonce
            )
         );
         return Optional.empty();
      } else {
         return nonceParseResult;
      }
   }

   private static Optional<String> parseNonceFromHeader(EncryptedJWT jwe) {
      try {
         Optional<String> customNonceHeader = Optional.ofNullable((String)jwe.getHeader().getCustomParam("nonce"));
         if (customNonceHeader.isEmpty()) {
            LOG.debug("The header of the JWE does not contain an entry with key \"nonce\".");
            return Optional.empty();
         } else {
            return customNonceHeader;
         }
      } catch (ClassCastException e) {
         LOG.debug(String.format("The nonce in the JWE header must be of type String: %s", e.getMessage()), e);
         return Optional.empty();
      }
   }

   private static Optional<String> parsePasswordFromPayload(EncryptedJWT jwe, JsonPointer passwordJsonPointer) {
      try {
         ObjectMapper mapper = new ObjectMapper();
         JsonNode payload = mapper.readTree(jwe.getJWTClaimsSet().toString());
         JsonNode password = payload.at(passwordJsonPointer);
         if (!password.isTextual()) {
            LOG.debug(
               String.format(
                  "Parsing the password from the JWE payload failed. Make sure the JSON Pointer (%s) resolves to the plaintext password of type JSON String.",
                  passwordJsonPointer
               )
            );
            return Optional.empty();
         } else {
            return Optional.of(password.textValue());
         }
      } catch (ParseException | JsonProcessingException e) {
         LOG.debug(String.format("Parsing the JWE claim set into a JSON Node failed: %s", e.getMessage()), e);
         return Optional.empty();
      }
   }
}