package com.atlassian.bitbucket.mesh.rpc;

import com.google.protobuf.*;
import io.grpc.*;
import io.grpc.MethodDescriptor.PrototypeMarshaller;
import io.grpc.protobuf.lite.GrpcProtoLiteBackdoor;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BooleanSupplier;

import static java.util.Objects.requireNonNull;

public class ZeroCopyMarshaller<T extends MessageLite> implements MethodDescriptor.Marshaller<T> {

    private final BooleanSupplier enabledSupplier;
    private final PrototypeMarshaller<T> protobufMarshaller;
    private final Parser<T> protobufParser;

    public ZeroCopyMarshaller(BooleanSupplier enabledSupplier, PrototypeMarshaller<T> protobufMarshaller) {
        requireNonNull(enabledSupplier, "enabledSupplier");
        this.protobufMarshaller = requireNonNull(protobufMarshaller, "protobufMarshaller");

        T prototype = protobufMarshaller.getMessagePrototype();
        if (prototype == null) {
            //Some types such as Void can have a null message prototype.
            //For these, we don't perform zero-copy optimizations
            this.enabledSupplier = () -> false;
            protobufParser = null;
        } else {
            this.enabledSupplier = enabledSupplier;
            //noinspection unchecked
            protobufParser = (Parser<T>) prototype.getParserForType();
        }
    }

    @Nonnull
    @Override
    public T parse(@Nonnull InputStream stream) {
        try {
            if (enabledSupplier.getAsBoolean() && stream instanceof Detachable && stream instanceof KnownLength &&
                    stream instanceof HasByteBuffer && ((HasByteBuffer) stream).byteBufferSupported()) {
                int size = stream.available();
                //Detach the stream. The application is now responsible for closing the InputStream when the message
                //has been processed. Register the detached stream with ZeroCopyStreamCloser, which will close it
                //after onMessage completes.
                InputStream detached = ((Detachable) stream).detach();
                ZeroCopyStreamCloser.register(detached);

                //Mark the stream to prevent the underlying ByteBuf instances to be released after they've been read
                //(see CompositeReadableBuffer.advanceBuffer)
                detached.mark(size);
                List<ByteString> fragments = new ArrayList<>();
                while (detached.available() != 0) {
                    ByteBuffer buffer = ((HasByteBuffer) detached).getByteBuffer();
                    if (buffer == null) {
                        break;
                    }
                    fragments.add(UnsafeByteOperations.unsafeWrap(buffer));
                    //The buffer position and the _stream_ position are independent, so we have to skip
                    //bytes on the stream to move forward to the next underlying buffer. The previous
                    //buffer still has its own position, so the data is still available there
                    detached.skip(buffer.remaining());
                }

                //Create an immutable CodedInputStream with aliasing enabled to allow zero-copy
                //unmarshalling of ByteString fields
                CodedInputStream cis = ByteString.copyFrom(fragments).newCodedInput();
                cis.enableAliasing(true);
                cis.setSizeLimit(Integer.MAX_VALUE);
                return protobufParser.parseFrom(cis, GrpcProtoLiteBackdoor.getExtensionRegistry());
            }
        } catch (InvalidProtocolBufferException ipbe) {
            throw Status.INTERNAL.withCause(ipbe).asRuntimeException();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        //Cannot or should not apply zero-copy optimizations, fall back to the default parser
        return protobufMarshaller.parse(stream);
    }

    @Override
    public InputStream stream(T value) {
        return GrpcProtoLiteBackdoor.newProtoInputStream(value, protobufParser);
    }
}
