package com.sun.identity.saml2.key;

import com.airlock.iam.opensso.jaxb.entityconfig.BaseConfigType;
import com.airlock.iam.opensso.jaxb.metadata.KeyDescriptorType;
import com.airlock.iam.opensso.jaxb.metadata.KeyTypes;
import com.airlock.iam.opensso.jaxb.metadata.RoleDescriptorType;
import com.airlock.iam.opensso.jaxb.xmlenc.EncryptionMethodType;
import com.airlock.iam.opensso.jaxb.xmlsig.KeyInfoType;
import com.airlock.iam.opensso.jaxb.xmlsig.X509DataElement;
import com.airlock.iam.opensso.jaxb.xmlsig.X509DataType;
import com.google.common.annotations.VisibleForTesting;
import com.sun.identity.common.SystemConfigurationUtil;
import com.sun.identity.saml.xmlsig.KeyProvider;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.meta.SAML2MetaUtils;
import jakarta.annotation.Nullable;
import java.io.ByteArrayInputStream;
import java.lang.reflect.InvocationTargetException;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Collections;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;

public class KeyUtil {
   private static KeyProvider kp;
   private static final Hashtable<String, EncInfo> ENC_INFO = new Hashtable<>();
   private static final Hashtable<String, List<X509Certificate>> SIG_CERTIFICATES = new Hashtable<>();

   private KeyUtil() {
   }

   public static KeyProvider getKeyProviderInstance() {
      return kp;
   }

   public static PrivateKey getDecryptionKey(BaseConfigType baseConfig) {
      Map<String, List<String>> map = SAML2MetaUtils.getAttributes(baseConfig);
      List<String> list = map.get("encryptionCertAlias");
      PrivateKey decryptionKey = null;
      if (list != null && !list.isEmpty()) {
         String alias = list.getFirst();
         if (alias != null && !alias.isEmpty() && kp != null) {
            decryptionKey = kp.getPrivateKey(alias);
         }
      }

      return decryptionKey;
   }

   public static List<X509Certificate> getVerificationCerts(@Nullable RoleDescriptorType roleDescriptorType, String entityID, String role) {
      String classMethod = "KeyUtil.getVerificationCert: ";
      String index = entityID.trim() + "|" + role;
      List<X509Certificate> certs = SIG_CERTIFICATES.get(index);
      if (certs != null) {
         return certs;
      } else if (roleDescriptorType == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Null RoleDescriptorType input for entityID=" + entityID + " in " + role + " role.");
         return Collections.emptyList();
      } else {
         List<KeyDescriptorType> kd = getKeyDescriptors(roleDescriptorType, KeyTypes.SIGNING);
         if (kd.isEmpty()) {
            SAML2SDKUtils.DEBUG.error(classMethod + "No signing KeyDescriptor for entityID=" + entityID + " in " + role + " role.");
            return Collections.emptyList();
         } else {
            certs = kd.stream().map(KeyUtil::getCert).filter(Objects::nonNull).toList();
            if (certs.isEmpty()) {
               SAML2SDKUtils.DEBUG.error(classMethod + "No signing cert for entityID=" + entityID + " in " + role + " role.");
               return Collections.emptyList();
            } else {
               SIG_CERTIFICATES.put(index, certs);
               return certs;
            }
         }
      }
   }

   public static EncInfo getEncInfo(RoleDescriptorType roled, String entityID, String role) {
      String classMethod = "KeyUtil.getEncInfo: ";
      if (SAML2SDKUtils.DEBUG.messageEnabled()) {
         SAML2SDKUtils.DEBUG.message(classMethod + "Entering... \nEntityID=" + entityID + "\nRole=" + role);
      }

      String index = entityID.trim() + "|" + role;
      EncInfo encInfo = ENC_INFO.get(index);
      if (encInfo != null) {
         return encInfo;
      }

      if (roled == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Null RoleDescriptorType input for entityID=" + entityID + " in " + role + " role.");
         return null;
      }

      KeyDescriptorType kd = getKeyDescriptor(roled, KeyTypes.ENCRYPTION);
      if (kd == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "No encryption KeyDescriptor for entityID=" + entityID + " in " + role + " role.");
         return null;
      }

      X509Certificate cert = getCert(kd);
      if (cert == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "No encryption cert for entityID=" + entityID + " in " + role + " role.");
         return null;
      }

      List<EncryptionMethodType> emList = kd.getEncryptionMethod();
      String algorithm = null;
      int keySize = 0;
      if (emList != null && !emList.isEmpty()) {
         EncryptionMethodType em = emList.getFirst();
         if (em != null) {
            algorithm = em.getAlgorithm();
            List<?> cList = em.getContent();
            if (cList != null) {
               for (Object cObject : cList) {
                  if (cObject instanceof EncryptionMethodType.KeySize keySizeValue) {
                     keySize = ((BigInteger)keySizeValue.getValue()).intValue();
                     break;
                  }
               }
            }
         }
      }

      if (algorithm == null || algorithm.isEmpty()) {
         algorithm = "http://www.w3.org/2001/04/xmlenc#aes128-cbc";
         keySize = 128;
      }

      PublicKey pk = cert.getPublicKey();
      if (pk != null) {
         encInfo = new EncInfo(pk, algorithm, keySize);
      }

      if (encInfo != null) {
         ENC_INFO.put(index, encInfo);
      }

      return encInfo;
   }

   public static KeyDescriptorType getKeyDescriptor(RoleDescriptorType roled, KeyTypes usage) {
      List<KeyDescriptorType> keyDescriptors = getKeyDescriptors(roled, usage);
      return keyDescriptors.isEmpty() ? null : keyDescriptors.getFirst();
   }

   private static List<KeyDescriptorType> getKeyDescriptors(RoleDescriptorType roleDescriptorType, KeyTypes usage) {
      List<KeyDescriptorType> keyDescriptors = roleDescriptorType.getKeyDescriptor();
      List<KeyDescriptorType> keyDescriptorsWithCorrectUsage = keyDescriptors.stream()
         .filter(kd -> getUsage(kd).toLowerCase(Locale.ROOT).equals(usage.value()))
         .toList();
      return !keyDescriptorsWithCorrectUsage.isEmpty() ? keyDescriptorsWithCorrectUsage : keyDescriptorsWithoutUsages(keyDescriptors);
   }

   private static List<KeyDescriptorType> keyDescriptorsWithoutUsages(List<KeyDescriptorType> keyDescriptors) {
      return keyDescriptors.stream().filter(kd -> getUsage(kd).isEmpty()).toList();
   }

   private static String getUsage(KeyDescriptorType kd) {
      return Optional.ofNullable(kd.getUse()).map(KeyTypes::value).map(String::trim).orElse("");
   }

   public static X509Certificate getCert(KeyDescriptorType kd) {
      String classMethod = "KeyUtil.getCert: ";
      KeyInfoType ki = kd.getKeyInfo();
      if (ki == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "No KeyInfo.");
         return null;
      }

      Iterator<?> it = ki.getContent().iterator();
      X509DataElement data = null;

      while (data == null && it.hasNext()) {
         Object content = it.next();
         if (content instanceof X509DataElement) {
            data = (X509DataElement)content;
         }
      }

      if (data == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "No X509DataElement.");
         return null;
      }

      it = ((X509DataType)data.getValue()).getX509IssuerSerialOrX509SKIOrX509SubjectName().iterator();
      X509DataType.X509Certificate cert = null;

      while (cert == null && it.hasNext()) {
         Object content = it.next();
         if (content instanceof X509DataType.X509Certificate) {
            cert = (X509DataType.X509Certificate)content;
         }
      }

      if (cert == null) {
         SAML2SDKUtils.DEBUG.error(classMethod + "No X509Certificate.");
         return null;
      }

      byte[] bt = (byte[])cert.getValue();

      CertificateFactory cf;
      try {
         cf = CertificateFactory.getInstance("X.509");
      } catch (CertificateException ce) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Unable to get CertificateFactory for X.509 type", ce);
         return null;
      }

      ByteArrayInputStream bais = new ByteArrayInputStream(bt);
      X509Certificate retCert = null;

      try {
         while (bais.available() > 0) {
            retCert = (X509Certificate)cf.generateCertificate(bais);
         }

         return retCert;
      } catch (CertificateException ce) {
         SAML2SDKUtils.DEBUG.error(classMethod + "Unable to generate certificate from byte array input stream.", ce);
         return null;
      }
   }

   @VisibleForTesting
   static Hashtable<String, List<X509Certificate>> getCertificateCache() {
      return SIG_CERTIFICATES;
   }

   static {
      try {
         kp = (KeyProvider)Class.forName(
               SystemConfigurationUtil.getProperty("com.sun.identity.saml.xmlsig.keyprovider.class", "com.sun.identity.saml.xmlsig.JKSKeyProvider")
            )
            .getDeclaredConstructor()
            .newInstance();
      } catch (ClassNotFoundException cnfe) {
         SAML2SDKUtils.DEBUG.error("KeyUtil static block: Couldn't find the class.", cnfe);
         kp = null;
      } catch (InstantiationException | InvocationTargetException | NoSuchMethodException ie) {
         SAML2SDKUtils.DEBUG.error("KeyUtil static block: Couldn't instantiate the key provider instance.", ie);
         kp = null;
      } catch (IllegalAccessException iae) {
         SAML2SDKUtils.DEBUG.error("KeyUtil static block: Couldn't access the default constructor.", iae);
         kp = null;
      }
   }
}