/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.arrow.vector.ipc;

import static java.util.Arrays.asList;
import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.Channels;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.ipc.message.ArrowMessage;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.ipc.message.IpcOption;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.MetadataVersion;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Test;

public class MessageSerializerTest {

  public static ArrowBuf buf(BufferAllocator alloc, byte[] bytes) {
    ArrowBuf buffer = alloc.buffer(bytes.length);
    buffer.writeBytes(bytes);
    return buffer;
  }

  public static byte[] array(ArrowBuf buf) {
    byte[] bytes = new byte[checkedCastToInt(buf.readableBytes())];
    buf.readBytes(bytes);
    return bytes;
  }

  private int intToByteRoundtrip(int v, byte[] bytes) {
    MessageSerializer.intToBytes(v, bytes);
    return MessageSerializer.bytesToInt(bytes);
  }

  @Test
  public void testIntToBytes() {
    byte[] bytes = new byte[4];
    int[] values = new int[] {1, 15, 1 << 8, 1 << 16, Integer.MAX_VALUE};
    for (int v : values) {
      assertEquals(intToByteRoundtrip(v, bytes), v);
    }
  }

  @Test
  public void testWriteMessageBufferAligned() throws IOException {
    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    WriteChannel out = new WriteChannel(Channels.newChannel(outputStream));

    // This is not a valid Arrow Message, only to test writing and alignment
    ByteBuffer buffer = ByteBuffer.allocate(8).order(ByteOrder.nativeOrder());
    buffer.putInt(1);
    buffer.putInt(2);
    buffer.flip();

    int bytesWritten = MessageSerializer.writeMessageBuffer(out, 8, buffer);
    assertEquals(16, bytesWritten);

    buffer.rewind();
    buffer.putInt(3);
    buffer.flip();
    bytesWritten = MessageSerializer.writeMessageBuffer(out, 4, buffer);
    assertEquals(16, bytesWritten);

    ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray());
    ReadChannel in = new ReadChannel(Channels.newChannel(inputStream));
    ByteBuffer result = ByteBuffer.allocate(32).order(ByteOrder.nativeOrder());
    in.readFully(result);
    result.rewind();

    // First message continuation, size, and 2 int values
    assertEquals(MessageSerializer.IPC_CONTINUATION_TOKEN, result.getInt());
    // message length is represented in little endian
    result.order(ByteOrder.LITTLE_ENDIAN);
    assertEquals(8, result.getInt());
    result.order(ByteOrder.nativeOrder());
    assertEquals(1, result.getInt());
    assertEquals(2, result.getInt());

    // Second message continuation, size, 1 int value and 4 bytes padding
    assertEquals(MessageSerializer.IPC_CONTINUATION_TOKEN, result.getInt());
    // message length is represented in little endian
    result.order(ByteOrder.LITTLE_ENDIAN);
    assertEquals(8, result.getInt());
    result.order(ByteOrder.nativeOrder());
    assertEquals(3, result.getInt());
    assertEquals(0, result.getInt());
  }

  @Test
  public void testSchemaMessageSerialization() throws IOException {
    Schema schema = testSchema();
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    long size = MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema);
    assertEquals(size, out.toByteArray().length);

    ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
    Schema deserialized =
        MessageSerializer.deserializeSchema(new ReadChannel(Channels.newChannel(in)));
    assertEquals(schema, deserialized);
    assertEquals(1, deserialized.getFields().size());
  }

  @Test
  public void testSchemaDictionaryMessageSerialization() throws IOException {
    DictionaryEncoding dictionary = new DictionaryEncoding(9L, false, new ArrowType.Int(8, true));
    Field field =
        new Field("test", new FieldType(true, ArrowType.Utf8.INSTANCE, dictionary, null), null);
    Schema schema = new Schema(Collections.singletonList(field));
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    long size = MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema);
    assertEquals(size, out.toByteArray().length);

    ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
    Schema deserialized =
        MessageSerializer.deserializeSchema(new ReadChannel(Channels.newChannel(in)));
    assertEquals(schema, deserialized);
  }

  @Test
  public void testSerializeRecordBatchV4() throws IOException {
    byte[] validity = new byte[] {(byte) 255, 0};
    // second half is "undefined"
    byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};

    BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE);
    ArrowBuf validityb = buf(alloc, validity);
    ArrowBuf valuesb = buf(alloc, values);

    ArrowRecordBatch batch =
        new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb));

    // avoid writing legacy ipc format by default
    IpcOption option = new IpcOption(false, MetadataVersion.V4);
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch, option);

    ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
    ReadChannel channel = new ReadChannel(Channels.newChannel(in));
    ArrowMessage deserialized = MessageSerializer.deserializeMessageBatch(channel, alloc);
    assertEquals(ArrowRecordBatch.class, deserialized.getClass());
    verifyBatch((ArrowRecordBatch) deserialized, validity, values);
  }

  @Test
  public void testSerializeRecordBatchV5() throws Exception {
    byte[] validity = new byte[] {(byte) 255, 0};
    // second half is "undefined"
    byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};

    BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE);
    ArrowBuf validityb = buf(alloc, validity);
    ArrowBuf valuesb = buf(alloc, values);

    ArrowRecordBatch batch =
        new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb));

    // avoid writing legacy ipc format by default
    IpcOption option = new IpcOption(false, MetadataVersion.V5);
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch, option);
    validityb.close();
    valuesb.close();
    batch.close();

    {
      ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
      ReadChannel channel = new ReadChannel(Channels.newChannel(in));
      ArrowMessage deserialized = MessageSerializer.deserializeMessageBatch(channel, alloc);
      assertEquals(ArrowRecordBatch.class, deserialized.getClass());
      verifyBatch((ArrowRecordBatch) deserialized, validity, values);
      deserialized.close();
    }

    {
      byte[] validBytes = out.toByteArray();
      byte[] missingBytes = Arrays.copyOfRange(validBytes, /*from=*/ 0, validBytes.length - 1);

      ByteArrayInputStream in = new ByteArrayInputStream(missingBytes);
      ReadChannel channel = new ReadChannel(Channels.newChannel(in));

      assertThrows(
          IOException.class, () -> MessageSerializer.deserializeMessageBatch(channel, alloc));
    }

    alloc.close();
  }

  public static Schema testSchema() {
    return new Schema(
        asList(
            new Field(
                "testField",
                FieldType.nullable(new ArrowType.Int(8, true)),
                Collections.<Field>emptyList())));
  }

  // Verifies batch contents matching test schema.
  public static void verifyBatch(ArrowRecordBatch batch, byte[] validity, byte[] values) {
    assertTrue(batch != null);
    List<ArrowFieldNode> nodes = batch.getNodes();
    assertEquals(1, nodes.size());
    ArrowFieldNode node = nodes.get(0);
    assertEquals(16, node.getLength());
    assertEquals(8, node.getNullCount());
    List<ArrowBuf> buffers = batch.getBuffers();
    assertEquals(2, buffers.size());
    assertArrayEquals(validity, MessageSerializerTest.array(buffers.get(0)));
    assertArrayEquals(values, MessageSerializerTest.array(buffers.get(1)));
  }
}
