/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed 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.
 */

#include <cmath>
#include <vector>

#include <fbjni/fbjni.h>

#include "expect.h"

using namespace facebook::jni;

local_ref<jbooleanArray> testMakeBoolArray(alias_ref<jclass>, jint size) {
  return make_boolean_array(size);
}

local_ref<jbyteArray> testMakeByteArray(alias_ref<jclass>, jint size) {
  return make_byte_array(size);
}

local_ref<jcharArray> testMakeCharArray(alias_ref<jclass>, jint size) {
  return make_char_array(size);
}

local_ref<jshortArray> testMakeShortArray(alias_ref<jclass>, jint size) {
  return make_short_array(size);
}

local_ref<jintArray> testMakeIntArray(alias_ref<jclass>, jint size) {
  return make_int_array(size);
}

local_ref<jlongArray> testMakeLongArray(alias_ref<jclass>, jint size) {
  return make_long_array(size);
}

local_ref<jfloatArray> testMakeFloatArray(alias_ref<jclass>, jint size) {
  return make_float_array(size);
}

local_ref<jdoubleArray> testMakeDoubleArray(alias_ref<jclass>, jint size) {
  return make_double_array(size);
}

jboolean testGetSetBooleanArray(
    alias_ref<jclass>,
    alias_ref<jbooleanArray> array) {
  EXPECT(array);

  auto n = array->size();
  EXPECT(n == 2);

  auto vec = std::vector<jboolean>(n);
  array->getRegion(0, n, vec.data());
  auto smartbuf = array->getRegion(0, n);

  EXPECT(!vec[0] && vec[1]);
  EXPECT(!smartbuf[0] && smartbuf[1]);

  for (auto i = 0u; i < n; ++i) {
    smartbuf[i] = !smartbuf[i];
  }

  array->setRegion(0, n, smartbuf.get());

  return JNI_TRUE;
}

jboolean testPinBooleanArray(
    alias_ref<jclass>,
    alias_ref<jbooleanArray> array) {
  EXPECT(array);

  auto n = static_cast<jboolean>(array->size());
  auto pinned = array->pin();

  EXPECT(!pinned[0] && pinned[1]);

  for (auto i = 0; i < n; ++i) {
    pinned[i] = !pinned[i];
  }

  return JNI_TRUE;
}

constexpr double kEps = 1e-3;

template <typename JArrayType>
jboolean testGetSetArray(alias_ref<jclass>, alias_ref<JArrayType> array) {
  EXPECT(array);
  int n = array->size();
  auto vec = std::vector<typename jtype_traits<JArrayType>::entry_type>(n);
  array->getRegion(0, n, vec.data());
  auto smartbuf = array->getRegion(0, n);

  for (auto i = 0; i < n; ++i) {
    EXPECT(std::abs(static_cast<double>(vec[i] - i)) < kEps);
    EXPECT(std::abs(static_cast<double>(smartbuf[i] - i)) < kEps);
  }

  for (auto i = 0; i < n; ++i) {
    smartbuf[i] *= 2;
  }

  array->setRegion(0, n, smartbuf.get());

  return JNI_TRUE;
}

template <typename JArrayType>
jboolean testPinArray(alias_ref<jclass>, alias_ref<JArrayType> array) {
  EXPECT(array);

  int n = array->size();
  auto pinned = array->pin();

  for (auto i = 0; i < n; ++i) {
    EXPECT(std::abs(static_cast<double>(pinned[i] - i)) < kEps);
  }

  for (auto i = 0; i < n; ++i) {
    pinned[i] *= 2;
  }

  return JNI_TRUE;
}

template <typename JArrayType>
jboolean testPinArrayRegion(alias_ref<jclass>, alias_ref<JArrayType> array) {
  EXPECT(array);
  EXPECT(array->size() > 5);

  int splits[] = {0, 1, 3, static_cast<int>(array->size())};
  for (int i = 0; i < 3; i++) {
    auto pinned = array->pinRegion(splits[i], splits[i + 1] - splits[i]);
    for (int j = 0; j < static_cast<int>(pinned.size()); j++) {
      EXPECT(std::abs(static_cast<double>(pinned[j] - j - splits[i])) < kEps);
      pinned[j] *= 2;
    }
  }
  return JNI_TRUE;
}

template <typename JArrayType>
jboolean testPinArrayCritical(alias_ref<jclass>, alias_ref<JArrayType> array) {
  EXPECT(array);

  int n = array->size();
  auto pinned = array->pinCritical();

  for (auto i = 0; i < n; ++i) {
    EXPECT(std::abs(static_cast<double>(pinned[i] - i)) < kEps);
  }

  for (auto i = 0; i < n; ++i) {
    pinned[i] *= 2;
  }

  return JNI_TRUE;
}

jboolean testIndexOutOfBoundsInRegions(alias_ref<jclass>) {
  constexpr auto N = 7;
  constexpr auto TOO_MUCH = 10;
  constexpr auto NEGATIVE = -1;

  auto array = make_int_array(N);

  try {
    auto buf = array->getRegion(TOO_MUCH, N);
    EXPECT(false);
  } catch (JniException&) {
  }

  try {
    auto buf = array->getRegion(NEGATIVE, N);
    EXPECT(false);
  } catch (JniException&) {
  }

  try {
    auto vec = std::vector<jint>(TOO_MUCH);
    array->setRegion(0, vec.size(), vec.data());
    EXPECT(false);
  } catch (JniException&) {
  }

  try {
    auto vec = std::vector<jint>(1);
    array->setRegion(NEGATIVE, vec.size(), vec.data());
    EXPECT(false);
  } catch (JniException&) {
  }

  return JNI_TRUE;
}

jboolean TestBooleanArrayIndexing(
    alias_ref<jobject> self,
    alias_ref<jbooleanArray> input,
    jint idx) {
  auto array = input->pin();
  jboolean value = array[idx];
  return value;
}

jint TestIntegerArrayIndexing(
    alias_ref<jobject> self,
    alias_ref<jintArray> input,
    jint idx) {
  auto array = input->pin();
  jint value = array[idx];
  return value;
}

jsize TestIntegerArraySize(
    alias_ref<jobject> self,
    alias_ref<jintArray> input) {
  auto array = input->pin();
  jsize size = array.size();
  return size;
}

alias_ref<jintArray> TestIntegerArrayIncrement(
    alias_ref<jobject> self,
    alias_ref<jintArray> input) {
  auto array = input->pin();
  for (size_t ii = 0; ii < array.size(); ii++) {
    array[ii]++;
  }
  return input;
}

void TestIntegerArrayMoveAssignment(
    alias_ref<jobject> self,
    alias_ref<jintArray> input) {
  auto array = input->pin();
  array[0] = 0;
  array.release();
}

jboolean isPinnedArrayACopy(alias_ref<jobject>, alias_ref<jintArray> input) {
  return input->pin().isCopy();
}

jboolean testCopiedPinnedArray(alias_ref<jobject>, alias_ref<jintArray> input) {
  EXPECT(input->size() > 0);
  input->pin()[0] = 100;
  auto pin = input->pin();
  EXPECT(pin.isCopy());
  EXPECT(pin[0] == input->pin()[0]);

  pin[0] = 200;
  EXPECT(input->pin()[0] == 100);

  pin.commit();
  EXPECT(input->pin()[0] == 200);
  pin[0] = 300;
  EXPECT(input->pin()[0] == 200);
  pin.abort();
  EXPECT(input->pin()[0] == 200);
  pin = input->pin();
  pin[0] = 400;
  pin.release();
  EXPECT(input->pin()[0] == 400);
  return JNI_TRUE;
}

jboolean testNonCopiedPinnedArray(
    alias_ref<jobject>,
    alias_ref<jintArray> input) {
  EXPECT(input->size() > 0);
  auto pin = input->pin();
  EXPECT(!pin.isCopy());
  EXPECT(pin[0] == input->pin()[0]);
  pin.commit();
  EXPECT(pin[0] == input->pin()[0]);
  pin[0] = 100;
  EXPECT(pin[0] == input->pin()[0]);
  input->pin()[0] = 200;
  EXPECT(pin[0] == input->pin()[0]);
  pin.abort();
  return JNI_TRUE;
}

void RegisterPrimitiveArrayTests() {
  registerNatives(
      "com/facebook/jni/PrimitiveArrayTests",
      {
          makeNativeMethod("nativeTestMakeBooleanArray", testMakeBoolArray),
          makeNativeMethod("nativeTestMakeByteArray", testMakeByteArray),
          makeNativeMethod("nativeTestMakeCharArray", testMakeCharArray),
          makeNativeMethod("nativeTestMakeShortArray", testMakeShortArray),
          makeNativeMethod("nativeTestMakeIntArray", testMakeIntArray),
          makeNativeMethod("nativeTestMakeLongArray", testMakeLongArray),
          makeNativeMethod("nativeTestMakeFloatArray", testMakeFloatArray),
          makeNativeMethod("nativeTestMakeDoubleArray", testMakeDoubleArray),

          makeNativeMethod(
              "nativeTestGetSetBooleanArray", testGetSetBooleanArray),
          makeNativeMethod(
              "nativeTestGetSetByteArray", testGetSetArray<jbyteArray>),
          makeNativeMethod(
              "nativeTestGetSetCharArray", testGetSetArray<jcharArray>),
          makeNativeMethod(
              "nativeTestGetSetShortArray", testGetSetArray<jshortArray>),
          makeNativeMethod(
              "nativeTestGetSetIntArray", testGetSetArray<jintArray>),
          makeNativeMethod(
              "nativeTestGetSetLongArray", testGetSetArray<jlongArray>),
          makeNativeMethod(
              "nativeTestGetSetFloatArray", testGetSetArray<jfloatArray>),
          makeNativeMethod(
              "nativeTestGetSetDoubleArray", testGetSetArray<jdoubleArray>),

          makeNativeMethod("nativeTestPinBooleanArray", testPinBooleanArray),
          makeNativeMethod("nativeTestPinByteArray", testPinArray<jbyteArray>),
          makeNativeMethod("nativeTestPinCharArray", testPinArray<jcharArray>),
          makeNativeMethod(
              "nativeTestPinShortArray", testPinArray<jshortArray>),
          makeNativeMethod("nativeTestPinIntArray", testPinArray<jintArray>),
          makeNativeMethod("nativeTestPinLongArray", testPinArray<jlongArray>),
          makeNativeMethod(
              "nativeTestPinFloatArray", testPinArray<jfloatArray>),
          makeNativeMethod(
              "nativeTestPinDoubleArray", testPinArray<jdoubleArray>),

          makeNativeMethod(
              "nativeTestPinByteArrayRegion", testPinArrayRegion<jbyteArray>),
          makeNativeMethod(
              "nativeTestPinCharArrayRegion", testPinArrayRegion<jcharArray>),
          makeNativeMethod(
              "nativeTestPinShortArrayRegion", testPinArrayRegion<jshortArray>),
          makeNativeMethod(
              "nativeTestPinIntArrayRegion", testPinArrayRegion<jintArray>),
          makeNativeMethod(
              "nativeTestPinLongArrayRegion", testPinArrayRegion<jlongArray>),
          makeNativeMethod(
              "nativeTestPinFloatArrayRegion", testPinArrayRegion<jfloatArray>),
          makeNativeMethod(
              "nativeTestPinDoubleArrayRegion",
              testPinArrayRegion<jdoubleArray>),

          makeNativeMethod(
              "nativeTestPinByteArrayCritical",
              testPinArrayCritical<jbyteArray>),
          makeNativeMethod(
              "nativeTestPinCharArrayCritical",
              testPinArrayCritical<jcharArray>),
          makeNativeMethod(
              "nativeTestPinShortArrayCritical",
              testPinArrayCritical<jshortArray>),
          makeNativeMethod(
              "nativeTestPinIntArrayCritical", testPinArrayCritical<jintArray>),
          makeNativeMethod(
              "nativeTestPinLongArrayCritical",
              testPinArrayCritical<jlongArray>),
          makeNativeMethod(
              "nativeTestPinFloatArrayCritical",
              testPinArrayCritical<jfloatArray>),
          makeNativeMethod(
              "nativeTestPinDoubleArrayCritical",
              testPinArrayCritical<jdoubleArray>),

          makeNativeMethod(
              "nativeTestIndexOutOfBoundsInRegions",
              testIndexOutOfBoundsInRegions),

          makeNativeMethod(
              "nativeTestBooleanArrayIndexing", TestBooleanArrayIndexing),
          makeNativeMethod(
              "nativeTestIntegerArrayIndexing", TestIntegerArrayIndexing),
          makeNativeMethod("nativeTestIntegerArraySize", TestIntegerArraySize),
          makeNativeMethod(
              "nativeTestIntegerArrayIncrement", TestIntegerArrayIncrement),
          makeNativeMethod(
              "nativeTestIntegerArrayMoveAssignment",
              TestIntegerArrayMoveAssignment),

          makeNativeMethod("nativeIsPinnedArrayACopy", isPinnedArrayACopy),
          makeNativeMethod(
              "nativeTestCopiedPinnedArray", testCopiedPinnedArray),
          makeNativeMethod(
              "nativeTestNonCopiedPinnedArray", testNonCopiedPinnedArray),
      });
}