package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyMarshaller;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyServerCallListener;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.protobuf.MessageLite;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.protobuf.lite.ProtoLiteUtils;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import static org.hamcrest.CoreMatchers.not;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertNotNull;
import static org.mockito.Answers.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class ZeroCopyServerInterceptorTest {

    @Mock
    private GrpcConfig config;

    @SuppressWarnings("unchecked")
    @Test
    public void testDoesNotOverrideRequestMarshallerForMethodOptionWhenDisabled() {
        when(config.isZeroCopy()).thenReturn(false);

        //Create a copy of the original method descriptor so we don't dirty the (static) original in the test
        MethodDescriptor<RpcLfsReadObjectRequest, RpcStreamChunk> descriptor = methodDescriptor(
                LfsServiceGrpc.getReadObjectMethod(),
                RpcLfsReadObjectRequest.getDefaultInstance(),
                RpcStreamChunk.getDefaultInstance());
        ServerCall<RpcLfsReadObjectRequest, RpcStreamChunk> call = mock(ServerCall.class);
        lenient().when(call.getMethodDescriptor()).thenReturn(descriptor);

        ZeroCopyServerInterceptor interceptor = new ZeroCopyServerInterceptor(config);
        ServerCall.Listener<RpcLfsReadObjectRequest> listener =
                interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));

        assertNotNull(listener);
        assertThat(listener, not(instanceOf(ZeroCopyServerCallListener.class)));
        assertThat(descriptor.getRequestMarshaller(), not(instanceOf(ZeroCopyMarshaller.class)));

        //Verify that the zeroCopy setting can be changed at runtime
        when(config.isZeroCopy()).thenReturn(true);

        //Mock another call and verify that it _is_ intercepted
        listener = interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));
        assertNotNull(listener);
        assertThat(listener, instanceOf(ZeroCopyServerCallListener.class));
        assertThat(descriptor.getRequestMarshaller(), instanceOf(ZeroCopyMarshaller.class));
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testDoesNotOverrideRequestMarshallerForServiceOptionWhenDisabled() {
        when(config.isZeroCopy()).thenReturn(false);

        //Create a copy of the original method descriptor so we don't dirty the (static) original in the test
        MethodDescriptor<RpcHttpBackendFragment, RpcHostingWriteStreamFragment> descriptor = methodDescriptor(
                HostingServiceGrpc.getHttpBackendMethod(),
                RpcHttpBackendFragment.getDefaultInstance(),
                RpcHostingWriteStreamFragment.getDefaultInstance());
        ServerCall<RpcHttpBackendFragment, RpcHostingWriteStreamFragment> call = mock(ServerCall.class);
        lenient().when(call.getMethodDescriptor()).thenReturn(descriptor);

        ZeroCopyServerInterceptor interceptor = new ZeroCopyServerInterceptor(config);
        ServerCall.Listener<RpcHttpBackendFragment> listener =
                interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));

        assertNotNull(listener);
        assertThat(listener, not(instanceOf(ZeroCopyServerCallListener.class)));
        assertThat(descriptor.getRequestMarshaller(), not(instanceOf(ZeroCopyMarshaller.class)));

        //Verify that the zeroCopy setting can be changed at runtime
        when(config.isZeroCopy()).thenReturn(true);

        //Mock another call and verify that it _is_ intercepted
        listener = interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));
        assertNotNull(listener);
        assertThat(listener, instanceOf(ZeroCopyServerCallListener.class));
        assertThat(descriptor.getRequestMarshaller(), instanceOf(ZeroCopyMarshaller.class));
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testDoesNotOverrideRequestMarshallerIfNoMethodOrServiceOptionsConfigured() {
        when(config.isZeroCopy()).thenReturn(true);

        //Create a copy of the original method descriptor so we don't dirty the (static) original in the test
        MethodDescriptor<RpcCheckFormatRequest, RpcCheckFormatResponse> descriptor = methodDescriptor(
                RefServiceGrpc.getCheckFormatMethod(),
                RpcCheckFormatRequest.getDefaultInstance(),
                RpcCheckFormatResponse.getDefaultInstance());
        ServerCall<RpcCheckFormatRequest, RpcCheckFormatResponse> call = mock(ServerCall.class);
        lenient().when(call.getMethodDescriptor()).thenReturn(descriptor);

        ZeroCopyServerInterceptor interceptor = new ZeroCopyServerInterceptor(config);
        ServerCall.Listener<RpcCheckFormatRequest> listener =
                interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));

        assertNotNull(listener);
        assertThat(listener, not(instanceOf(ZeroCopyServerCallListener.class)));
        assertThat(descriptor.getRequestMarshaller(), not(instanceOf(ZeroCopyMarshaller.class)));
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testOverridesRequestMarshallerForMethodOptionWhenEnabled() {
        when(config.isZeroCopy()).thenReturn(true);

        //Create a copy of the original method descriptor so we don't dirty the (static) original in the test
        MethodDescriptor<RpcLfsReadObjectRequest, RpcStreamChunk> descriptor = methodDescriptor(
                LfsServiceGrpc.getReadObjectMethod(),
                RpcLfsReadObjectRequest.getDefaultInstance(),
                RpcStreamChunk.getDefaultInstance());
        ServerCall<RpcLfsReadObjectRequest, RpcStreamChunk> call = mock(ServerCall.class);
        when(call.getMethodDescriptor()).thenReturn(descriptor);

        ZeroCopyServerInterceptor interceptor = new ZeroCopyServerInterceptor(config);
        ServerCall.Listener<RpcLfsReadObjectRequest> listener =
                interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));

        assertNotNull(listener);
        assertThat(listener, instanceOf(ZeroCopyServerCallListener.class));
        assertThat(descriptor.getRequestMarshaller(), instanceOf(ZeroCopyMarshaller.class));
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testOverridesRequestMarshallerForServiceOptionWhenEnabled() {
        when(config.isZeroCopy()).thenReturn(true);

        //Create a copy of the original method descriptor so we don't dirty the (static) original in the test
        MethodDescriptor<RpcHttpBackendFragment, RpcHostingWriteStreamFragment> descriptor = methodDescriptor(
                HostingServiceGrpc.getHttpBackendMethod(),
                RpcHttpBackendFragment.getDefaultInstance(),
                RpcHostingWriteStreamFragment.getDefaultInstance());
        ServerCall<RpcHttpBackendFragment, RpcHostingWriteStreamFragment> call = mock(ServerCall.class);
        when(call.getMethodDescriptor()).thenReturn(descriptor);

        ZeroCopyServerInterceptor interceptor = new ZeroCopyServerInterceptor(config);
        ServerCall.Listener<RpcHttpBackendFragment> listener =
                interceptor.interceptCall(call, new Metadata(), mock(ServerCallHandler.class, RETURNS_DEEP_STUBS));

        assertNotNull(listener);
        assertThat(listener, instanceOf(ZeroCopyServerCallListener.class));
        assertThat(descriptor.getRequestMarshaller(), instanceOf(ZeroCopyMarshaller.class));
    }

    /**
     * Creates a {@link MethodDescriptor} as it generated by gRPC code generation. This method creates a copy of the
     * generated method descriptor and replaces the marshaller with the standard protobuf marshallers in to ensure
     * the method descriptor starts in the expected state (if a wired test runs before this unit test,
     * ZeroCopyServerInterceptor may already have replaced the marshaller).
     */
    private <ReqT extends MessageLite, RespT extends MessageLite> MethodDescriptor<ReqT, RespT> methodDescriptor(
            MethodDescriptor<ReqT, RespT> descriptor,
            ReqT defaultRequest, RespT defaultResponse) {
        MethodDescriptor.Marshaller<ReqT> requestMarshaller = ProtoLiteUtils.marshaller(defaultRequest);
        MethodDescriptor.Marshaller<RespT> responseMarshaller = ProtoLiteUtils.marshaller(defaultResponse);
        return descriptor.toBuilder(requestMarshaller, responseMarshaller).build();
    }
}