package org.infinispan.protostream;

import org.infinispan.protostream.config.Configuration;
import org.infinispan.protostream.impl.BaseMarshallerDelegate;
import org.infinispan.protostream.impl.RawProtoStreamReaderImpl;
import org.infinispan.protostream.impl.RawProtoStreamWriterImpl;
import org.infinispan.protostream.impl.SerializationContextImpl;
import org.infinispan.protostream.impl.WrappedMessageMarshaller;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

/**
 * @author anistor@redhat.com
 * @since 1.0
 */
public final class ProtobufUtil {

   private static final String WRAPPING_DEFINITIONS_RES = "/org/infinispan/protostream/message-wrapping.proto";

   public static SerializationContext newSerializationContext(Configuration configuration) {
      SerializationContextImpl serializationContext = new SerializationContextImpl(configuration);

      try {
         serializationContext.registerProtoFiles(FileDescriptorSource.fromResources(WRAPPING_DEFINITIONS_RES));
      } catch (IOException e) {
         throw new RuntimeException("Failed to initialize serialization context", e);
      } catch (DescriptorParserException e) {
         throw new RuntimeException("Failed to initialize serialization context", e);
      }

      serializationContext.registerMarshaller(new WrappedMessageMarshaller());

      return serializationContext;
   }

   private static <A> void writeTo(SerializationContext ctx, RawProtoStreamWriter out, A t) throws IOException {
      if (t == null) {
         throw new IllegalArgumentException("Object to marshall cannot be null");
      }
      BaseMarshallerDelegate marshallerDelegate = ((SerializationContextImpl) ctx).getMarshallerDelegate(t.getClass());
      marshallerDelegate.marshall(null, t, null, out);
      out.flush();
   }

   public static void writeTo(SerializationContext ctx, OutputStream out, Object t) throws IOException {
      writeTo(ctx, RawProtoStreamWriterImpl.newInstance(out), t);
   }

   public static byte[] toByteArray(SerializationContext ctx, Object t) throws IOException {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();
      writeTo(ctx, baos, t);
      return baos.toByteArray();
   }

   private static <A> A readFrom(SerializationContext ctx, RawProtoStreamReader in, Class<A> clazz) throws IOException {
      BaseMarshallerDelegate<A> marshallerDelegate = ((SerializationContextImpl) ctx).getMarshallerDelegate(clazz);
      return marshallerDelegate.unmarshall(null, null, in);
   }

   public static <A> A readFrom(SerializationContext ctx, InputStream in, Class<A> clazz) throws IOException {
      return readFrom(ctx, RawProtoStreamReaderImpl.newInstance(in), clazz);
   }

   public static <A> A fromByteArray(SerializationContext ctx, byte[] bytes, Class<A> clazz) throws IOException {
      return readFrom(ctx, RawProtoStreamReaderImpl.newInstance(bytes), clazz);
   }

   public static <A> A fromByteArray(SerializationContext ctx, byte[] bytes, int offset, int length, Class<A> clazz) throws IOException {
      return readFrom(ctx, RawProtoStreamReaderImpl.newInstance(bytes, offset, length), clazz);
   }

   /**
    * Parses a top-level message that was wrapped according to the org.infinispan.protostream.WrappedMessage proto
    * definition.
    *
    * @param ctx
    * @param bytes
    * @return
    * @throws IOException
    */
   public static Object fromWrappedByteArray(SerializationContext ctx, byte[] bytes) throws IOException {
      return fromWrappedByteArray(ctx, bytes, 0, bytes.length);
   }

   public static Object fromWrappedByteArray(SerializationContext ctx, byte[] bytes, int offset, int length) throws IOException {
      ByteArrayInputStream bais = new ByteArrayInputStream(bytes, offset, length);
      return WrappedMessageMarshaller.readWrappedMessage(ctx, RawProtoStreamReaderImpl.newInstance(bais));
   }

   public static byte[] toWrappedByteArray(SerializationContext ctx, Object t) throws IOException {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();
      WrappedMessageMarshaller.writeWrappedMessage(ctx, RawProtoStreamWriterImpl.newInstance(baos), t);
      return baos.toByteArray();
   }
}

@@ -149,16 +149,80 @@
 legate;%0A
+import org.infinispan.protostream.impl.ByteArrayOutputStreamEx;%0A
 import o
@@ -270,32 +270,32 @@
 reamReaderImpl;%0A
-
 import org.infin
@@ -633,16 +633,44 @@
 tStream;
+%0Aimport java.nio.ByteBuffer;
 %0A%0A/**%0A *
@@ -2380,32 +2380,268 @@
 eArray();%0A   %7D%0A%0A
+   public static ByteBuffer toByteBuffer(SerializationContext ctx, Object t) throws IOException %7B%0A      ByteArrayOutputStreamEx baos = new ByteArrayOutputStreamEx();%0A      writeTo(ctx, baos, t);%0A      return baos.getByteBuffer();%0A   %7D%0A%0A
    private stati
@@ -3313,32 +3313,110 @@
 , clazz);%0A   %7D%0A%0A
+   //todo %5Banistor%5D what happens with remaining trailing bytes? signal error?%0A
    public static
@@ -3635,24 +3635,239 @@
 azz);%0A   %7D%0A%0A
+   public static %3CA%3E A fromByteBuffer(SerializationContext ctx, ByteBuffer byteBuffer, Class%3CA%3E clazz) throws IOException %7B%0A      return readFrom(ctx, RawProtoStreamReaderImpl.newInstance(byteBuffer), clazz);%0A   %7D%0A%0A
    /**%0A    *
@@ -4577,32 +4577,267 @@
 e(bais));%0A   %7D%0A%0A
+   public static Object fromWrappedByteBuffer(SerializationContext ctx, ByteBuffer byteBuffer) throws IOException %7B%0A      return WrappedMessageMarshaller.readWrappedMessage(ctx, RawProtoStreamReaderImpl.newInstance(byteBuffer));%0A   %7D%0A%0A
    public static
@@ -5084,24 +5084,24 @@
 (baos), t);%0A
-
       return
@@ -5125,11 +5125,329 @@
 ();%0A   %7D
+%0A%0A   public static ByteBuffer toWrappedByteBuffer(SerializationContext ctx, Object t) throws IOException %7B%0A      ByteArrayOutputStreamEx baos = new ByteArrayOutputStreamEx();%0A      WrappedMessageMarshaller.writeWrappedMessage(ctx, RawProtoStreamWriterImpl.newInstance(baos), t);%0A      return baos.getByteBuffer();%0A   %7D
 %0A%7D%0A

Add preliminary ByteBuffer support