package com.atlassian.bitbucket.mesh.network;

import io.grpc.Channel;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.stub.AbstractStub;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Function;

import static java.util.Objects.requireNonNull;

class GrpcMeshChannel implements MeshChannel {

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

    private final ManagedChannel channel;
    private final AtomicBoolean closed;
    private final Consumer<GrpcMeshChannel> closeListener;
    private final MeshNode node;

    GrpcMeshChannel(@Nonnull MeshNode node, @Nonnull ManagedChannelBuilder<?> channelBuilder,
                    @Nullable Consumer<GrpcMeshChannel> closeListener) {
        this.closeListener = closeListener;
        this.node = requireNonNull(node, "node");

        channel = requireNonNull(channelBuilder, "channelBuilder")
                .build();
        closed = new AtomicBoolean();
    }

    @Override
    public boolean close() {
        if (closed.compareAndSet(false, true)) {
            if (closeListener != null) {
                try {
                    closeListener.accept(this);
                } catch (RuntimeException e) {
                    log.warn("Close listener failed for {} channel. Proceeding to shut down the channel", node, e);
                }
            }

            //This intentionally does not do a shutdownNow, to allow any in-flight RPCs over the channel
            //to complete cleanly. After this is called the channel won't accept any new RPCs
            channel.shutdown();

            return true;
        }
        return false;
    }

    @Override
    public <S extends AbstractStub<S>> S newStub(@Nonnull Function<Channel, S> stubFactory) {
        return requireNonNull(stubFactory, "stubFactory").apply(channel)
                .withOption(OPT_CHANNEL, this)
                .withOption(OPT_TARGET_NODE, node);
    }

    @Override
    public String toString() {
        return node.getId() + " (" + node.getRpcUrl() + ")";
    }
}
