package com.atlassian.bitbucket.mesh.network;

import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.util.ConcurrencyUtils;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.ClientInterceptor;
import io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.NegotiationType;
import io.grpc.netty.NettyChannelBuilder;
import io.netty.handler.ssl.SslContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.event.EventListener;

import javax.annotation.Nonnull;
import javax.net.ssl.SSLException;
import java.net.URI;
import java.util.Set;
import java.util.concurrent.*;

import static com.atlassian.bitbucket.mesh.MeshConstants.EXECUTOR_RPC_CLIENT;
import static java.util.Objects.requireNonNull;

public class GrpcChannelRegistry implements ChannelRegistry {

    private static final Logger log = LoggerFactory.getLogger(GrpcChannelRegistry.class);

    private final ConcurrentMap<String, GrpcMeshChannel> channelsById;
    private final GrpcConfig config;
    private final ExecutorService executor;
    private final ClientInterceptor[] interceptors;
    private final SslContext sslContext;

    public GrpcChannelRegistry(GrpcConfig config, ClientInterceptor[] interceptors) {
        this.config = config;
        this.interceptors = interceptors;

        channelsById = new ConcurrentHashMap<>();
        executor = Executors.newCachedThreadPool(
                ConcurrencyUtils.createDaemonThreadFactory(EXECUTOR_RPC_CLIENT));
        sslContext = buildSslContext();
    }

    @Nonnull
    @Override
    public MeshChannel getChannel(@Nonnull MeshNode node) {
        requireNonNull(node, "node");

        return channelsById.computeIfAbsent(node.getId(),
                id -> new GrpcMeshChannel(node, channelBuilder(node), channel -> onChannelClosed(id, channel)));
    }

    @EventListener(TopologyUpdatedEvent.class)
    public void onTopologyUpdated() {
        clear();
    }

    public void shutdown() {
        clear();

        if (MoreExecutors.shutdownAndAwaitTermination(executor, 10L, TimeUnit.SECONDS)) {
            log.debug("The request executor has shutdown");
        } else {
            log.warn("Failed to shutdown the request executor; there may be active requests");
        }
    }

    void clear() {
        //Trying to copy and then clear the map would leave a window for a channel to racily be added between
        //the two. Instead, we create a copy of the map's keys and then loop over the map removing each. Each
        //channel we successfully remove is closed
        Set<String> nodeIds = ImmutableSet.copyOf(channelsById.keySet());
        for (String nodeId : nodeIds) {
            GrpcMeshChannel channel = channelsById.remove(nodeId);
            if (channel != null) {
                log.debug("{}: Closing channel", nodeId);
                channel.close();
            }
        }
    }

    private void onChannelClosed(String id, GrpcMeshChannel channel) {
        //Only remove the channel from the map if it's the exact same instance
        //noinspection ObjectEquality
        channelsById.computeIfPresent(id, (ignored, value) -> value == channel ? null : value);
    }

    private static SslContext buildSslContext() {
        try {
            return GrpcSslContexts.forClient().build();
        } catch (IllegalStateException | SSLException e) {
            log.error("Unable to build SSL context", e);
        }
        return null;
    }

    private NettyChannelBuilder channelBuilder(MeshNode node) {
        URI uri = URI.create(node.getRpcUrl());

        String host = uri.getHost();
        int port = uri.getPort();
        boolean tls = !"http".equals(uri.getScheme()); //Use TLS for any scheme other than explicit http://

        log.debug("{}: Creating gRPC channel for {}:{}", node.getId(), host, port);
        return NettyChannelBuilder.forAddress(host, port)
                .executor(executor)
                //Add the common interceptors
                .intercept(interceptors)
                .maxInboundMessageSize(config.getMaxMessageSize())
                .maxInboundMetadataSize(config.getMaxMetadataSize())
                .negotiationType(tls ? NegotiationType.TLS : NegotiationType.PLAINTEXT)
                //This executor handles certain blocking operations performed by gRPC itself (like DNS lookups,
                //just as an example). It may never get used, but if it does we don't want to use the default
                .offloadExecutor(executor)
                .sslContext(tls ? sslContext : null);
    }
}
