/*
 * Copyright (C) 2012 Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.airlift.drift.transport.apache.client;

import com.google.common.net.HostAndPort;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.ListeningScheduledExecutorService;
import io.airlift.drift.transport.client.MethodInvoker;
import io.airlift.drift.transport.client.MethodInvokerFactory;
import io.airlift.security.pem.PemReader;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TTransportFactory;

import javax.annotation.PreDestroy;
import javax.inject.Inject;
import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.util.Optional;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
import static com.google.common.util.concurrent.MoreExecutors.shutdownAndAwaitTermination;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static java.lang.Math.toIntExact;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.Executors.newCachedThreadPool;
import static java.util.concurrent.Executors.newScheduledThreadPool;
import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;

public class ApacheThriftMethodInvokerFactory<I>
        implements MethodInvokerFactory<I>, Closeable
{
    private final Function<I, ApacheThriftClientConfig> clientConfigurationProvider;

    private final ListeningExecutorService executorService;
    private final ListeningScheduledExecutorService delayService;
    private final HostAndPort defaultSocksProxy;

    public static ApacheThriftMethodInvokerFactory<?> createStaticApacheThriftMethodInvokerFactory(ApacheThriftClientConfig clientConfig)
    {
        return createStaticApacheThriftMethodInvokerFactory(clientConfig, new ApacheThriftConnectionFactoryConfig());
    }

    public static ApacheThriftMethodInvokerFactory<?> createStaticApacheThriftMethodInvokerFactory(
            ApacheThriftClientConfig clientConfig,
            ApacheThriftConnectionFactoryConfig factoryConfig)
    {
        requireNonNull(clientConfig, "clientConfig is null");
        return new ApacheThriftMethodInvokerFactory<>(factoryConfig, clientIdentity -> clientConfig);
    }

    @Inject
    public ApacheThriftMethodInvokerFactory(ApacheThriftConnectionFactoryConfig factoryConfig, Function<I, ApacheThriftClientConfig> clientConfigurationProvider)
    {
        requireNonNull(factoryConfig, "factoryConfig is null");

        ThreadFactory threadFactory = daemonThreadsNamed("drift-client-%s");
        if (factoryConfig.getThreadCount() == null) {
            executorService = listeningDecorator(newCachedThreadPool(threadFactory));
            delayService = listeningDecorator(newSingleThreadScheduledExecutor(daemonThreadsNamed("drift-client-delay-%s")));
        }
        else {
            delayService = listeningDecorator(newScheduledThreadPool(factoryConfig.getThreadCount(), threadFactory));
            executorService = delayService;
        }
        this.clientConfigurationProvider = requireNonNull(clientConfigurationProvider, "clientConfigurationProvider is null");
        this.defaultSocksProxy = factoryConfig.getSocksProxy();
    }

    @Override
    public MethodInvoker createMethodInvoker(I clientIdentity)
    {
        ApacheThriftClientConfig config = clientConfigurationProvider.apply(clientIdentity);
        if (config.getSocksProxy() == null) {
            config.setSocksProxy(defaultSocksProxy);
        }

        TTransportFactory transportFactory;
        switch (config.getTransport()) {
            case UNFRAMED:
                transportFactory = new TTransportFactory();
                break;
            case FRAMED:
                transportFactory = new TFramedTransport.Factory(toIntExact(config.getMaxFrameSize().toBytes()));
                break;
            default:
                throw new IllegalArgumentException("Unknown transport: " + config.getTransport());
        }

        TProtocolFactory protocolFactory;
        switch (config.getProtocol()) {
            case BINARY:
                protocolFactory = new TBinaryProtocol.Factory();
                break;
            case COMPACT:
                protocolFactory = new TCompactProtocol.Factory();
                break;
            default:
                throw new IllegalArgumentException("Unknown protocol: " + config.getProtocol());
        }

        Optional<SSLContext> sslContext = Optional.empty();
        if (config.isSslEnabled()) {
            sslContext = Optional.of(createSslContext(config));
        }

        return new ApacheThriftMethodInvoker(
                executorService,
                delayService,
                transportFactory,
                protocolFactory,
                config.getConnectTimeout(),
                config.getRequestTimeout(),
                Optional.ofNullable(config.getSocksProxy()),
                sslContext);
    }

    private static SSLContext createSslContext(ApacheThriftClientConfig config)
    {
        try {
            KeyStore trustStore = PemReader.loadTrustStore(config.getTrustCertificate());
            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            trustManagerFactory.init(trustStore);

            KeyManager[] keyManagers = null;
            if (config.getKey() != null) {
                Optional<String> keyPassword = Optional.ofNullable(config.getKeyPassword());
                KeyStore keyStore = PemReader.loadKeyStore(config.getTrustCertificate(), config.getKey(), keyPassword);
                KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
                keyManagerFactory.init(keyStore, new char[0]);
                keyManagers = keyManagerFactory.getKeyManagers();
            }

            SSLContext sslContext = SSLContext.getInstance("TLS");
            sslContext.init(keyManagers, trustManagerFactory.getTrustManagers(), null);
            return sslContext;
        }
        catch (IOException | GeneralSecurityException e) {
            throw new IllegalArgumentException("Unable to load SSL keys", e);
        }
    }

    @PreDestroy
    @Override
    public void close()
    {
        shutdownAndAwaitTermination(executorService, 5, TimeUnit.MINUTES);
    }
}
