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

import com.airlock.iam.common.domain.model.jwt.SignatureVerificationInformation;
import com.airlock.iam.common.domain.model.jwt.signature.JwsAlgorithmEc;
import com.airlock.iam.common.domain.model.jwt.signature.JwsAlgorithmRsa;
import com.airlock.iam.common.domain.model.jwt.signature.SignatureAlgorithm;
import com.airlock.iam.common.domain.service.JwksPublicKeyRegistry;
import com.airlock.iam.common.domain.service.jwt.JwtSettingsValidationError;
import com.airlock.iam.common.domain.service.jwt.signature.PrivateKeyJwtSigner;
import com.airlock.iam.common.infrastructure.service.JwksProvider;
import com.airlock.iam.core.misc.plugin.config.validation.KeystoreValidator;
import com.airlock.iam.core.misc.plugin.config.validation.OpenedKeystoreValidatorHelper;
import com.airlock.iam.plugin.api.annotation.config.ConfigurationValidator;
import com.airlock.iam.plugin.api.annotation.lifecycle.AfterConfigActivation;
import com.airlock.iam.plugin.api.annotation.lifecycle.BeforePluginDeactivation;
import com.airlock.iam.plugin.api.annotation.property.EnumProperty;
import com.airlock.iam.plugin.api.annotation.property.FileProperty;
import com.airlock.iam.plugin.api.annotation.property.StringProperty;
import com.airlock.iam.plugin.api.config.ConfigurationContext;
import com.airlock.iam.plugin.api.config.ConfigurationEnvironment;
import com.airlock.iam.plugin.api.config.Environment;
import com.airlock.iam.plugin.api.config.PluginConfiguration;
import com.airlock.iam.plugin.api.exception.PropertyConfigurationException;
import com.airlock.iam.plugin.api.exception.PropertyConstraintViolationException;
import com.airlock.iam.plugin.api.valuereader.EnumPropertyValueReader;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.File;
import java.security.PrivateKey;
import java.util.List;
import java.util.Objects;
import java.util.Optional;

public abstract class PrivateKeyJwtSignatureConfig implements JwksPublicKeyRegistry {
   @EnumProperty(displayName = "Algorithm", description = "Private key based signature algorithm to use.", optional = true, defaultValue = "RS256")
   private PrivateKeyJwtSignatureConfig.Algorithm algorithm;
   public static final String ALGORITHM_DISPLAY_NAME = "Algorithm";
   public static final String ALGORITHM_FIELD = "algorithm";
   @FileProperty(
      displayName = "Keystore File",
      description = "Keystore file name containing the certificate and key used to sign the JWT.",
      optional = false,
      exampleValues = {"keystore.jks", "/keys/keystore.p12"},
      assureExistence = true,
      assureReadable = true,
      isFile = true
   )
   private File keystoreFile;
   public static final String KEYSTORE_FILE_DISPLAY_NAME = "Keystore File";
   public static final String KEYSTORE_FILE_FIELD = "keystoreFile";
   @StringProperty(displayName = "Keystore Password", description = "The password used to open the keystore.", optional = true, isSensitive = true)
   private String keystorePassword;
   public static final String KEYSTORE_PASSWORD_DISPLAY_NAME = "Keystore Password";
   public static final String KEYSTORE_PASSWORD_FIELD = "keystorePassword";
   @StringProperty(
      displayName = "Signing Key Alias",
      description = "The alias of the key used to sign the JWT. This field can be omitted if the keystore only contains one private key entry.",
      optional = true,
      exampleValues = "alias"
   )
   private String signingKeyAlias;
   public static final String SIGNINGKEY_ALIAS_DISPLAY_NAME = "Signing Key Alias";
   public static final String SIGNINGKEY_ALIAS_FIELD = "signingKeyAlias";
   @StringProperty(
      displayName = "Signing Key Password",
      description = "The password used to retrieve the key from the keystore. This password can be the same as the keystore password.",
      optional = false,
      isSensitive = true
   )
   private String signingKeyPassword;
   public static final String SIGNINGKEY_PASSWORD_DISPLAY_NAME = "Signing Key Password";
   public static final String SIGNINGKEY_PASSWORD_FIELD = "signingKeyPassword";
   private PrivateKeyJwtSigner signer;

   protected PrivateKeyJwtSignatureConfig(
      PrivateKeyJwtSignatureConfig.Algorithm algorithm, File keystoreFile, String keystorePassword, String signingKeyAlias, String signingKeyPassword
   ) {
      this.algorithm = algorithm;
      this.keystoreFile = keystoreFile;
      this.keystorePassword = keystorePassword;
      this.signingKeyAlias = signingKeyAlias;
      this.signingKeyPassword = signingKeyPassword;
   }

   protected abstract boolean isIncludeKid();

   @ConfigurationValidator(localOnly = true)
   static void validatePrivateKeyAndCertificate(
      PluginConfiguration config, ConfigurationEnvironment environment, ConfigurationContext context, Environment validationEnvironment
   ) throws PropertyConfigurationException {
      if (validationEnvironment == Environment.CURRENT_VM) {
         PrivateKeyJwtSignatureConfig.Algorithm algorithm = (PrivateKeyJwtSignatureConfig.Algorithm)EnumPropertyValueReader.effectiveConfigEnum(
            config, "algorithm", environment, context
         );
         SignatureAlgorithm signatureAlgorithm = algorithm.getAlgorithm();
         OpenedKeystoreValidatorHelper validator = KeystoreValidator.getAutodetectingKeystoreValidator(
               "keystoreFile", "keystorePassword", config, environment, context
            )
            .validateKeystore();
         PrivateKey privateKey = validator.matchingPrivateKeyWithSpecificPasswordOrElseThrow(
            "signingKeyAlias", "signingKeyPassword", Optional.of(signatureAlgorithm.getPrivateKeyFilter())
         );
         Optional<JwtSettingsValidationError> validationError = PrivateKeyJwtSigner.validate(signatureAlgorithm, privateKey);
         if (validationError.isPresent()) {
            throw new PropertyConstraintViolationException(
               "signingKeyAlias", config, environment, context, validationError.get().getMessage(), validationError.get().getCause().orElse(null)
            );
         }

         validator.matchingX509CertificateOrElseThrow("signingKeyAlias", Optional.of(signatureAlgorithm.getCertificateFilter()));
      }
   }

   @AfterConfigActivation
   void initSignatureVerificationAndSigner() {
      this.signer = PrivateKeyJwtSigner.builder(this.algorithm.getAlgorithm(), this.keystoreFile)
         .keyStorePassword(this.keystorePassword)
         .privateKeyAlias(this.signingKeyAlias)
         .privateKeyPassword(this.signingKeyPassword)
         .includeKid(this.isIncludeKid())
         .build();
      this.registerSigner();
   }

   public SignatureVerificationInformation getSignatureVerificationInformation() {
      return Objects.requireNonNull(this.signer, "AfterConfigActivation has not been called").getSignatureVerificationInformation().orElseThrow();
   }

   private void registerSigner() {
      JwksProvider.register(this);
   }

   @BeforePluginDeactivation
   void deregisterSigner() {
      JwksProvider.deregister(this);
   }

   @Override
   public List<SignatureVerificationInformation> getSigningCertificatesForJwks() {
      return Objects.requireNonNull(this.signer, "AfterConfigActivation has not been called").getSignatureVerificationInformation().stream().toList();
   }

   @SuppressFBWarnings(justification = "generated code")
   public PrivateKeyJwtSignatureConfig.Algorithm getAlgorithm() {
      return this.algorithm;
   }

   @SuppressFBWarnings(justification = "generated code")
   public File getKeystoreFile() {
      return this.keystoreFile;
   }

   @SuppressFBWarnings(justification = "generated code")
   public String getKeystorePassword() {
      return this.keystorePassword;
   }

   @SuppressFBWarnings(justification = "generated code")
   public String getSigningKeyAlias() {
      return this.signingKeyAlias;
   }

   @SuppressFBWarnings(justification = "generated code")
   public String getSigningKeyPassword() {
      return this.signingKeyPassword;
   }

   @SuppressFBWarnings(justification = "generated code")
   public PrivateKeyJwtSigner getSigner() {
      return this.signer;
   }

   @SuppressFBWarnings(justification = "generated code")
   public PrivateKeyJwtSignatureConfig() {
   }

   public enum Algorithm {
      RS256(JwsAlgorithmRsa.RS256),
      RS384(JwsAlgorithmRsa.RS384),
      RS512(JwsAlgorithmRsa.RS512),
      PS256(JwsAlgorithmRsa.PS256),
      PS384(JwsAlgorithmRsa.PS384),
      PS512(JwsAlgorithmRsa.PS512),
      ES256(JwsAlgorithmEc.ES256),
      ES384(JwsAlgorithmEc.ES384),
      ES512(JwsAlgorithmEc.ES512);

      private final SignatureAlgorithm algorithm;

      @SuppressFBWarnings(justification = "generated code")
      public SignatureAlgorithm getAlgorithm() {
         return this.algorithm;
      }

      @SuppressFBWarnings(justification = "generated code")
      Algorithm(final SignatureAlgorithm algorithm) {
         this.algorithm = algorithm;
      }
   }
}