package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.grpc.FragmentResponseObserver;
import com.atlassian.bitbucket.mesh.process.OsProcess;
import com.atlassian.bitbucket.mesh.process.ProcessCompletionHandler;
import com.atlassian.bitbucket.mesh.process.ProcessStdioHandler;
import com.atlassian.bitbucket.mesh.process.ProcessSummary;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamFragment;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamResult;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamStart;
import com.google.protobuf.ByteString;
import com.zaxxer.nuprocess.NuProcess;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.ClientCallStreamObserver;
import org.apache.commons.lang3.StringUtils;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.junit.MockitoJUnitRunner;
import org.mockito.stubbing.Answer;

import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

import static io.grpc.Status.DEADLINE_EXCEEDED;
import static io.grpc.Status.INTERNAL;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.arrayContaining;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class FragmentResponseObserverTest {

    @Mock
    private ProcessCompletionHandler completionHandler;
    @Mock
    private ClientCallStreamObserver<Object> requestStream;
    @InjectMocks
    private FragmentResponseObserver<Void> responseObserver;
    @Mock
    private ProcessStdioHandler<Void> stdioHandler;

    @Test
    public void testBeforeStart() {
        responseObserver.beforeStart(requestStream);

        verify(stdioHandler).onPreStart(isNotNull());
    }

    @Test
    public void testOnCompleted() {
        responseObserver.onCompleted();

        verifyNoInteractions(completionHandler, requestStream, stdioHandler);
    }

    @Test
    public void testOnErrorForCancel() {
        doAnswer(invocation -> {
            OsProcess process = invocation.getArgument(0);
            process.cancel();

            return null;
        }).when(stdioHandler).onStart(any());

        responseObserver.beforeStart(requestStream);
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStart(RpcStreamStart.newBuilder()
                        .setCommandLine(ByteString.copyFromUtf8("git rev-list HEAD")))
                .build());
        responseObserver.onError(Status.CANCELLED
                .withDescription("The request was canceled by a handler")
                .asRuntimeException());
        responseObserver.asResult();

        verify(completionHandler).onComplete(argThat(ProcessSummary::isCanceled));
        verify(requestStream).cancel(eq("The request was canceled by a handler"), isNull());
    }

    @Test
    public void testOnErrorForTimeout() {
        responseObserver.beforeStart(requestStream);
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStart(RpcStreamStart.newBuilder()
                        .setCommandLine(ByteString.copyFromUtf8("git rev-list HEAD")))
                .build());
        responseObserver.onError(DEADLINE_EXCEEDED
                .withDescription("Simulated timeout")
                .asRuntimeException());
        responseObserver.asResult();

        verify(completionHandler).onComplete(argThat(ProcessSummary::isTimedOut));
    }

    @Test
    public void testOnErrorWithoutOnStart() {
        StatusRuntimeException expected = INTERNAL
                .withDescription("Some unexpected failure")
                .asRuntimeException();

        responseObserver.beforeStart(requestStream);
        responseObserver.onError(expected);
        responseObserver.asResult();

        ArgumentCaptor<ProcessSummary> summaryCaptor = ArgumentCaptor.forClass(ProcessSummary.class);
        verify(completionHandler).onComplete(summaryCaptor.capture());

        ProcessSummary summary = summaryCaptor.getValue();
        Throwable thrown = summary.getThrown();
        assertNotNull("An exception should have been passed to the exit handler", thrown);
        assertSame(expected, thrown);
        assertThat(thrown.getSuppressed(), arrayContaining(instanceOf(IllegalStateException.class)));
    }

    @Test
    public void testOnNextWithPartiallyConsumedStdout() {
        responseObserver.beforeStart(requestStream);

        ByteArrayOutputStream consumed = new ByteArrayOutputStream(1000);
        doAnswer(new Answer<Void>() {

            private final byte[] temp = new byte[100];

            @Override
            public Void answer(InvocationOnMock invocation) throws Exception {
                //Simulate draining the buffer in 100 byte chunks. Each callback, we read as many 100 byte
                //chunks as we can, and then potentially leave some number of bytes in the buffer
                ByteBuffer buffer = invocation.getArgument(0);
                while (buffer.remaining() >= 100) {
                    buffer.get(temp);
                    consumed.write(temp);
                }

                //On the final invocation, drain the whole buffer
                boolean closed = invocation.getArgument(1);
                if (closed && buffer.hasRemaining()) {
                    int remaining = buffer.remaining();
                    buffer.get(temp, 0, remaining);
                    consumed.write(temp, 0, remaining);
                }

                return null;
            }
        }).when(stdioHandler).onStdout(any(), anyBoolean());

        //Simulate writing the output to the callback in 5 275-byte chunks, which the handler will
        //consume in increments divisible by 100. That means:
        //1. Drain 200 bytes, 75 left
        //2. Drain 300 bytes, 50 left
        //3. Drain 300 bytes, 25 left
        //4. Drain 300 bytes, 0 left (This will let the proxy be nulled)
        //5. Drain 275 bytes (Final callback always drains the whole buffer)
        String data = StringUtils.repeat('a', 1375);
        byte[] output = data.getBytes(StandardCharsets.UTF_8);
        for (int i = 0; i < 5; i++) {
            ByteBuffer slice = ByteBuffer.wrap(output, i * 275, 275);

            responseObserver.onNext(RpcStreamFragment.newBuilder()
                    .setStdout(RpcStreamChunk.newBuilder()
                            .setClosed(i == 4) //Closed on the final iteration
                            .setData(ByteString.copyFrom(slice)))
                    .build());
        }

        assertEquals(data, new String(consumed.toByteArray(), StandardCharsets.UTF_8));
    }

    /**
     * Tests the scenario where the handler doesn't consume anything from the buffer until it gets closed == true.
     */
    @Test
    public void testOnNextWithProxiedBufferFollowedByClosed() {
        String testOutput = "Wed 22 Jan 2020 16:15:18 GMT\n";
        byte[] testOutputBytes = testOutput.getBytes(StandardCharsets.UTF_8);
        ByteBuffer expectedBuffer = ByteBuffer.wrap(testOutputBytes);
        expectedBuffer.flip();
        ByteBuffer expectedProxyBuffer = ByteBuffer.allocate(NuProcess.BUFFER_CAPACITY)
                .put(testOutputBytes);

        responseObserver.beforeStart(requestStream);
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStdout(RpcStreamChunk.newBuilder()
                        .setData(ByteString.copyFromUtf8(testOutput)))
                .build());
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStdout(RpcStreamChunk.newBuilder().setClosed(true))
                .build());
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setResult(RpcStreamResult.newBuilder())
                .build());
        responseObserver.asResult();

        verify(stdioHandler).onPreStart(any());
        verify(stdioHandler).onStdout(eq(expectedBuffer), eq(false));
        verify(stdioHandler).onStdout(eq(expectedProxyBuffer), eq(true));
    }

    @Test
    public void testOnNextWithResultAfterOutput() {
        String stderr = "fatal: not a git repository\n";

        responseObserver.beforeStart(requestStream);
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStart(RpcStreamStart.newBuilder()
                        .setCommandLine(ByteString.copyFromUtf8("git version")))
                .build());
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStderr(RpcStreamChunk.newBuilder()
                        .setClosed(true)
                        .setData(ByteString.copyFrom(stderr, StandardCharsets.UTF_8)))
                .build());
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setStdout(RpcStreamChunk.newBuilder()
                        .setClosed(true)
                        .setData(ByteString.EMPTY))
                .build());
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setResult(RpcStreamResult.newBuilder()
                        .setExitCode(1))
                .build());
        responseObserver.asResult();

        verify(stdioHandler).onExit(eq(1));

        ArgumentCaptor<ProcessSummary> summaryCaptor = ArgumentCaptor.forClass(ProcessSummary.class);
        verify(completionHandler).onComplete(summaryCaptor.capture());
    }

    @Test
    public void testOnNextWithResultButNoStart() {
        responseObserver.beforeStart(requestStream);
        responseObserver.onNext(RpcStreamFragment.newBuilder()
                .setResult(RpcStreamResult.newBuilder()
                        .setExitCode(1))
                .build());
        responseObserver.asResult();

        verify(stdioHandler).onExit(eq(1));

        ArgumentCaptor<ProcessSummary> summaryCaptor = ArgumentCaptor.forClass(ProcessSummary.class);
        verify(completionHandler).onComplete(summaryCaptor.capture());

        ProcessSummary summary = summaryCaptor.getValue();
        assertEquals(1, summary.getExitCode());
        assertThat(summary.getThrown(), instanceOf(IllegalStateException.class));
    }
}