package rx.broadcast;

import com.google.protobuf.ByteString;
import com.google.protobuf.DescriptorProtos.DescriptorProto;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
import com.google.protobuf.DescriptorProtos.FieldOptions;
import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Descriptors.FileDescriptor;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Parser;
import org.jetbrains.annotations.NotNull;

import java.util.concurrent.atomic.AtomicInteger;

public class CausalOrderProtobufSerializer<T> implements Serializer<VectorTimestamped<T>> {
    private static final String MESSAGE_NAME = "VectorTimestamped";

    private static final String IDS_FIELD_NAME = "ids";

    private static final String TIMESTAMPS_FIELD_NAME = "timestamps";

    private static final String VALUE_FIELD_NAME = "value";

    private static final int IDS_FIELD_NUMBER = 1;

    private static final int TIMESTAMPS_FIELD_NUMBER = 2;

    private static final int VALUE_FIELD_NUMBER = 3;

    private final FieldDescriptor ids;

    private final FieldDescriptor timestamps;

    private final FieldDescriptor value;

    private final DynamicMessage.Builder messageBuilder;

    private final Parser<DynamicMessage> messageParser;

    private final Serializer<T> objectSerializer;

    public CausalOrderProtobufSerializer(@NotNull final Serializer<T> objectSerializer) {
        this.objectSerializer = objectSerializer;
        try {
            final FileDescriptorProto timestampedMessageFile = FileDescriptorProto.newBuilder()
                .addMessageType(
                    DescriptorProto.newBuilder()
                        .setName(MESSAGE_NAME)
                        .addField(
                            FieldDescriptorProto.newBuilder()
                                .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
                                .setName(IDS_FIELD_NAME)
                                .setNumber(IDS_FIELD_NUMBER)
                                .setType(FieldDescriptorProto.Type.TYPE_BYTES))
                        .addField(
                            FieldDescriptorProto.newBuilder()
                                .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
                                .setName(TIMESTAMPS_FIELD_NAME)
                                .setNumber(TIMESTAMPS_FIELD_NUMBER)
                                .setOptions(FieldOptions.newBuilder()
                                    .setPacked(true))
                                .setType(FieldDescriptorProto.Type.TYPE_UINT64))
                        .addField(
                            FieldDescriptorProto.newBuilder()
                                .setName(VALUE_FIELD_NAME)
                                .setNumber(VALUE_FIELD_NUMBER)
                                .setType(FieldDescriptorProto.Type.TYPE_BYTES)))
                .build();
            final Descriptor message = FileDescriptor
                .buildFrom(timestampedMessageFile, new FileDescriptor[0])
                .findMessageTypeByName(MESSAGE_NAME);
            this.ids = message.findFieldByName(IDS_FIELD_NAME);
            this.timestamps = message.findFieldByName(TIMESTAMPS_FIELD_NAME);
            this.value = message.findFieldByName(VALUE_FIELD_NAME);
            this.messageBuilder = DynamicMessage.newBuilder(message);
            this.messageParser = messageBuilder.buildPartial().getParserForType();
        } catch (final DescriptorValidationException e) {
            throw new RuntimeException(e);
        }
    }

    @NotNull
    @Override
    public final VectorTimestamped<T> decode(@NotNull final byte[] data) {
        try {
            final DynamicMessage message = messageParser.parseFrom(data);
            final byte[] bytes = ((ByteString) message.getField(this.value)).toByteArray();
            final T value = objectSerializer.decode(bytes);
            final int idsCount = message.getRepeatedFieldCount(this.ids);
            final int timestampsCount = message.getRepeatedFieldCount(this.timestamps);
            final Sender[] ids = new Sender[idsCount];
            final long[] timestamps = new long[timestampsCount];

            for (int i = 0; i < idsCount; i++) {
                ids[i] = new Sender(((ByteString) message.getRepeatedField(this.ids, i)).toByteArray());
            }
            for (int i = 0; i < timestampsCount; i++) {
                timestamps[i] = (long) message.getRepeatedField(this.timestamps, i);
            }

            return new VectorTimestamped<>(value, new VectorTimestamp(ids, timestamps));
        } catch (final InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
    }

    @NotNull
    @Override
    public final byte[] encode(@NotNull final VectorTimestamped<T> data) {
        final DynamicMessage.Builder builder = messageBuilder
            .setField(this.value, objectSerializer.encode(data.value));

        final AtomicInteger i = new AtomicInteger(0);
        data.timestamp.stream().forEachOrdered(entry -> {
            builder.setRepeatedField(this.ids, i.get(), entry.id);
            builder.setRepeatedField(this.timestamps, i.get(), entry.timestamp);
            i.incrementAndGet();
        });

        return builder.build().toByteArray();
    }
}

@@ -721,59 +721,8 @@
 l;%0A%0A
-import java.util.concurrent.atomic.AtomicInteger;%0A%0A
 publ
@@ -1232,16 +1232,65 @@
 R = 3;%0A%0A
+    private final Descriptor messageDescriptor;%0A%0A
     priv
@@ -1414,66 +1414,8 @@
 e;%0A%0A
-    private final DynamicMessage.Builder messageBuilder;%0A%0A
     
@@ -3296,32 +3296,78 @@
 (MESSAGE_NAME);%0A
+            this.messageDescriptor = message;%0A
             this
@@ -3372,32 +3372,42 @@
 is.ids = message
+Descriptor
 .findFieldByName
@@ -3453,32 +3453,42 @@
 stamps = message
+Descriptor
 .findFieldByName
@@ -3544,16 +3544,26 @@
  message
+Descriptor
 .findFie
@@ -3614,21 +3614,20 @@
 .message
-Build
+Pars
 er = Dyn
@@ -3656,18 +3656,27 @@
 (message
+Descriptor
 )
-;
 %0A       
@@ -3684,58 +3684,44 @@
     
-this.messageParser = messageBuilder.buildPartial()
+    .buildPartial()%0A                
 .get
@@ -5165,22 +5165,52 @@
 r = 
-m
+DynamicM
 essage
+.new
 Builder
+(messageDescriptor)
 %0A   
@@ -5283,62 +5283,8 @@
 );%0A%0A
-        final AtomicInteger i = new AtomicInteger(0);%0A
     
@@ -5349,35 +5349,35 @@
         builder.
-set
+add
 RepeatedField(th
@@ -5388,25 +5388,27 @@
 ds, 
-i.get(), entry.id
+entry.id.byteBuffer
 );%0A 
@@ -5426,19 +5426,19 @@
 builder.
-set
+add
 Repeated
@@ -5463,17 +5463,8 @@
 mps,
- i.get(),
  ent
@@ -5482,41 +5482,8 @@
 p);%0A
-            i.incrementAndGet();%0A
     

Fix CausalOrderProtobufSerializer repeated fields