package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.AbstractGrpcServiceTest;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.test.annotation.DirtiesRepo;
import com.atlassian.bitbucket.mesh.test.annotation.Repo;
import com.google.protobuf.ByteString;
import com.google.protobuf.Timestamp;
import io.grpc.ManagedChannel;
import io.grpc.stub.StreamObserver;
import org.apache.commons.lang3.StringUtils;
import org.junit.Test;
import org.springframework.beans.factory.annotation.Autowired;

import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.Random;
import java.util.concurrent.TimeUnit;

/**
 * Class that can be used to benchmark the performance of GitDataService, which is useful when experimenting with
 * different buffer sizes, flow control settings and marshallers. The class name intentionally does not end with Test
 * to ensure the benchmark is not run as part of the normal unit tests.
 */
public class GrpcGitDataServiceBenchmark extends AbstractGrpcServiceTest {

    @Autowired
    private ManagedChannel managedChannel;
    @Repo
    private Repository repository;

    @DirtiesRepo
    @Test
    public void testUploadLargeFile() throws Exception {
        final int loops = 50;
        final int warmUp = 10;
        long maxDuration = 0L;
        long minDuration = Long.MAX_VALUE;
        long totalDuration = 0L;
        int messageSize = 268435456;

        //I've experimented with different 'chunk' sizes and 32kb messages seems to be the sweet spot in terms of
        //throughput. Interestingly, if the chunk size is just a little too small or too high, throughput drops down
        //significantly. Some numbers to back it up:
        // - 64kb:      {min = 3351280260ns, max = 3745184828ns, avg = 3473720503ns} (yes 10x as high as 32kb)
        // - 32kb + 64: {min = 291052404ns, max = 698749069ns, avg = 371808950ns}
        // - 32kb:      {min = 278052693ns, max = 660331051ns, avg = 355985144ns}
        // - 32kb - 4:  {min = 288330004ns, max = 627816192ns, avg = 358949170ns}
        // - 32kb - 64: {min = 540334488ns, max = 3267285811ns, avg = 1883779360ns}
        // - 16kb:      {min = 295587829ns, max = 912885831ns, avg = 517399602ns}
        // -  8kb:      {min = 736087294ns, max = 2364473914ns, avg = 1318965307ns}
        byte[] bytes = new byte[32 * 1024];
        int chunks = (messageSize - 1) / bytes.length + 1;
        Random random = new Random(messageSize);
        random.nextBytes(bytes);

        for (int i = 0; i < loops + warmUp; ++i) {
            long start = System.nanoTime();
            UploadResponseObserver<RpcUploadFileFragment, RpcUploadFileResponse> responseObserver =
                    new UploadResponseObserver<>();
            StreamObserver<RpcUploadFileFragment> requestObserver = stub().uploadFile(responseObserver);

            //Send the upload file request
            RpcUploadFileResponse response = responseObserver.sendRequest(RpcUploadFileFragment.newBuilder()
                    .setRequest(RpcUploadFileRequest.newBuilder()
                            .setRepository(repository.getId())
                            .setFile(RpcOsFile.newBuilder()
                                    .setLastModified(Timestamp.newBuilder().setSeconds(1623659279L))
                                    .setPath(ByteString.copyFromUtf8("pack/pack-23cd17a83701deb0c649a2e1ce6e68f66a7cd172.pack"))
                                    .setSize(268435456L)) // 256M
                            .setType(RpcUploadType.UPLOAD_TYPE_GIT_OBJECTS))
                    .build())
                    .get(5, TimeUnit.SECONDS);

            assertTrue(response.getWant());
            //Stream the file to the server in chunks of bytes.length each
            for (int j = 0; j < chunks; ++j) {
                boolean closed = j == chunks - 1;
                responseObserver.sendChunk(RpcUploadFileFragment.newBuilder()
                        .setChunk(RpcStreamChunk.newBuilder()
                                .setData(closed ?
                                        ByteStringUtils.unsafeWrap(bytes, 0, (messageSize - (j * bytes.length))) :
                                        ByteStringUtils.unsafeWrap(bytes))
                                .setClosed(closed))
                        .build());
            }
            requestObserver.onCompleted();

            responseObserver.awaitCompleted(5, TimeUnit.SECONDS);
            long duration = System.nanoTime() - start;
            System.out.println("iteration [" + StringUtils.leftPad(Integer.toString(i - warmUp + 1), 3, '0') + "/" +
                    loops + "]:\t" + (duration / 1_000_000.0) + "ms");
            if (i >= warmUp) {
                totalDuration += duration;
                if (duration > maxDuration) {
                    maxDuration = duration;
                }
                if (duration < minDuration) {
                    minDuration = duration;
                }
            }

            Path expectedPath = repository.getPath().resolve("objects/pack/pack-23cd17a83701deb0c649a2e1ce6e68f66a7cd172.pack");
            assertTrue(expectedPath + " is not a regular file", Files.isRegularFile(expectedPath));

            //Verify the attributes
            BasicFileAttributes attributes = Files.readAttributes(expectedPath, BasicFileAttributes.class);
            assertEquals(messageSize, attributes.size());
            assertEquals(1623659279000L, attributes.lastModifiedTime().toMillis());
            Files.delete(expectedPath);

            //Sleep for a bit to let the system CPU throttle down again
            Thread.sleep(250L);
        }

        System.out.println("Time taken: {min = " + minDuration + "ns, max = " + maxDuration + "ns, avg = " +
                (totalDuration / loops) + "ns}");
    }

    private GitDataServiceGrpc.GitDataServiceStub stub() {
        return GitDataServiceGrpc.newStub(managedChannel);
    }
}
