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

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BaseValueVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.holders.NullableIntHolder;
import org.apache.arrow.vector.holders.NullableVarBinaryHolder;
import org.apache.arrow.vector.types.UnionMode;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.junit.jupiter.api.Test;

public class TestDenseUnionBufferSize {
  @Test
  public void testBufferSize() {
    try (BufferAllocator allocator = new RootAllocator();
        DenseUnionVector duv =
            new DenseUnionVector(
                "duv",
                allocator,
                FieldType.nullable(new ArrowType.Union(UnionMode.Dense, null)),
                null)) {

      byte aTypeId = 42;
      byte bTypeId = 7;

      duv.addVector(
          aTypeId,
          new IntVector("a", FieldType.notNullable(new ArrowType.Int(32, true)), allocator));
      duv.addVector(
          bTypeId,
          new VarBinaryVector("b", FieldType.notNullable(new ArrowType.Binary()), allocator));

      NullableIntHolder intHolder = new NullableIntHolder();
      NullableVarBinaryHolder varBinaryHolder = new NullableVarBinaryHolder();

      int aCount = BaseValueVector.INITIAL_VALUE_ALLOCATION + 1;
      for (int i = 0; i < aCount; i++) {
        duv.setTypeId(i, aTypeId);
        duv.setSafe(i, intHolder);
      }

      int bCount = 2;
      for (int i = 0; i < bCount; i++) {
        duv.setTypeId(i + aCount, bTypeId);
        duv.setSafe(i + aCount, varBinaryHolder);
      }

      int count = aCount + bCount;
      duv.setValueCount(count);

      // will not necessarily see an error unless bounds checking is on.
      assertDoesNotThrow(duv::getBufferSize);

      IntVector intVector = duv.getIntVector(aTypeId);
      VarBinaryVector varBinaryVector = duv.getVarBinaryVector(bTypeId);

      long overhead = DenseUnionVector.TYPE_WIDTH + DenseUnionVector.OFFSET_WIDTH;

      assertEquals(
          overhead * count + intVector.getBufferSize() + varBinaryVector.getBufferSize(),
          duv.getBufferSize());

      assertEquals(
          overhead * (aCount + 1)
              + intVector.getBufferSizeFor(aCount)
              + varBinaryVector.getBufferSizeFor(1),
          duv.getBufferSizeFor(aCount + 1));
    }
  }
}
