package com.airlock.iam.security.util.crypto;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Optional;
import java.util.function.Predicate;
import javax.security.auth.x500.X500Principal;

public final class KeyStoreUtils {
   public static final Predicate<X509Certificate> RSA_PUBLIC_KEY_CERTIFICATE_FILTER = cert -> "RSA".equalsIgnoreCase(cert.getPublicKey().getAlgorithm());
   public static final Predicate<X509Certificate> EC_PUBLIC_KEY_CERTIFICATE_FILTER = cert -> "EC".equalsIgnoreCase(cert.getPublicKey().getAlgorithm());
   public static final Predicate<PrivateKey> RSA_PRIVATE_KEY_FILTER = key -> "RSA".equalsIgnoreCase(key.getAlgorithm());
   public static final Predicate<PrivateKey> EC_PRIVATE_KEY_FILTER = key -> "EC".equalsIgnoreCase(key.getAlgorithm());

   private KeyStoreUtils() {
   }

   public static KeyStore createKeyStoreFromFile(File keyStorePath, Optional<String> keyStorePassword, String keyStoreType) throws KeyStoreException {
      return loadKeyStore(keyStorePath, keyStorePassword, keyStoreType);
   }

   public static KeyStore createKeyStoreFromInputStream(InputStream keyStoreStream, String keyStoreName, Optional<String> keyStorePassword, String keyStoreType) throws KeyStoreException {
      return loadKeyStore(keyStoreStream, keyStoreName, keyStorePassword, keyStoreType);
   }

   public static KeyStore createKeyStoreFromFile(File keyStorePath, Optional<String> keyStorePassword, KeyStoreUtils.KeyStoreType keyStoreType) throws KeyStoreException {
      return keyStoreType == KeyStoreUtils.KeyStoreType.AUTO
         ? autoLoadKeyStore(keyStorePath, keyStorePassword)
         : loadKeyStore(keyStorePath, keyStorePassword, keyStoreType.name());
   }

   private static KeyStore autoLoadKeyStore(File keyStorePath, Optional<String> keyStorePassword) throws KeyStoreException {
      Preconditions.checkNotNull(keyStorePath);

      try {
         return loadKeyStore(keyStorePath, keyStorePassword, KeyStoreUtils.KeyStoreType.JKS.name());
      } catch (KeyStoreException jksException) {
         try {
            return loadKeyStore(keyStorePath, keyStorePassword, KeyStoreUtils.KeyStoreType.PKCS12.name());
         } catch (KeyStoreException pkcs12Exception) {
            if (!keyStorePath.getName().endsWith(".p12") && !keyStorePath.getName().endsWith(".pfx")) {
               throw jksException;
            } else {
               throw pkcs12Exception;
            }
         }
      }
   }

   private static KeyStore loadKeyStore(File keyStorePath, Optional<String> keyStorePassword, String keyStoreType) throws KeyStoreException {
      try {
         return loadKeyStore(Files.newInputStream(keyStorePath.toPath()), keyStorePath.getName(), keyStorePassword, keyStoreType);
      } catch (IOException e) {
         throw new KeyStoreException("IO error loading keystore: \"" + keyStorePath.getName() + "\"", e);
      }
   }

   private static KeyStore loadKeyStore(InputStream keyStoreStream, String keyStoreName, Optional<String> keyStorePassword, String keyStoreType) throws KeyStoreException {
      try {
         KeyStore keyStore = KeyStore.getInstance(keyStoreType);
         loadKeyStoreInternal(keyStoreStream, keyStorePassword, keyStore);
         return keyStore;
      } catch (KeyStoreException e) {
         throw new KeyStoreException("invalid keystore: \"" + keyStoreName + "\"", e);
      } catch (NoSuchAlgorithmException e) {
         throw new KeyStoreException("algorithm missing for keystore: \"" + keyStoreName + "\"", e);
      } catch (CertificateException e) {
         throw new KeyStoreException("cannot load certificate for alias: \"" + keyStoreName + "\"", e);
      } catch (IOException e) {
         throw new KeyStoreException("IO error loading keystore: \"" + keyStoreName + "\"", e);
      }
   }

   private static void loadKeyStoreInternal(InputStream keyStoreStream, Optional<String> keyStorePassword, KeyStore keyStore) throws IOException, NoSuchAlgorithmException, CertificateException {
      byte[] keyStoreBytes = keyStoreStream.readAllBytes();

      try {
         keyStore.load(new ByteArrayInputStream(keyStoreBytes), keyStorePassword.orElse("").toCharArray());
      } catch (IOException | NoSuchAlgorithmException | CertificateException e) {
         if (!keyStorePassword.isEmpty()) {
            throw e;
         }

         keyStore.load(new ByteArrayInputStream(keyStoreBytes), null);
      }
   }

   public static List<X509Certificate> getAllX509Certificates(KeyStore keyStore) throws KeyStoreException {
      Preconditions.checkNotNull(keyStore);
      List<X509Certificate> certsFromKeyStore = Lists.newArrayList();
      Enumeration<String> aliases = keyStore.aliases();

      while (aliases.hasMoreElements()) {
         getX509Certificate(keyStore, aliases.nextElement()).ifPresent(certsFromKeyStore::add);
      }

      return certsFromKeyStore;
   }

   public static Optional<X509Certificate> getX509Certificate(KeyStore keyStore, String certificateAlias) throws KeyStoreException {
      return Optional.ofNullable(keyStore.getCertificate(certificateAlias)).filter(X509Certificate.class::isInstance).map(X509Certificate.class::cast);
   }

   public static Optional<X509Certificate> getX509Certificate(KeyStore keyStore, X500Principal principal) throws KeyStoreException {
      for (String alias : Collections.list(keyStore.aliases())) {
         if (keyStore.getCertificate(alias) instanceof X509Certificate x509Certificate && x509Certificate.getSubjectX500Principal().equals(principal)) {
            return Optional.of(x509Certificate);
         }
      }

      return Optional.empty();
   }

   public static Optional<X509Certificate> getUniqueX509Certificate(
      KeyStore keyStore, Optional<String> certificateAlias, Optional<Predicate<X509Certificate>> filter
   ) throws KeyStoreException {
      Predicate<X509Certificate> certFilter = filter.orElse(x -> true);
      if (certificateAlias.isPresent()) {
         return getX509Certificate(keyStore, certificateAlias.get()).filter(certFilter);
      }

      List<X509Certificate> certificates = getAllX509Certificates(keyStore).stream().filter(certFilter).toList();
      return certificates.size() == 1 ? Optional.of(certificates.getFirst()) : Optional.empty();
   }

   public static List<PrivateKey> getAllPrivateKeys(KeyStore keyStore, Optional<String> privateKeyPassword) throws KeyStoreException {
      List<PrivateKey> privateKeys = Lists.newArrayList();
      Enumeration<String> aliases = keyStore.aliases();

      while (aliases.hasMoreElements()) {
         try {
            getPrivateKey(keyStore, aliases.nextElement(), privateKeyPassword).ifPresent(privateKeys::add);
         } catch (UnrecoverableKeyException var5) {
         }
      }

      return privateKeys;
   }

   public static Optional<PrivateKey> getPrivateKey(KeyStore keyStore, String privateKeyAlias, Optional<String> privateKeyPassword) throws UnrecoverableKeyException, KeyStoreException {
      try {
         return Optional.ofNullable(keyStore.getKey(privateKeyAlias, privateKeyPassword.orElse("").toCharArray()))
            .filter(PrivateKey.class::isInstance)
            .map(PrivateKey.class::cast);
      } catch (NoSuchAlgorithmException e) {
         throw new KeyStoreException("there is a problem with the keystore", e);
      }
   }

   public static Optional<PrivateKey> getUniquePrivateKey(
      KeyStore keyStore, Optional<String> privateKeyAlias, Optional<String> privateKeyPassword, Optional<Predicate<PrivateKey>> filter
   ) throws UnrecoverableKeyException, KeyStoreException {
      Predicate<PrivateKey> keyFilter = filter.orElse(x -> true);
      if (privateKeyAlias.isPresent()) {
         return getPrivateKey(keyStore, privateKeyAlias.get(), privateKeyPassword).filter(keyFilter);
      }

      List<PrivateKey> privateKeys = getAllPrivateKeys(keyStore, privateKeyPassword).stream().filter(keyFilter).toList();
      return privateKeys.size() == 1 ? Optional.of(privateKeys.getFirst()) : Optional.empty();
   }

   public enum KeyStoreType {
      JKS,
      PKCS12,
      AUTO;
   }
}