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

import static org.assertj.core.api.Assertions.assertThat;

import java.math.BigInteger;
import java.util.Arrays;
import org.apache.parquet.io.api.Binary;
import org.apache.parquet.schema.PrimitiveType;
import org.junit.jupiter.api.Test;

/**
 * Some utility methods for generating Binary values that length if fixed (e.g. for FIXED_LEN_BYTE_ARRAY or INT96).
 */
public class FixedBinaryTestUtils {
  public static Binary getFixedBinary(int length, BigInteger bigInt) {
    byte[] array = bigInt.toByteArray();
    if (array.length == length) {
      return Binary.fromConstantByteArray(array);
    } else if (array.length < length) {
      byte[] padded = new byte[length];
      int paddingLength = length - array.length;
      if (bigInt.signum() < 0) {
        Arrays.fill(padded, 0, paddingLength, (byte) 0xFF);
      } else {
        Arrays.fill(padded, 0, paddingLength, (byte) 0x00);
      }
      System.arraycopy(array, 0, padded, paddingLength, array.length);
      return Binary.fromConstantByteArray(padded);
    } else {
      throw new IllegalArgumentException("Specified BigInteger (" + bigInt + ") is too long for fixed bytes ("
          + array.length + '>' + length + ')');
    }
  }

  public static Binary getFixedBinary(PrimitiveType type, BigInteger bigInt) {
    switch (type.getPrimitiveTypeName()) {
      case FIXED_LEN_BYTE_ARRAY:
        return getFixedBinary(type.getTypeLength(), bigInt);
      case INT96:
        return getFixedBinary(12, bigInt);
      case BINARY:
        return Binary.fromConstantByteArray(bigInt.toByteArray());
      default:
        throw new IllegalArgumentException("Type " + type + " cannot be represented by a Binary");
    }
  }

  @Test
  public void testGetFixedBinary() {
    assertThat(getFixedBinary(10, BigInteger.valueOf(Integer.MIN_VALUE)).getBytes())
        .isEqualTo(b(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x80, 0x00, 0x00, 0x00));
    assertThat(getFixedBinary(11, BigInteger.valueOf(-1)).getBytes())
        .isEqualTo(b(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF));
    assertThat(getFixedBinary(12, BigInteger.valueOf(0)).getBytes())
        .isEqualTo(b(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00));
    assertThat(getFixedBinary(13, BigInteger.valueOf(1)).getBytes())
        .isEqualTo(b(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01));
    assertThat(getFixedBinary(14, BigInteger.valueOf(Integer.MAX_VALUE)).getBytes())
        .isEqualTo(b(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7F, 0xFF, 0xFF, 0xFF));
  }

  private static byte[] b(int... bytes) {
    byte[] arr = new byte[bytes.length];
    for (int i = 0, n = arr.length; i < n; ++i) {
      arr[i] = (byte) bytes[i];
    }
    return arr;
  }
}
