/*
 * 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.table;

import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH;

import java.math.BigDecimal;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.DateMilliVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.DurationVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.GenerateSampleData;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.IntervalDayVector;
import org.apache.arrow.vector.IntervalMonthDayNanoVector;
import org.apache.arrow.vector.IntervalYearVector;
import org.apache.arrow.vector.LargeVarBinaryVector;
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.SmallIntVector;
import org.apache.arrow.vector.TimeMicroVector;
import org.apache.arrow.vector.TimeMilliVector;
import org.apache.arrow.vector.TimeNanoVector;
import org.apache.arrow.vector.TimeSecVector;
import org.apache.arrow.vector.TimeStampMicroTZVector;
import org.apache.arrow.vector.TimeStampMicroVector;
import org.apache.arrow.vector.TimeStampMilliTZVector;
import org.apache.arrow.vector.TimeStampMilliVector;
import org.apache.arrow.vector.TimeStampNanoTZVector;
import org.apache.arrow.vector.TimeStampNanoVector;
import org.apache.arrow.vector.TimeStampSecTZVector;
import org.apache.arrow.vector.TimeStampSecVector;
import org.apache.arrow.vector.TinyIntVector;
import org.apache.arrow.vector.UInt1Vector;
import org.apache.arrow.vector.UInt2Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.UInt8Vector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.complex.impl.NullableStructWriter;
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
import org.apache.arrow.vector.complex.writer.Float8Writer;
import org.apache.arrow.vector.complex.writer.IntWriter;
import org.apache.arrow.vector.holders.NullableUInt4Holder;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;

public class TestUtils {

  public static final String INT_VECTOR_NAME = "intCol";
  public static final String INT_VECTOR_NAME_1 = "intCol1";
  public static final String VARCHAR_VECTOR_NAME_1 = "varcharCol1";
  public static final String VARBINARY_VECTOR_NAME_1 = "varbinaryCol1";
  public static final String FIXEDBINARY_VECTOR_NAME_1 = "varbinaryCol1";
  public static final String INT_VECTOR_NAME_2 = "intCol2";
  public static final String INT_LIST_VECTOR_NAME = "int list vector";
  public static final String BIGINT_INT_MAP_VECTOR_NAME = "bigint-int map vector";
  public static final String STRUCT_VECTOR_NAME = "struct_vector";
  public static final String UNION_VECTOR_NAME = "union_vector";

  /**
   * Returns a list of two IntVectors to be used to instantiate Tables for testing. Each IntVector
   * has two values set.
   */
  static List<FieldVector> twoIntColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    IntVector v2 = new IntVector(INT_VECTOR_NAME_2, allocator);
    v2.allocateNew(2);
    v2.set(0, 3);
    v2.set(1, 4);
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  /**
   * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first
   * vector is an IntVector and the second is a VarCharVector. Each vector has two values set.
   */
  static List<FieldVector> intPlusVarcharColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    VarCharVector v2 = new VarCharVector(VARCHAR_VECTOR_NAME_1, allocator);
    v2.allocateNew(2);
    v2.set(0, "one".getBytes(StandardCharsets.UTF_8));
    v2.set(1, "two".getBytes(StandardCharsets.UTF_8));
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  /**
   * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first
   * vector is an IntVector and the second is a LargeVarCharVector. Each vector has two values set.
   */
  static List<FieldVector> intPlusLargeVarcharColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    LargeVarCharVector v2 = new LargeVarCharVector(VARCHAR_VECTOR_NAME_1, allocator);
    v2.allocateNew(2);
    v2.set(0, "one".getBytes(StandardCharsets.UTF_8));
    v2.set(1, "two".getBytes(StandardCharsets.UTF_8));
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  /**
   * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first
   * vector is an IntVector and the second is a VarBinaryVector. Each vector has two values set. The
   * large binary vectors values are "one" and "two" encoded with UTF-8
   */
  static List<FieldVector> intPlusVarBinaryColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    VarBinaryVector v2 = new VarBinaryVector(VARBINARY_VECTOR_NAME_1, allocator);
    v2.allocateNew(2);
    v2.set(0, "one".getBytes(StandardCharsets.UTF_8));
    v2.set(1, "two".getBytes(StandardCharsets.UTF_8));
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  /**
   * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first
   * vector is an IntVector and the second is a VarBinaryVector. Each vector has two values set. The
   * large binary vectors values are "one" and "two" encoded with UTF-8
   */
  static List<FieldVector> intPlusLargeVarBinaryColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    LargeVarBinaryVector v2 = new LargeVarBinaryVector(VARBINARY_VECTOR_NAME_1, allocator);
    v2.allocateNew(2);
    v2.set(0, "one".getBytes(StandardCharsets.UTF_8));
    v2.set(1, "two".getBytes(StandardCharsets.UTF_8));
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  /**
   * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first
   * vector is an IntVector and the second is a FixedSizeBinary vector. Each vector has two values
   * set. The large binary vectors values are "one" and "two" encoded with UTF-8
   */
  static List<FieldVector> intPlusFixedBinaryColumns(BufferAllocator allocator) {
    List<FieldVector> vectorList = new ArrayList<>();
    IntVector v1 = getSimpleIntVector(allocator);
    FixedSizeBinaryVector v2 = new FixedSizeBinaryVector(FIXEDBINARY_VECTOR_NAME_1, allocator, 3);
    v2.allocateNew(2);
    v2.set(0, "one".getBytes(StandardCharsets.UTF_8));
    v2.set(1, "two".getBytes(StandardCharsets.UTF_8));
    v2.setValueCount(2);
    vectorList.add(v1);
    vectorList.add(v2);
    return vectorList;
  }

  private static IntVector getSimpleIntVector(BufferAllocator allocator) {
    IntVector v1 = new IntVector(INT_VECTOR_NAME_1, allocator);
    v1.allocateNew(2);
    v1.set(0, 1);
    v1.set(1, 2);
    v1.setValueCount(2);
    return v1;
  }

  /**
   * Returns a list of fixed-width vectors for testing. It includes
   *
   * <ol>
   *   <li>all integral and floating point types
   *   <li>all basic times and timestamps (second, milli, micro, nano
   * </ol>
   *
   * The vector names are based on their type name (e.g. BigIntVector is called "bigInt_vector"
   */
  static List<FieldVector> fixedWidthVectors(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    numericVectors(vectors, allocator, rowCount);
    simpleTemporalVectors(vectors, allocator, rowCount);
    return vectors;
  }

  /**
   * Returns a list of all integral and floating point vectors. The vector names are based on their
   * type name (e.g. BigIntVector is called "bigInt_vector"
   */
  static List<FieldVector> numericVectors(
      List<FieldVector> vectors, BufferAllocator allocator, int rowCount) {
    vectors.add(new IntVector("int_vector", allocator));
    vectors.add(new BigIntVector("bigInt_vector", allocator));
    vectors.add(new SmallIntVector("smallInt_vector", allocator));
    vectors.add(new TinyIntVector("tinyInt_vector", allocator));
    vectors.add(new UInt1Vector("uInt1_vector", allocator));
    vectors.add(new UInt2Vector("uInt2_vector", allocator));
    vectors.add(new UInt4Vector("uInt4_vector", allocator));
    vectors.add(new UInt8Vector("uInt8_vector", allocator));
    vectors.add(new Float4Vector("float4_vector", allocator));
    vectors.add(new Float8Vector("float8_vector", allocator));
    vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount));
    return vectors;
  }

  static List<FieldVector> numericVectors(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    return numericVectors(vectors, allocator, rowCount);
  }

  static List<FieldVector> simpleTemporalVectors(
      List<FieldVector> vectors, BufferAllocator allocator, int rowCount) {
    vectors.add(new TimeSecVector("timeSec_vector", allocator));
    vectors.add(new TimeMilliVector("timeMilli_vector", allocator));
    vectors.add(new TimeMicroVector("timeMicro_vector", allocator));
    vectors.add(new TimeNanoVector("timeNano_vector", allocator));

    vectors.add(new TimeStampSecVector("timeStampSec_vector", allocator));
    vectors.add(new TimeStampMilliVector("timeStampMilli_vector", allocator));
    vectors.add(new TimeStampMicroVector("timeStampMicro_vector", allocator));
    vectors.add(new TimeStampNanoVector("timeStampNano_vector", allocator));

    vectors.add(new DateMilliVector("dateMilli_vector", allocator));
    vectors.add(new DateDayVector("dateDay_vector", allocator));

    vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount));
    return vectors;
  }

  static List<FieldVector> simpleTemporalVectors(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    return simpleTemporalVectors(vectors, allocator, rowCount);
  }

  static List<FieldVector> timezoneTemporalVectors(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    vectors.add(new TimeStampSecTZVector("timeStampSecTZ_vector", allocator, "UTC"));
    vectors.add(new TimeStampMilliTZVector("timeStampMilliTZ_vector", allocator, "UTC"));
    vectors.add(new TimeStampMicroTZVector("timeStampMicroTZ_vector", allocator, "UTC"));
    vectors.add(new TimeStampNanoTZVector("timeStampNanoTZ_vector", allocator, "UTC"));
    vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount));
    return vectors;
  }

  static List<FieldVector> intervalVectors(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    vectors.add(new IntervalDayVector("intervalDay_vector", allocator));
    vectors.add(new IntervalYearVector("intervalYear_vector", allocator));
    vectors.add(new IntervalMonthDayNanoVector("intervalMonthDayNano_vector", allocator));
    vectors.add(
        new DurationVector(
            "duration_vector",
            new FieldType(true, new ArrowType.Duration(TimeUnit.SECOND), null),
            allocator));
    vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount));
    return vectors;
  }

  /** Returns a list vector of ints. */
  static ListVector simpleListVector(BufferAllocator allocator) {
    ListVector listVector = ListVector.empty(INT_LIST_VECTOR_NAME, allocator);
    final int innerCount = 80; // total number of values
    final int outerCount = 8; // total number of values in the list vector itself
    final int listLength = innerCount / outerCount; // length of an individual list

    Types.MinorType type = Types.MinorType.INT;
    listVector.addOrGetVector(FieldType.nullable(type.getType()));

    listVector.allocateNew();
    IntVector dataVector = (IntVector) listVector.getDataVector();

    for (int i = 0; i < innerCount; i++) {
      dataVector.set(i, i);
    }
    dataVector.setValueCount(innerCount);

    for (int i = 0; i < outerCount; i++) {
      BitVectorHelper.setBit(listVector.getValidityBuffer(), i);
      listVector.getOffsetBuffer().setInt(i * OFFSET_WIDTH, i * listLength);
      listVector.getOffsetBuffer().setInt((i + 1) * OFFSET_WIDTH, (i + 1) * listLength);
    }
    listVector.setLastSet(outerCount - 1);
    listVector.setValueCount(outerCount);
    return listVector;
  }

  static StructVector simpleStructVector(BufferAllocator allocator) {
    final String INT_COL = "struct_int_child";
    final String FLT_COL = "struct_flt_child";
    StructVector structVector = StructVector.empty(STRUCT_VECTOR_NAME, allocator);
    final int size = 6; // number of structs

    NullableStructWriter structWriter = structVector.getWriter();
    structVector.addOrGet(
        INT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
    structVector.addOrGet(
        FLT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
    structVector.allocateNew();
    IntWriter intWriter = structWriter.integer(INT_COL);
    Float8Writer float8Writer = structWriter.float8(FLT_COL);

    for (int i = 0; i < size; i++) {
      structWriter.setPosition(i);
      structWriter.start();
      intWriter.writeInt(i);
      float8Writer.writeFloat8(i * .1);
      structWriter.end();
    }

    structWriter.setValueCount(size);

    return structVector;
  }

  /** Returns a MapVector of longs to doubles. */
  static MapVector simpleMapVector(BufferAllocator allocator) {
    MapVector mapVector = MapVector.empty(BIGINT_INT_MAP_VECTOR_NAME, allocator, false);
    mapVector.allocateNew();
    int count = 5;
    UnionMapWriter mapWriter = mapVector.getWriter();
    for (int i = 0; i < count; i++) {
      mapWriter.startMap();
      for (int j = 0; j < i + 1; j++) {
        mapWriter.startEntry();
        mapWriter.key().bigInt().writeBigInt(j);
        mapWriter.value().integer().writeInt(j);
        mapWriter.endEntry();
      }
      mapWriter.endMap();
    }
    mapWriter.setValueCount(count);
    return mapVector;
  }

  static List<FieldVector> decimalVector(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    vectors.add(
        new DecimalVector(
            "decimal_vector",
            new FieldType(true, new ArrowType.Decimal(38, 10, 128), null),
            allocator));
    vectors.forEach(vec -> generateDecimalData((DecimalVector) vec, rowCount));
    return vectors;
  }

  static List<FieldVector> bitVector(BufferAllocator allocator, int rowCount) {
    List<FieldVector> vectors = new ArrayList<>();
    vectors.add(new BitVector("bit_vector", allocator));
    vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount));
    return vectors;
  }

  /** Returns a UnionVector. */
  static UnionVector simpleUnionVector(BufferAllocator allocator) {
    final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
    uInt4Holder.value = 100;
    uInt4Holder.isSet = 1;

    UnionVector unionVector = new UnionVector(UNION_VECTOR_NAME, allocator, null, null);
    unionVector.allocateNew();

    // write some data
    unionVector.setType(0, Types.MinorType.UINT4);
    unionVector.setSafe(0, uInt4Holder);
    unionVector.setType(2, Types.MinorType.UINT4);
    unionVector.setSafe(2, uInt4Holder);
    unionVector.setValueCount(4);
    return unionVector;
  }

  /** Returns a DenseUnionVector. */
  static DenseUnionVector simpleDenseUnionVector(BufferAllocator allocator) {
    final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
    uInt4Holder.value = 100;
    uInt4Holder.isSet = 1;

    DenseUnionVector unionVector = new DenseUnionVector(UNION_VECTOR_NAME, allocator, null, null);
    unionVector.allocateNew();

    // write some data
    byte uint4TypeId =
        unionVector.registerNewTypeId(Field.nullable("", Types.MinorType.UINT4.getType()));
    unionVector.setTypeId(0, uint4TypeId);
    unionVector.setSafe(0, uInt4Holder);
    unionVector.setTypeId(2, uint4TypeId);
    unionVector.setSafe(2, uInt4Holder);
    unionVector.setValueCount(4);
    return unionVector;
  }

  private static void generateDecimalData(DecimalVector vector, int valueCount) {
    final BigDecimal even = new BigDecimal("0.0543278923");
    final BigDecimal odd = new BigDecimal("2.0543278923");
    for (int i = 0; i < valueCount; i++) {
      if (i % 2 == 0) {
        vector.setSafe(i, even);
      } else {
        vector.setSafe(i, odd);
      }
    }
    vector.setValueCount(valueCount);
  }
}
