package com.airlock.iam.common.infrastructure.service;

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.common.domain.model.jwt.SignatureVerificationInformation;
import com.airlock.iam.common.domain.model.jwt.signature.KeyId;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.Curve;
import com.nimbusds.jose.jwk.ECKey;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.KeyUse;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.ECKey.Builder;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.security.cert.X509Certificate;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.util.List;

public final class JWKSetExtractor {
   private final List<SignatureVerificationInformation> signatureVerificationInformation;

   public JWKSet getJwkSet() {
      return new JWKSet(this.signatureVerificationInformation.stream().map(this::jwkFromCertificate).distinct().toList());
   }

   private JWK jwkFromCertificate(SignatureVerificationInformation verificationInformation) {
      X509Certificate certificate = verificationInformation.getCertificate();
      if (certificate.getPublicKey() instanceof RSAPublicKey) {
         return createRSAKey(verificationInformation);
      } else if (certificate.getPublicKey() instanceof ECPublicKey) {
         return createECKey(verificationInformation);
      } else {
         throw new IllegalArgumentException("Invalid certificate to verify ID Tokens with (subject name: " + subjectName(certificate) + ").");
      }
   }

   private static String subjectName(X509Certificate certificate) {
      return LogUtil.logSafeQuote(certificate.getSubjectX500Principal().getName());
   }

   private static ECKey createECKey(SignatureVerificationInformation verificationInformation) {
      X509Certificate certificate = verificationInformation.getCertificate();
      ECPublicKey publicKey = (ECPublicKey)certificate.getPublicKey();
      Curve curve = Curve.forECParameterSpec(publicKey.getParams());
      if (curve == null) {
         throw new IllegalArgumentException(
            "Could not determine EC curve from the public key (subject name: " + subjectName(certificate) + "). Curve may not be supported."
         );
      } else {
         return new Builder(curve, publicKey)
            .keyUse(KeyUse.SIGNATURE)
            .algorithm(jwsAlgorithm(verificationInformation))
            .keyID(KeyId.keyId(certificate).keyId())
            .build();
      }
   }

   private static RSAKey createRSAKey(SignatureVerificationInformation verificationInformation) {
      X509Certificate certificate = verificationInformation.getCertificate();
      RSAPublicKey publicKey = (RSAPublicKey)certificate.getPublicKey();
      return new com.nimbusds.jose.jwk.RSAKey.Builder(publicKey)
         .keyUse(KeyUse.SIGNATURE)
         .algorithm(jwsAlgorithm(verificationInformation))
         .keyID(KeyId.keyId(certificate).keyId())
         .build();
   }

   private static JWSAlgorithm jwsAlgorithm(SignatureVerificationInformation verificationInformation) {
      return JWSAlgorithm.parse(verificationInformation.getAlgorithm().name());
   }

   @SuppressFBWarnings(justification = "generated code")
   public JWKSetExtractor(List<SignatureVerificationInformation> signatureVerificationInformation) {
      this.signatureVerificationInformation = signatureVerificationInformation;
   }
}