// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// IWYU pragma: private, include "third_party/jni_zero/jni_zero.h"

#ifndef JNI_ZERO_JNI_WRAPPERS_H_
#define JNI_ZERO_JNI_WRAPPERS_H_

#include <jni.h>

#include <array>
#include <cstdint>
#include <initializer_list>
#include <iterator>
#include <span>
#include <string_view>
#include <vector>

#include "third_party/jni_zero/java_refs.h"
#include "third_party/jni_zero/jni_methods.h"
#include "third_party/jni_zero/logging.h"

// Workaround for clang-format bug.
#define JNI_LIFETIME_BOUND [[clang::lifetimebound]]

// Wrapper used to receive int when calling Java from native.
// The wrapper disallows automatic conversion of long to int.
// This is to avoid a common anti-pattern where a Java int is used
// to receive a native pointer. Please use a Java long to receive
// native pointers, so that the code works on both 32-bit and 64-bit
// platforms. Note the wrapper allows other lossy conversions into
// int32_t (the underlying type of jint) that could be considered
// anti-patterns, such as from size_t.

// Checking is only done in debugging builds.

#ifdef NDEBUG

typedef int32_t JniIntWrapper;

// This inline is sufficiently trivial that it does not change the
// final code generated by g++.
inline int32_t as_jint(JniIntWrapper wrapper) {
  return wrapper;
}

#else

class JniIntWrapper {
 public:
  JniIntWrapper() : i_(0) {}
  JniIntWrapper(int i) : i_(i) {}
  JniIntWrapper(const JniIntWrapper& ji) : i_(ji.i_) {}
  template <class T>
  JniIntWrapper(const T& t) : i_(t) {}
  int32_t as_jint() const { return i_; }

 private:
  // If you get an "is private" error at the line below it is because you used
  // an implicit conversion to convert a long to an int when calling Java.
  // We disallow this, as a common anti-pattern allows converting a native
  // pointer (intptr_t) to a Java int. Please use a Java long to represent
  // a native pointer. If you want a lossy conversion, please use an
  // explicit conversion in your C++ code. Note an error is only seen when
  // compiling on a 64-bit platform, as intptr_t is indistinguishable from
  // int on 32-bit platforms.
  JniIntWrapper(long);
  int32_t i_;
};

inline int32_t as_jint(const JniIntWrapper& wrapper) {
  return wrapper.as_jint();
}

#endif  // NDEBUG

namespace jni_zero {

// JArrayViewBase is the base class for both JArrayViewCritical and JArrayView.
//
// JArrayView is a wrapper for a jobjectArray which supports a few library
// functions that get the length of the array and get one or all elements of the
// array. JArrayViewCritical is a wrapper for a primitive jarray (e.g.
// jbooleanArray, jintArray, etc.) backed by GetPrimitiveArrayCritical.
//
// JArrayView and JArrayViewCritical also support input iteration, allowing Java
// arrays to be iterated over with a range-based for loop, or used with
// <algorithm> functions that accept input iterators.
//
// The wrapper holds a local reference to the array and only queries the size of
// the array once, so must only be used as a stack-based object from the current
// thread.
//
// Note that this does *not* update the contents of the array if you mutate the
// returned array elements.
template <typename T>
class JArrayViewBase {
 public:
  ~JArrayViewBase() = default;

  // Get the number of elements in this JArray.
  int32_t length() const noexcept { return length_; }

  // Get the number of elements in this JArray.
  size_t size() const noexcept { return static_cast<size_t>(length_); }

  bool empty() const noexcept { return length_ == 0; }

 protected:
  JArrayViewBase(JNIEnv* env, JArray<T> array)
      : env_(env), array_(array), length_(env->GetArrayLength(array)) {}

  JArrayViewBase(JNIEnv* env, JArray<T> array, int32_t length)
      : env_(env), array_(array), length_(length) {}

  JNIEnv* env_;
  JArray<T> array_;
  int32_t length_;
};

template <typename T>
class JArrayView;

template <typename T>
class JArrayViewCritical;

// Wrapper for a jobjectArray.
template <typename T>
  requires internal::IsJobject<T>
class JArrayView<T> : public JArrayViewBase<T> {
  using JArrayViewBase<T>::env_;
  using JArrayViewBase<T>::array_;
  using JArrayViewBase<T>::length_;

 public:
  class iterator {
   public:
    // We can only be an input iterator, as all richer iterator types must
    // implement the multipass guarantee (always returning the same object for
    // the same iterator position), which is not practical when returning
    // temporary objects.
    using iterator_category = std::input_iterator_tag;

    using difference_type = ptrdiff_t;

    using value_type = ScopedJavaLocalRef<T>;

    // It doesn't make sense to return a reference type as the iterator creates
    // temporary wrapper objects when dereferenced. Fortunately, it's not
    // required that input iterators actually use references, and defining it
    // as value_type is valid.
    using reference = ScopedJavaLocalRef<T>;

    // This exists to make operator-> work as expected: its return value must
    // resolve to an actual pointer (otherwise the compiler just keeps calling
    // operator-> on the return value until it does), so we need an extra level
    // of indirection. This is sometimes called an "arrow proxy" or similar, and
    // this version is adapted from base/value_iterators.h.
    class pointer {
     public:
      explicit pointer(const ScopedJavaLocalRef<T>& ref) : ref_(ref) {}
      pointer(const pointer& ptr) = default;
      pointer& operator=(const pointer& ptr) = delete;
      ScopedJavaLocalRef<T>* operator->() { return &ref_; }

     private:
      ScopedJavaLocalRef<T> ref_;
    };

    iterator(const iterator&) = default;
    ~iterator() = default;

    iterator& operator=(const iterator&) = default;

    bool operator==(const iterator& other) const {
      JNI_ZERO_DCHECK(jarray_view_ == other.jarray_view_);
      return i_ == other.i_;
    }

    bool operator!=(const iterator& other) const {
      JNI_ZERO_DCHECK(jarray_view_ == other.jarray_view_);
      return i_ != other.i_;
    }

    ScopedJavaLocalRef<T> operator*() const {
      JNI_ZERO_DCHECK(i_ < jarray_view_->length_);
      return jarray_view_->Get(i_);
    }

    pointer operator->() const { return pointer(operator*()); }

    iterator& operator++() {
      JNI_ZERO_DCHECK(i_ < jarray_view_->length_);
      ++i_;
      return *this;
    }

    iterator operator++(int) {
      iterator old = *this;
      ++*this;
      return old;
    }

   private:
    iterator(const JArrayView<T>* jarray_view, int32_t i)
        : jarray_view_(jarray_view), i_(i) {}

    const JArrayView<T>* jarray_view_;
    int32_t i_;

    friend JArrayView;
  };

  JArrayView<T>(JNIEnv* env, JArray<T> array) : JArrayViewBase<T>(env, array) {}

  ~JArrayView<T>() = default;

  JArrayView(const JArrayView&) = delete;
  JArrayView(JArrayView&&) = delete;

  ScopedJavaLocalRef<T> Get(int32_t index) const {
    jobject obj = env_->GetObjectArrayElement(array_, index);
    return jni_zero::AdoptRef(env_, static_cast<T>(obj));
  }

  void CopyTo(std::vector<ScopedJavaLocalRef<T>>* buf) const {
    for (int32_t i = 0; i < length_; i++) {
      jobject obj = env_->GetObjectArrayElement(array_, i);
      buf->push_back(ScopedJavaLocalRef<T>::Adopt(env_, static_cast<T>(obj)));
    }
  }

  iterator begin() const { return iterator(this, 0); }

  iterator end() const { return iterator(this, length_); }

  friend iterator;
};

namespace internal {

template <typename T>
struct is_scoped_java_local_ref : std::false_type {};

template <typename T>
struct is_scoped_java_local_ref<ScopedJavaLocalRef<T>> : std::true_type {};

template <typename T>
concept IsScopedJavaLocalRef = is_scoped_java_local_ref<T>::value;

template <typename T>
struct _JniFuncMappings {
  // E.g. uint32_t -> int32_t
  using Canonical = typename _CanonicalJniPrimitiveType<T>::type;

  static JArray<T> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<T>>(
        _JniFuncMappings<Canonical>::NewArray(env, size));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<T> arr,
                             int32_t start,
                             int32_t len,
                             const T* buf) {
    _JniFuncMappings<Canonical>::SetArrayRegion(
        env, static_cast<JArray<Canonical>>(arr), start, len,
        reinterpret_cast<const Canonical*>(buf));
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<T> arr,
                             int32_t start,
                             int32_t len,
                             T* buf) {
    _JniFuncMappings<Canonical>::GetArrayRegion(
        env, static_cast<JArray<Canonical>>(arr), start, len,
        reinterpret_cast<Canonical*>(buf));
  }
};

template <>
struct _JniFuncMappings<bool> {
  static JArray<bool> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<bool>>(
        env->NewBooleanArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<bool> arr,
                             int32_t start,
                             int32_t len,
                             const bool* buf) {
    env->SetBooleanArrayRegion(arr, start, len,
                               reinterpret_cast<const uint8_t*>(buf));
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<bool> arr,
                             int32_t start,
                             int32_t len,
                             bool* buf) {
    env->GetBooleanArrayRegion(arr, start, len,
                               reinterpret_cast<uint8_t*>(buf));
  }
};

template <>
struct _JniFuncMappings<int8_t> {
  static JArray<int8_t> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<int8_t>>(
        env->NewByteArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<int8_t> arr,
                             int32_t start,
                             int32_t len,
                             const int8_t* buf) {
    env->SetByteArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<int8_t> arr,
                             int32_t start,
                             int32_t len,
                             int8_t* buf) {
    env->GetByteArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<uint16_t> {
  static JArray<uint16_t> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<uint16_t>>(
        env->NewCharArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<uint16_t> arr,
                             int32_t start,
                             int32_t len,
                             const uint16_t* buf) {
    env->SetCharArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<uint16_t> arr,
                             int32_t start,
                             int32_t len,
                             uint16_t* buf) {
    env->GetCharArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<int16_t> {
  static JArray<int16_t> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<int16_t>>(
        env->NewShortArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<int16_t> arr,
                             int32_t start,
                             int32_t len,
                             const int16_t* buf) {
    env->SetShortArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<int16_t> arr,
                             int32_t start,
                             int32_t len,
                             int16_t* buf) {
    env->GetShortArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<int32_t> {
  static JArray<int32_t> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<int32_t>>(
        env->NewIntArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<int32_t> arr,
                             int32_t start,
                             int32_t len,
                             const int32_t* buf) {
    env->SetIntArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<int32_t> arr,
                             int32_t start,
                             int32_t len,
                             int32_t* buf) {
    env->GetIntArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<int64_t> {
  static JArray<int64_t> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<int64_t>>(
        env->NewLongArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<int64_t> arr,
                             int32_t start,
                             int32_t len,
                             const int64_t* buf) {
    env->SetLongArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<int64_t> arr,
                             int32_t start,
                             int32_t len,
                             int64_t* buf) {
    env->GetLongArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<float> {
  static JArray<float> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<float>>(
        env->NewFloatArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<float> arr,
                             int32_t start,
                             int32_t len,
                             const float* buf) {
    env->SetFloatArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<float> arr,
                             int32_t start,
                             int32_t len,
                             float* buf) {
    env->GetFloatArrayRegion(arr, start, len, buf);
  }
};

template <>
struct _JniFuncMappings<double> {
  static JArray<double> NewArray(JNIEnv* env, size_t size) {
    return static_cast<JArray<double>>(
        env->NewDoubleArray(static_cast<jsize>(size)));
  }
  static void SetArrayRegion(JNIEnv* env,
                             JArray<double> arr,
                             int32_t start,
                             int32_t len,
                             const double* buf) {
    env->SetDoubleArrayRegion(arr, start, len, buf);
  }
  static void GetArrayRegion(JNIEnv* env,
                             JArray<double> arr,
                             int32_t start,
                             int32_t len,
                             double* buf) {
    env->GetDoubleArrayRegion(arr, start, len, buf);
  }
};

}  // namespace internal

// Wrapper for a primitive jarray (e.g. jbooleanArray, jintArray, etc.).
// Note: This uses GetPrimitiveArrayCritical / ReleasePrimitiveArrayCritical,
// so no JNI calls or slow operations should be made while this view is in
// scope.
template <typename T>
  requires internal::IsPrimitiveType<T>
class JArrayViewCritical<T> : public JArrayViewBase<T> {
  using JArrayViewBase<T>::env_;
  using JArrayViewBase<T>::array_;
  using JArrayViewBase<T>::length_;

 public:
  JArrayViewCritical<T>(JNIEnv* env, JArray<T> array)
      : JArrayViewBase<T>(env, array),
        data_(static_cast<T*>(env->GetPrimitiveArrayCritical(array, nullptr))) {
  }

  ~JArrayViewCritical<T>() {
    if (data_) {
      env_->ReleasePrimitiveArrayCritical(array_, data_, JNI_ABORT);
    }
  }

  JArrayViewCritical(const JArrayViewCritical&) = delete;
  JArrayViewCritical& operator=(const JArrayViewCritical&) = delete;
  JArrayViewCritical& operator=(JArrayViewCritical&&) = delete;
  JArrayViewCritical(JArrayViewCritical&& other) noexcept
      : JArrayViewBase<T>(other.env_, other.array_, other.length_),
        data_(std::exchange(other.data_, nullptr)) {
    other.array_ = nullptr;
  }

  T Get(int32_t index) const {
    JNI_ZERO_CHECK(index >= 0 && index < length_);
#pragma clang unsafe_buffer_usage begin
    return data_[index];
#pragma clang unsafe_buffer_usage end
  }

  T operator[](size_t index) const { return Get(static_cast<int32_t>(index)); }

  const T* data() const noexcept JNI_LIFETIME_BOUND { return data_; }

  const T* begin() const noexcept JNI_LIFETIME_BOUND { return data_; }

  const T* end() const noexcept JNI_LIFETIME_BOUND {
#pragma clang unsafe_buffer_usage begin
    return data_ + length_;
#pragma clang unsafe_buffer_usage end
  }

  std::span<const T> as_span() const noexcept JNI_LIFETIME_BOUND {
    return std::span<const T>(data_, static_cast<size_t>(length_));
  }

  std::string_view as_string_view() const noexcept JNI_LIFETIME_BOUND
    requires(sizeof(T) == 1)
  {
    return std::string_view(reinterpret_cast<const char*>(data_),
                            static_cast<size_t>(length_));
  }

 private:
  T* data_;
};

template <typename T>
  requires internal::IsPrimitiveType<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env, size_t size) {
  JArray<T> ret = static_cast<JArray<T>>(
      internal::_JniFuncMappings<T>::NewArray(env, size));
  CheckException(env);
  return jni_zero::AdoptRef(env, ret);
}

template <typename T>
  requires internal::IsJobject<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              size_t length,
                                              jclass cls) {
  JArray<T> ret = static_cast<JArray<T>>(
      env->NewObjectArray(static_cast<jsize>(length), cls, nullptr));
  CheckException(env);
  return jni_zero::AdoptRef(env, ret);
}

template <typename T>
  requires internal::IsJobject<T>
static ScopedJavaLocalRef<JArray<T>>
NewArray(JNIEnv* env, std::span<const ScopedJavaLocalRef<T>> buf, jclass cls) {
  int32_t length = static_cast<int32_t>(buf.size());
  JArray<T> ret =
      static_cast<JArray<T>>(env->NewObjectArray(length, cls, nullptr));
  CheckException(env);
  for (size_t i = 0; i < buf.size(); i++) {
    env->SetObjectArrayElement(ret, static_cast<int32_t>(i), buf[i].obj());
  }
  return jni_zero::AdoptRef(env, ret);
}

template <typename T, typename U>
  requires(!internal::IsScopedJavaLocalRef<U>)
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              std::span<const U> buf,
                                              jclass cls) {
  int32_t length = static_cast<int32_t>(buf.size());
  JArray<T> ret =
      static_cast<JArray<T>>(env->NewObjectArray(length, cls, nullptr));
  CheckException(env);
  for (size_t i = 0; i < buf.size(); i++) {
    env->SetObjectArrayElement(ret, static_cast<int32_t>(i),
                               ToJniType(env, buf[i]).obj());
  }
  return jni_zero::AdoptRef(env, ret);
}

template <typename T>
  requires internal::IsPrimitiveType<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              std::span<const T> buf) {
  size_t size = buf.size();
  JArray<T> ret = static_cast<JArray<T>>(
      internal::_JniFuncMappings<T>::NewArray(env, size));
  CheckException(env);
  internal::_JniFuncMappings<T>::SetArrayRegion(
      env, ret, 0, static_cast<int32_t>(size), buf.data());
  return jni_zero::AdoptRef(env, ret);
}

extern jclass g_object_class;

template <typename T>
static ScopedJavaLocalRef<JArray<jobject>> NewObjectArray(
    JNIEnv* env,
    std::span<const T> buf) {
  return NewArray<jobject>(env, buf, g_object_class);
}

inline ScopedJavaLocalRef<JArray<jobject>> NewObjectArray(JNIEnv* env,
                                                          size_t length) {
  return NewArray<jobject>(env, length, g_object_class);
}

extern jclass g_string_class;

template <typename T>
static ScopedJavaLocalRef<JArray<jstring>> NewStringArray(
    JNIEnv* env,
    std::span<const T> buf) {
  return NewArray<jstring>(env, buf, g_string_class);
}

inline ScopedJavaLocalRef<JArray<jstring>> NewStringArray(JNIEnv* env,
                                                          size_t length) {
  return NewArray<jstring>(env, length, g_string_class);
}

// Below are overloads that take a std::vector instead of std::span.
// These overloads are needed because template deduction happens before implicit
// conversions, so the compiler cannot match std::vector<T> to std::span<T>.
template <typename T>
  requires internal::IsJobject<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(
    JNIEnv* env,
    const std::vector<ScopedJavaLocalRef<T>>& buf,
    jclass cls) {
  return NewArray<T>(env, std::span<const ScopedJavaLocalRef<T>>(buf), cls);
}

template <typename T, typename U>
  requires(!internal::IsScopedJavaLocalRef<U>)
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              const std::vector<U>& buf,
                                              jclass cls) {
  return NewArray<T, U>(env, std::span<const U>(buf), cls);
}

template <typename T>
  requires internal::IsPrimitiveType<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              const std::vector<T>& buf) {
  if constexpr (std::is_same_v<T, bool>) {
    // A std::vector<bool> stores each bool using 1 bit rather than 1 byte,
    // and thus we are not allowed to directly convert std::vector<bool> to
    // std::span<bool>.
    if (buf.size() <= 1024) {
      std::array<bool, 1024> bool_arr;
      for (size_t i = 0; i < buf.size(); ++i) {
        bool_arr[i] = buf[i];
      }
      return NewArray<bool>(env,
                            std::span<const bool>(bool_arr.data(), buf.size()));
    }
    std::vector<uint8_t> bool_vec;
    bool_vec.reserve(buf.size());
    for (bool b : buf) {
      bool_vec.push_back(static_cast<uint8_t>(b));
    }
    return NewArray<bool>(
        env,
        std::span<const bool>(reinterpret_cast<const bool*>(bool_vec.data()),
                              bool_vec.size()));
  } else {
    return NewArray<T>(env, std::span<const T>(buf));
  }
}

template <typename T>
static ScopedJavaLocalRef<JArray<jobject>> NewObjectArray(
    JNIEnv* env,
    const std::vector<T>& buf) {
  return NewObjectArray<T>(env, std::span<const T>(buf));
}

template <typename T>
static ScopedJavaLocalRef<JArray<jstring>> NewStringArray(
    JNIEnv* env,
    const std::vector<T>& buf) {
  return NewStringArray<T>(env, std::span<const T>(buf));
}

// Below are overloads that take std::initializer_list instead of std::span.
// These overloads are needed because template deduction happens before implicit
// conversions, so the compiler cannot match std::initializer_list<T> to
// std::span<T>.
template <typename T>
  requires internal::IsJobject<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(
    JNIEnv* env,
    std::initializer_list<ScopedJavaLocalRef<T>> buf,
    jclass cls) {
  return NewArray<T>(env, std::span<const ScopedJavaLocalRef<T>>(buf), cls);
}

template <typename T, typename U>
  requires(!internal::IsScopedJavaLocalRef<U>)
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              std::initializer_list<U> buf,
                                              jclass cls) {
  return NewArray<T, U>(env, std::span<const U>(buf), cls);
}

template <typename T>
  requires internal::IsPrimitiveType<T>
static ScopedJavaLocalRef<JArray<T>> NewArray(JNIEnv* env,
                                              std::initializer_list<T> buf) {
  return NewArray<T>(env, std::span<const T>(buf));
}

template <typename T>
static ScopedJavaLocalRef<JArray<jobject>> NewObjectArray(
    JNIEnv* env,
    std::initializer_list<T> buf) {
  return NewObjectArray<T>(env, std::span<const T>(buf));
}

template <typename T>
static ScopedJavaLocalRef<JArray<jstring>> NewStringArray(
    JNIEnv* env,
    std::initializer_list<T> buf) {
  return NewStringArray<T>(env, std::span<const T>(buf));
}

}  // namespace jni_zero

#undef JNI_LIFETIME_BOUND

#endif  // JNI_ZERO_JNI_WRAPPERS_H_