package com.atlassian.bitbucket.mesh.rpc;

import com.atlassian.bitbucket.mesh.rpc.v1.RpcDiffRequest;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcGitTimeouts;
import com.google.protobuf.ByteString;
import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.Descriptors;
import io.grpc.MethodDescriptor.PrototypeMarshaller;
import io.grpc.internal.CompositeReadableBuffer;
import io.grpc.internal.ReadableBuffers;
import io.grpc.netty.GrpcNettyBackdoor;
import io.grpc.protobuf.lite.ProtoLiteUtils;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BooleanSupplier;

import static org.junit.Assert.*;

public class ZeroCopyMarshallerTest {

    @Before
    public void setUp() {
        System.setProperty("io.netty.maxDirectMemory", Long.toString(1024 * 1024L));
    }

    @After
    public void tearDown() {
        System.clearProperty("io.netty.maxDirectMemory");
    }

    @Test
    public void testFallbackForCompositeDataFrame() throws IOException {
        doTestZeroCopyMarshalling(() -> false, 0, 3);
    }

    @Test
    public void testFallbackForUnsupportedStreamType() throws IOException {
        RpcDiffRequest request = createDiffRequest();
        byte[] encoded = encode(request);

        RpcDiffRequest parsed = createMarshaller(() -> true).parse(new ByteArrayInputStream(encoded));
        assertEquals(request, parsed);
    }

    @Test
    public void testZeroCopyDisabled() throws IOException {
        doTestZeroCopyMarshalling(() -> false, 0, 1);
    }

    @Test
    public void testZeroCopyForCompositeDataFrame() throws IOException {
        doTestZeroCopyMarshalling(() -> true, 0, 3);
    }

    @Test
    public void testZeroCopyNonZeroOffset() throws IOException {
        doTestZeroCopyMarshalling(() -> true, 20, 1);
    }

    @Test
    public void testZeroCopyZeroOffset() throws IOException {
        doTestZeroCopyMarshalling(() -> true, 0, 1);
    }

    private static RpcDiffRequest createDiffRequest() {
        RpcDiffRequest.Builder builder = RpcDiffRequest.newBuilder()
                .addPaths(ByteString.copyFromUtf8("src/main/test"))
                .setAncestor(ByteString.copyFromUtf8("ancestor"))
                .setRev(ByteString.copyFromUtf8("revision"));
        //Construct the message just like we'd construct it in the application code (setting timeouts and repository
        //using setField in GitRequestHelper), in case it changes the field order in the message
        Descriptors.FieldDescriptor timeoutsField = builder.getDescriptorForType().findFieldByName("timeouts");
        builder.setField(timeoutsField, RpcGitTimeouts.newBuilder()
                .setExecution(120)
                .setIdle(10)
                .build());
        Descriptors.FieldDescriptor repoField = builder.getDescriptorForType().findFieldByName("repository");
        builder.setField(repoField, "p/0001/hierarchy/r/repo-id");

        return builder.build();
    }

    private static SimulatedFrame simulateDeframe(byte[] payload, int garbagePrefixLength, int fragments) {
        //Simulates one or more Netty ByteBuf's being processed to construct a frame of gRPC data. To get as close
        //as possible to the 'real' thing, we simulate gRPCs internal flow:
        //
        //- Netty's ByteBuf is wrapped in a NettyReadableBuffer, which increases the ref count
        //- MessageDeframer.deframe adds the incoming ByteBuf to a CompositeReadableBuffer for delivery on
        //  another thread.
        SimulatedFrame frame = new SimulatedFrame();
        int index = 0;
        int bytesPerFragment = payload.length / fragments;

        //Special handling for the first buffer to simulate starting the read from a non-zero index.
        ByteBuf buf = Unpooled.directBuffer(bytesPerFragment + garbagePrefixLength);
        for (int i = 0; i < garbagePrefixLength; ++i) {
            buf.writeByte(i);
        }
        buf.writeBytes(payload, 0, bytesPerFragment);
        index += bytesPerFragment;
        buf.skipBytes(garbagePrefixLength);

        //Simulate MessageDeframer.deframe adding the incoming ByteBuf for future delivery
        frame.add(buf);

        //Repeat for the subsequent fragments
        for (int i = 1; i < fragments; ++i) {
            int length = i == fragments - 1 ? payload.length - index : bytesPerFragment;
            buf = Unpooled.directBuffer(length);
            buf.writeBytes(payload, index, length);
            index += length;

            frame.add(buf);
        }

        return frame;
    }

    private static ZeroCopyMarshaller<RpcDiffRequest> createMarshaller(BooleanSupplier enabledSupplier) {
        return new ZeroCopyMarshaller<>(enabledSupplier,
                (PrototypeMarshaller<RpcDiffRequest>) ProtoLiteUtils.marshaller(RpcDiffRequest.getDefaultInstance()));
    }

    private static void doTestZeroCopyMarshalling(BooleanSupplier enabledSupplier,
                                                  int garbagePrefixLength,
                                                  int frameFragments) throws IOException {
        RpcDiffRequest request = createDiffRequest();
        byte[] encoded = encode(request);

        //Here we simulate grpc-netty's data parsing flow; MessageDeframer.deframe adds incoming Netty buffers to a
        //CompositeReadableBuffer for delivery on another thread (when enough data is available).
        //frame contains both the constructed CompositeReadableBuffer and the backing ByteBuf instance so we can
        //verify when they're released.
        SimulatedFrame simulatedFrame = simulateDeframe(encoded, garbagePrefixLength, frameFragments);

        //MessageDeframer.deliver() is called from a grpc-nio-worker thread, which converts the frame into an
        //ReadableBuffers.BufferInputStream. This transfers the 'ownership' of the frame; the frame is closed when
        //the InputStream is closed.
        try (InputStream inputStream = ReadableBuffers.openStream(simulatedFrame.frame, true)) {
            //- The resulting stream is wrapped in a SingleMessageProducer which is handed to the ServerStreamListener
            //  for processing.
            //- On yet another thread (grpc-server-*), the ServerStreamListener parses the message, delivers the parsed
            //  message to the gRPC service implementation and closes the stream, which closes the underlying
            //  CompositeReadableBuffer, which closes the encapsulated NettyReadableBuffer, which releases the
            //  underlying Netty ByteBuf(s).
            RpcDiffRequest parsed = createMarshaller(enabledSupplier).parse(inputStream);
            assertEquals(request, parsed);

            //Without zero-copy, CompositeReadableBuffer discards and closes its nested ReadableBuffer once it has
            //fully read it. With zero-copy, we need to retain these buffers
            if (enabledSupplier.getAsBoolean()) {
                //Verify that the buffer has not yet been released
                assertFalse("Buffer was released, which could cause data corruption", simulatedFrame.isReleased());
            } else {
                assertTrue("Buffer was expected to be released when zero-copy is disabled", simulatedFrame.isReleased());
            }
        }

        if (enabledSupplier.getAsBoolean()) {
            //Closing the input stream should NOT release the underlying buffers iff zero-copy is enabled
            assertFalse(simulatedFrame.isReleased());

            //Closing the detached stream should release the underlying buffers
            ZeroCopyStreamCloser.close();
        }
        assertTrue(simulatedFrame.isReleased());
    }

    private static byte[] encode(RpcDiffRequest request) throws IOException {
        //Construct a message and serialize it to bytes
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        CodedOutputStream cos = CodedOutputStream.newInstance(bos);

        request.writeTo(cos);
        cos.flush();

        return bos.toByteArray();
    }

    private static class SimulatedFrame {

        private final List<ByteBuf> buffers;
        private final CompositeReadableBuffer frame;

        SimulatedFrame() {
            buffers = new ArrayList<>();
            frame = new CompositeReadableBuffer();
        }

        void add(ByteBuf buf) {
            buffers.add(buf);
            frame.addBuffer(GrpcNettyBackdoor.newReadableBuffer(buf.retain()));
            //gRPC delivers the frame on a different thread, so release the underlying Netty buffer to
            //match the application behavior.
            buf.release();
        }

        boolean isReleased() {
            return buffers.stream().allMatch(buf -> buf.refCnt() == 0);
        }
    }
}