package ru.yandex.qe.ssl;

import java.security.PrivateKey;
import java.security.cert.PKIXBuilderParameters;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;

import com.google.common.base.Throwables;

import ru.yandex.qe.http.certificates.PemParser;

/**
 * @author rurikk
 */
public class SslSupport {
    private PrivateKey privateKey;
    private X509Certificate certificate;
    private List<X509Certificate> trustedCertificates = new ArrayList<>();
    private CrlFetcher crlFetcher;
    private Pattern peerCnPattern;

    public PrivateKey getPrivateKey() {
        return privateKey;
    }

    public X509Certificate getCertificate() {
        return certificate;
    }

    public SslSupport setPemKey(String privateKeyPem, String certificatePem) {
        privateKey = PemParser.parseRsa(privateKeyPem);
        certificate = PemParser.parseX509(certificatePem);
        return this;
    }

    public List<X509Certificate> getTrustedCertificates() {
        return trustedCertificates;
    }

    public SslSupport addTrusted(String certificatePems) {
        trustedCertificates.addAll(PemParser.parseX509s(certificatePems));
        return this;
    }

    public CrlFetcher getCrlFetcher() {
        return crlFetcher;
    }

    public SslSupport setCrlFetcher(CrlFetcher crlFetcher) {
        this.crlFetcher = crlFetcher;
        return this;
    }

    public Pattern getPeerCnPattern() {
        return peerCnPattern;
    }

    public SslSupport setPeerCnPattern(String cnPattern) {
        this.peerCnPattern = Pattern.compile(cnPattern);
        return this;
    }

    // bundle format used by certificator:
    // <private key>
    // <certificate>
    // <CA chain>
    public SslSupport setCertificateBundle(String bundleContents) {
        List<Object> items = PemParser.parseAll(bundleContents);
        privateKey = items.stream()
                .filter(PrivateKey.class::isInstance)
                .map(PrivateKey.class::cast)
                .findFirst().orElseThrow(() -> new IllegalArgumentException("Key not found"));
        List<X509Certificate> certificates = items.stream()
                .filter(X509Certificate.class::isInstance)
                .map(X509Certificate.class::cast)
                .collect(Collectors.toList());

        // Own certificate and CA certificate
        if (certificates.size() < 2) {
            String message = String.format("Expected at least 2 certificates, found %s", certificates.size());
            throw new IllegalArgumentException(message);
        }
        this.certificate = certificates.get(0);
        trustedCertificates = certificates.subList(1, certificates.size());
        return this;
    }

    public KeyManagerFactory keyManagerFactory() {
        try {
            return SslUtils.keyManagerFactory(privateKey, certificate);
        } catch (Exception e) {
            throw Throwables.propagate(e);
        }
    }

    public TrustManagerFactory trustManagerFactory() {
        try {
            return SslUtils.trustManagerFactory(trustedCertificates, this::addVerifiers);
        } catch (Exception e) {
            throw Throwables.propagate(e);
        }
    }

    public SSLContext buildContext() {
        if (crlFetcher != null) {
            crlFetcher.setTrustedCertificates(trustedCertificates);
        }
        return SslUtils.sslContext(keyManagerFactory(), trustManagerFactory());
    }

    public void stopCrlFetcher() {
        if (crlFetcher != null) {
            crlFetcher.stop();
        }
    }

    private void addVerifiers(PKIXBuilderParameters params) {
        addRevocationCheck(params);
        addPatternCheck(params);
    }

    private void addRevocationCheck(PKIXBuilderParameters params) {
        if (crlFetcher != null) {
            params.setRevocationEnabled(true);
            params.addCertStore(SslUtils.fetchingCertStore(crlFetcher));
        } else {
            params.setRevocationEnabled(false);
        }
    }

    private void addPatternCheck(PKIXBuilderParameters params) {
        if (peerCnPattern != null) {
            params.addCertPathChecker(new CnPatternChecker(peerCnPattern));
        }
    }
}
