/*
 * 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 org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Integration test for reading/writing {@link org.apache.arrow.vector.VectorSchemaRoot} with large
 * (more than 2GB) buffers by {@link ArrowReader} and {@link ArrowWriter}.. To run this test, please
 * make sure there is at least 8GB free memory, and 8GB free.disk space in the system.
 */
public class ITTestIPCWithLargeArrowBuffers {

  private static final Logger logger =
      LoggerFactory.getLogger(ITTestIPCWithLargeArrowBuffers.class);

  // 4GB buffer size
  static final long BUFFER_SIZE = 4 * 1024 * 1024 * 1024L;

  static final int DICTIONARY_VECTOR_SIZE = (int) (BUFFER_SIZE / BigIntVector.TYPE_WIDTH);

  static final int ENCODED_VECTOR_SIZE = (int) (BUFFER_SIZE / IntVector.TYPE_WIDTH);

  static final String FILE_NAME = "largeArrowData.data";

  static final long DICTIONARY_ID = 123L;

  static final ArrowType.Int ENCODED_VECTOR_TYPE = new ArrowType.Int(32, true);

  static final DictionaryEncoding DICTIONARY_ENCODING =
      new DictionaryEncoding(DICTIONARY_ID, false, ENCODED_VECTOR_TYPE);

  static final FieldType ENCODED_FIELD_TYPE =
      new FieldType(true, ENCODED_VECTOR_TYPE, DICTIONARY_ENCODING, null);

  static final Field ENCODED_VECTOR_FIELD = new Field("encoded vector", ENCODED_FIELD_TYPE, null);

  private void testWriteLargeArrowData(boolean streamMode) throws IOException {
    // simulate encoding big int as int
    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        BigIntVector dictVector = new BigIntVector("dic vector", allocator);
        FileOutputStream out = new FileOutputStream(FILE_NAME);
        IntVector encodedVector = (IntVector) ENCODED_VECTOR_FIELD.createVector(allocator)) {

      // prepare dictionary provider.
      DictionaryProvider.MapDictionaryProvider provider =
          new DictionaryProvider.MapDictionaryProvider();
      Dictionary dictionary = new Dictionary(dictVector, DICTIONARY_ENCODING);
      provider.put(dictionary);

      // populate the dictionary vector
      dictVector.allocateNew(DICTIONARY_VECTOR_SIZE);
      for (int i = 0; i < DICTIONARY_VECTOR_SIZE; i++) {
        dictVector.set(i, i);
      }
      dictVector.setValueCount(DICTIONARY_VECTOR_SIZE);
      assertTrue(dictVector.getDataBuffer().capacity() > Integer.MAX_VALUE);
      logger.trace("Populating dictionary vector finished");

      // populate the encoded vector
      encodedVector.allocateNew(ENCODED_VECTOR_SIZE);
      for (int i = 0; i < ENCODED_VECTOR_SIZE; i++) {
        encodedVector.set(i, i % DICTIONARY_VECTOR_SIZE);
      }
      encodedVector.setValueCount(ENCODED_VECTOR_SIZE);
      assertTrue(encodedVector.getDataBuffer().capacity() > Integer.MAX_VALUE);
      logger.trace("Populating encoded vector finished");

      // build vector schema root and write data.
      try (VectorSchemaRoot root =
              new VectorSchemaRoot(
                  Arrays.asList(ENCODED_VECTOR_FIELD),
                  Arrays.asList(encodedVector),
                  ENCODED_VECTOR_SIZE);
          ArrowWriter writer =
              streamMode
                  ? new ArrowStreamWriter(root, provider, out)
                  : new ArrowFileWriter(root, provider, out.getChannel())) {
        writer.start();
        writer.writeBatch();
        writer.end();
        logger.trace("Writing data finished");
      }
    }

    assertTrue(new File(FILE_NAME).exists());
  }

  private void testReadLargeArrowData(boolean streamMode) throws IOException {
    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        FileInputStream in = new FileInputStream(FILE_NAME);
        ArrowReader reader =
            streamMode
                ? new ArrowStreamReader(in, allocator)
                : new ArrowFileReader(in.getChannel(), allocator)) {

      // verify schema
      Schema readSchema = reader.getVectorSchemaRoot().getSchema();
      assertEquals(1, readSchema.getFields().size());
      assertEquals(ENCODED_VECTOR_FIELD, readSchema.getFields().get(0));
      logger.trace("Verifying schema finished");

      // verify vector schema root
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();

      assertEquals(ENCODED_VECTOR_SIZE, root.getRowCount());
      assertEquals(1, root.getFieldVectors().size());
      assertTrue(root.getFieldVectors().get(0) instanceof IntVector);

      IntVector encodedVector = (IntVector) root.getVector(0);
      for (int i = 0; i < ENCODED_VECTOR_SIZE; i++) {
        assertEquals(i % DICTIONARY_VECTOR_SIZE, encodedVector.get(i));
      }
      logger.trace("Verifying encoded vector finished");

      // verify dictionary
      Map<Long, Dictionary> dictVectors = reader.getDictionaryVectors();
      assertEquals(1, dictVectors.size());
      Dictionary dictionary = dictVectors.get(DICTIONARY_ID);
      assertNotNull(dictionary);

      assertTrue(dictionary.getVector() instanceof BigIntVector);
      BigIntVector dictVector = (BigIntVector) dictionary.getVector();
      assertEquals(DICTIONARY_VECTOR_SIZE, dictVector.getValueCount());
      for (int i = 0; i < DICTIONARY_VECTOR_SIZE; i++) {
        assertEquals(i, dictVector.get(i));
      }
      logger.trace("Verifying dictionary vector finished");

      // ensure no more data available
      assertFalse(reader.loadNextBatch());
    } finally {
      File dataFile = new File(FILE_NAME);
      dataFile.delete();
      assertFalse(dataFile.exists());
    }
  }

  @Test
  public void testIPC() throws IOException {
    logger.trace("Start testing reading/writing large arrow stream data");
    testWriteLargeArrowData(true);
    testReadLargeArrowData(true);
    logger.trace("Finish testing reading/writing large arrow stream data");

    logger.trace("Start testing reading/writing large arrow file data");
    testWriteLargeArrowData(false);
    testReadLargeArrowData(false);
    logger.trace("Finish testing reading/writing large arrow file data");
  }
}
