package com.airlock.iam.factor.infrastructure.service.fido;

import com.airlock.iam.base.misc.result.Result;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.core.misc.util.ExceptionUtils;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.elog.api.LogTypes;
import com.airlock.iam.factor.domain.model.fido.CertificateChainVerifier;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.SetMultimap;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SignatureException;
import java.security.cert.CertificateException;
import java.security.cert.CertificateExpiredException;
import java.security.cert.CertificateNotYetValidException;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import javax.security.auth.x500.X500Principal;
import lombok.NonNull;
import org.apache.commons.lang3.Validate;
import org.bouncycastle.asn1.x500.X500Name;

public final class CertificateChainVerifierImpl implements CertificateChainVerifier {
   private static final ELogger LOG = ELoggerFactory.getLogger(CertificateChainVerifierImpl.class);
   private final SetMultimap<X500Name, X509Certificate> trustedCertificates;
   private final TimeProvider timeProvider;
   private final int maximumCertificateChainLength;

   private CertificateChainVerifierImpl(
      @NonNull Set<X509Certificate> trustedCertificates, @NonNull TimeProvider timeProvider, int maximumCertificateChainLength
   ) {
      if (trustedCertificates == null) {
         throw new NullPointerException("trustedCertificates is marked non-null but is null");
      }

      if (timeProvider == null) {
         throw new NullPointerException("timeProvider is marked non-null but is null");
      }

      Validate.isTrue(maximumCertificateChainLength >= 1, "maximum certificate chain length must be at least 1 but was '%d'", maximumCertificateChainLength);
      this.trustedCertificates = trustedCertificates.stream()
         .collect(ImmutableSetMultimap.toImmutableSetMultimap(cert -> new X500Name(cert.getSubjectX500Principal().getName()), cert -> cert));
      this.timeProvider = timeProvider;
      this.maximumCertificateChainLength = maximumCertificateChainLength;
   }

   static CertificateChainVerifier certificateChainVerifier(
      Set<X509Certificate> trustedCertificates, TimeProvider timeProvider, int maximumCertificateChainLength
   ) {
      return new CertificateChainVerifierImpl(trustedCertificates, timeProvider, maximumCertificateChainLength);
   }

   @Override
   public Result<X509Certificate, String> findTrustAnchor(List<X509Certificate> certificateChain) {
      if (certificateChain.isEmpty()) {
         return Result.failure("Verification of certificate chain failed because it does not contain any certificate");
      }

      if (certificateChain.size() > this.maximumCertificateChainLength) {
         return Result.failure(
            String.format(
               "Verification of certificate chain failed because it is too long. Expected certificate chain to contain at most '%d' certificates, but got '%d'",
               this.maximumCertificateChainLength,
               certificateChain.size()
            )
         );
      }

      if (this.trustedCertificates.isEmpty()) {
         return Result.failure("Verification of certificate chain failed because no certificates are trusted.");
      }

      Optional<String> certificateChainBroken = this.verifyCertificateChain(certificateChain);
      return certificateChainBroken.isPresent() ? Result.failure(certificateChainBroken.get()) : this.findTrustedCertificate(certificateChain);
   }

   private Optional<String> verifyCertificateChain(List<X509Certificate> certificateChain) {
      for (int i = 0; i < certificateChain.size(); i++) {
         X509Certificate currentCertificate = certificateChain.get(i);
         Optional<X509Certificate> nextCertificate = i < certificateChain.size() - 1 ? Optional.of(certificateChain.get(i + 1)) : Optional.empty();
         Optional<String> failure = this.verifyCertificateValidity(currentCertificate)
            .or(() -> this.verifySignatureWithNextCertificate(currentCertificate, nextCertificate));
         if (failure.isPresent()) {
            return failure;
         }
      }

      return Optional.empty();
   }

   private Optional<String> verifySignatureWithNextCertificate(X509Certificate currentCertificate, Optional<X509Certificate> nextCertificate) {
      return nextCertificate.isPresent() && !this.isSigned(currentCertificate, nextCertificate.get())
         ? Optional.of(
            String.format(
               "Verification of certification chain failed because chain is broken! Certificate %s was not signed by next certificate in the chain %s",
               this.log(currentCertificate),
               nextCertificate.get()
            )
         )
         : Optional.empty();
   }

   private Result<X509Certificate, String> findTrustedCertificate(List<X509Certificate> certificateChain) {
      for (X509Certificate certificate : certificateChain) {
         Optional<X509Certificate> trustedCertificate = this.isSignedByOneOf(certificate, this.validTrustedCertificates(certificate.getIssuerX500Principal()));
         if (trustedCertificate.isPresent()) {
            LOG.info(String.format("Certificate chain is valid and is derived from trust anchor %s", this.log(trustedCertificate.get())));
            return Result.result(trustedCertificate.get());
         }
      }

      return Result.failure(
         "Verification of certification chain failed because none of the certificates in the chain was signed by a valid trusted certificate"
      );
   }

   private Set<X509Certificate> validTrustedCertificates(X500Principal issuer) {
      Map<Boolean, Set<X509Certificate>> certificates = this.trustedCertificates
         .get(new X500Name(issuer.getName()))
         .stream()
         .collect(Collectors.partitioningBy(this::isValid, Collectors.toUnmodifiableSet()));
      certificates.get(false)
         .forEach(
            invalidCertificate -> LOG.warn(LogTypes.CONFIGURATION_ERROR, String.format("Trusted certificate %s is not valid", this.log(invalidCertificate)))
         );
      return certificates.get(true);
   }

   private boolean isValid(X509Certificate certificate) {
      return this.verifyCertificateValidity(certificate).isEmpty();
   }

   private Optional<String> verifyCertificateValidity(X509Certificate certificate) {
      try {
         certificate.checkValidity(this.timeProvider.getCurrentTime());
         return Optional.empty();
      } catch (CertificateExpiredException | CertificateNotYetValidException e) {
         return Optional.of(String.format("Certificate %s is not valid. Cause: %s", this.log(certificate), ExceptionUtils.getNestedThrowableString(e)));
      }
   }

   private Optional<X509Certificate> isSignedByOneOf(X509Certificate certificate, Set<X509Certificate> intermediateCertificates) {
      return intermediateCertificates.stream().filter(intermediateCertificate -> this.isSigned(certificate, intermediateCertificate)).findAny();
   }

   private boolean isSigned(X509Certificate certificate, X509Certificate intermediateCertificate) {
      try {
         certificate.verify(intermediateCertificate.getPublicKey());
         return true;
      } catch (CertificateException | NoSuchAlgorithmException | InvalidKeyException | NoSuchProviderException e) {
         LOG.debug(String.format("Signature of %s could not be verified. Cause: %s", this.log(certificate), ExceptionUtils.getNestedThrowableString(e)));
         return false;
      } catch (SignatureException e) {
         LOG.debug(String.format("Certificate %s was not signed by %s", this.log(certificate), this.log(intermediateCertificate)));
         return false;
      }
   }

   private String log(X509Certificate certificate) {
      return String.format(
         "X509Certificate[Issuer: '%s', Subject: '%s', Serial Number: '%d']",
         certificate.getIssuerX500Principal().getName(),
         certificate.getSubjectX500Principal().getName(),
         certificate.getSerialNumber()
      );
   }
}