package com.atlassian.bitbucket.mesh.grpc;

import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.*;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import io.micrometer.core.instrument.Timer;

import static com.atlassian.bitbucket.mesh.metrics.MetricsConstants.*;

/**
 * {@link ServerInterceptor} that measures the time spent servicing RPC operations. For client-streaming RPC methods,
 * the number of received messages per RPC operation is also measured.
 */
public class MetricsServerInterceptor implements ServerInterceptor {

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
                                                                 ServerCallHandler<ReqT, RespT> next) {

        ServerCall.Listener<ReqT> listener = next.startCall(call, headers);
        return new MetricsServerCallListener<>(call.getMethodDescriptor(), listener);
    }

    private static class MetricsServerCallListener<ReqT>
            extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {

        private final Timer.Sample sample;
        private final boolean streamingClient;
        private final Tags tags;

        private int messagesReceived;

        protected MetricsServerCallListener(MethodDescriptor<ReqT, ?> methodDescriptor,
                                            ServerCall.Listener<ReqT> delegate) {
            super(delegate);

            sample = Timer.start();
            streamingClient = !methodDescriptor.getType().clientSendsOneMessage();
            String serviceName = methodDescriptor.getServiceName();
            String methodName = methodDescriptor.getBareMethodName();
            if (serviceName != null && methodName != null) {
                int index = serviceName.lastIndexOf('.');
                if (index != -1) {
                    serviceName = serviceName.substring(index + 1);
                }
                tags = Tags.of(Tag.of(TAG_SERVICE, serviceName), Tag.of(TAG_METHOD, methodName));
            } else {
                tags = Tags.empty();
            }

            Context.current().addListener(ignored -> reportMetrics(), MoreExecutors.directExecutor());
        }

        @Override
        public void onMessage(ReqT message) {
            ++messagesReceived;
            super.onMessage(message);
        }

        protected void reportMetrics() {
            sample.stop(Metrics.timer(TIMER_RPC_CALLS, tags));
            if (streamingClient) {
                //Only track how many messages were received if we have a streaming client
                Metrics.summary(SUMMARY_RPC_MESSAGES, tags).record(messagesReceived);
            }
        }
    }
}
