package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.network.MeshChannel;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcError;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcErrorUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcWrongTarget;
import io.grpc.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.atlassian.bitbucket.mesh.network.MeshChannel.OPT_CHANNEL;
import static com.atlassian.bitbucket.mesh.network.MeshChannel.OPT_TARGET_NODE;

/**
 * Interceptor that handles a variety of errors reported by a Mesh node:
 * <ul>
 *     <li>closes an existing gRPC {@link ManagedChannel} if the server in question reports to have a different ID than
 *     we expected.</li>
 * </ul>
 */
public class ErrorHandlingClientInterceptor implements ClientInterceptor {

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

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
                                                               CallOptions callOptions, Channel next) {
        ClientCall<ReqT, RespT> delegate = next.newCall(method, callOptions);

        MeshChannel channel = callOptions.getOption(OPT_CHANNEL);
        MeshNode target = callOptions.getOption(OPT_TARGET_NODE);
        if (channel == null || target == null) {
            return delegate;
        }
        return new ErrorHandlingCall<>(delegate, channel, target);
    }

    private void onError(MeshChannel channel, MeshNode node, RpcError error) {
        if (error.hasWrongTarget()) {
            onWrongTarget(channel, node, error.getWrongTarget());
        }
    }

    private void onWrongTarget(MeshChannel channel, MeshNode node, RpcWrongTarget error) {
        //The RPC was sent to a different node than we expected; this can happen if a node comes online, and grpc
        //resolves the hostname to a stale IP address that happens to now be used by another Mesh node. This can
        //happen when docker-compose or kubernetes restarts a number of Mesh services and shuffle the IP allocation
        //around.
        if (channel != null && channel.close()) {
            log.warn("Closed mis-configured channel for {}. The Mesh node at the resolved address has ID {} " +
                    "instead of {}. Has the IP address changed?", node, error.getActualId(), error.getTargetId());
        }
    }

    private class ErrorHandlingCall<ReqT, RespT> extends ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT> {

        private final MeshChannel channel;
        private final MeshNode target;

        ErrorHandlingCall(ClientCall<ReqT, RespT> delegate, MeshChannel channel, MeshNode target) {
            super(delegate);

            this.channel = channel;
            this.target = target;
        }

        @Override
        public void start(Listener<RespT> listener, Metadata headers) {
            delegate().start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(listener) {

                @Override
                public void onClose(Status status, Metadata trailers) {
                    try {
                        if (!status.isOk()) {
                            RpcError error = RpcErrorUtils.fromMetadata(trailers);
                            if (error != null) {
                                onError(channel, target, error);
                            }
                        }
                    } catch (RuntimeException e) {
                        log.warn("Failed to handle RPC error", e);
                    }
                    delegate().onClose(status, trailers);
                }
            }, headers);
        }
    }
}