// 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_JAVA_REFS_H_
#define JNI_ZERO_JAVA_REFS_H_

#include <jni.h>

#include <concepts>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>

#include "third_party/jni_zero/compiler_specific.h"
#include "third_party/jni_zero/jni_export.h"
#include "third_party/jni_zero/logging.h"

// Forward declaration of template class that contains @CalledByNative methods.
// Must live in a custom / unique namespace to ensure it doesn't collide with
// namespaces used from @JniType strings.
namespace jni_zero_internal {
template <typename T>
class _CalledByNatives;
template <typename T>
class _CalledByNativesStatics;
}  // namespace jni_zero_internal

namespace jni_zero::internal {
template <typename T>
struct _JniFuncMappings;

template <typename T>
concept IsJobject =
    std::derived_from<std::remove_pointer_t<T>, std::remove_pointer_t<jobject>>;

// The "jobject" type is basically the "any" type, so we use this concept to
// allow implicit casts from jobject -> subclass.
template <typename T, typename U>
concept IsConvertibleJObject =
    std::is_convertible_v<U, T> || std::same_as<U, jobject>;

template <typename T>
concept IsPrimitiveType = std::is_arithmetic_v<T>;

template <typename T>
struct _CanonicalJniPrimitiveType {
  using type = T;
};

template <>
struct _CanonicalJniPrimitiveType<uint8_t> {
  using type = int8_t;
};

template <>
struct _CanonicalJniPrimitiveType<char> {
  using type = int8_t;
};

template <>
struct _CanonicalJniPrimitiveType<char16_t> {
  using type = uint16_t;
};

template <>
struct _CanonicalJniPrimitiveType<uint32_t> {
  using type = int32_t;
};

template <>
struct _CanonicalJniPrimitiveType<uint64_t> {
  using type = int64_t;
};

template <typename T>
struct _JArrayElementType {
  using type = void;
};

template <>
struct _JArrayElementType<jbooleanArray> {
  using type = bool;
};

template <>
struct _JArrayElementType<jbyteArray> {
  using type = int8_t;
};

template <>
struct _JArrayElementType<jcharArray> {
  using type = uint16_t;
};

template <>
struct _JArrayElementType<jshortArray> {
  using type = int16_t;
};

template <>
struct _JArrayElementType<jintArray> {
  using type = int32_t;
};

template <>
struct _JArrayElementType<jlongArray> {
  using type = int64_t;
};

template <>
struct _JArrayElementType<jfloatArray> {
  using type = float;
};

template <>
struct _JArrayElementType<jdoubleArray> {
  using type = double;
};

template <>
struct _JArrayElementType<jobjectArray> {
  using type = jobject;
};

template <typename T>
  requires internal::IsJobject<T>
class _JObjectArray : public _jobjectArray {};

template <typename T>
struct _JArrayHelper;

template <typename T>
  requires internal::IsJobject<T>
struct _JArrayHelper<T> {
  using type = _JObjectArray<T>;
};

template <>
struct _JArrayHelper<jobject> {
  using type = _jobjectArray;
};

template <>
struct _JArrayHelper<bool> {
  using type = _jbooleanArray;
};

template <>
struct _JArrayHelper<int8_t> {
  using type = _jbyteArray;
};

template <>
struct _JArrayHelper<uint16_t> {
  using type = _jcharArray;
};

template <>
struct _JArrayHelper<int16_t> {
  using type = _jshortArray;
};

template <>
struct _JArrayHelper<int32_t> {
  using type = _jintArray;
};

template <>
struct _JArrayHelper<int64_t> {
  using type = _jlongArray;
};

template <>
struct _JArrayHelper<float> {
  using type = _jfloatArray;
};

template <>
struct _JArrayHelper<double> {
  using type = _jdoubleArray;
};

// Subclasses the canonical JNI array struct (e.g. _jbyteArray for uint8_t)
// so that non-canonical primitive types (like uint8_t, uint32_t, char) retain
// their distinct C++ element type rather than collapsing into type aliases,
// while remaining implicitly convertible to their underlying JNI array type.
template <typename T>
  requires internal::IsPrimitiveType<T>
class _JPrimitiveArray
    : public _JArrayHelper<typename _CanonicalJniPrimitiveType<T>::type>::type {
};

// Primary template: non-canonical primitive types use _JPrimitiveArray to
// preserve their element type.
template <typename T>
struct _JArrayHelper {
  using type = _JPrimitiveArray<T>;
};

// Maps _JPrimitiveArray<T>* to its non-canonical element type T, allowing
// CreateViewCritical() to deduce and return JArrayViewCritical<T> directly.
template <typename T>
  requires internal::IsPrimitiveType<T>
struct _JArrayElementType<_JPrimitiveArray<T>*> {
  using type = T;
};

template <typename T>
using _JArray = typename _JArrayHelper<T>::type;

template <typename T>
using JArray = _JArray<T>*;
}  // namespace jni_zero::internal

using ::jni_zero::internal::JArray;
#define _JNI_ZERO_JArray_DEFINED

namespace jni_zero {

// Creates a new local reference frame, in which at least a given number of
// local references can be created. Note that local references already created
// in previous local frames are still valid in the current local frame.
class JNI_ZERO_COMPONENT_BUILD_EXPORT ScopedJavaLocalFrame {
 public:
  explicit ScopedJavaLocalFrame(JNIEnv* env);
  ScopedJavaLocalFrame(JNIEnv* env, int capacity);

  ScopedJavaLocalFrame(const ScopedJavaLocalFrame&) = delete;
  ScopedJavaLocalFrame& operator=(const ScopedJavaLocalFrame&) = delete;

  ~ScopedJavaLocalFrame();

 private:
  // This class is only good for use on the thread it was created on so
  // it's safe to cache the non-threadsafe JNIEnv* inside this object.
  JNIEnv* env_;
};

// Forward declare the generic java reference template class.
template <typename T = jobject>
  requires internal::IsJobject<T>
class JavaRef;

template <typename T>
concept IsJavaRef =
    std::is_base_of_v<jni_zero::JavaRef<jobject>, std::remove_cvref_t<T>>;

// Forward declaration of the JArrayView and JArrayViewCritical classes.
template <typename T>
class JArrayView;

template <typename T>
class JArrayViewCritical;

namespace internal {

// Concept to check if the _CalledByNatives<T> specialization is defined.
template <typename T>
concept HasCalledByNatives =
    requires { sizeof(jni_zero_internal::_CalledByNatives<T>); };
}  // namespace internal

// Template specialization of JavaRef, which acts as the base class for all
// other JavaRef<> template types. This allows you to e.g. pass
// ScopedJavaLocalRef<jstring> into a function taking const JavaRef<jobject>&
template <>
class JNI_ZERO_COMPONENT_BUILD_EXPORT JNI_ZERO_TRIVIAL_ABI JavaRef<jobject> {
 public:
  // Initializes a null reference.
  constexpr JavaRef() {}

  // Allow nullptr to be converted to JavaRef. This avoids having to declare an
  // empty JavaRef just to pass null to a function, and makes C++ "nullptr" and
  // Java "null" equivalent.
  constexpr JavaRef(std::nullptr_t) {}

  JavaRef(const JavaRef&) = delete;
  JavaRef& operator=(const JavaRef&) = delete;

  // Public to allow destruction of null JavaRef objects.
  ~JavaRef() {}

  // TODO(torne): maybe rename this to get() for consistency with unique_ptr
  // once there's fewer unnecessary uses of it in the codebase.
  jobject obj() const { return obj_; }

  explicit operator bool() const { return obj_ != nullptr; }

  // Deprecated. Just use bool conversion.
  // TODO(torne): replace usage and remove this.
  bool is_null() const { return obj_ == nullptr; }

  // Create a JavaRef that is not automatically released. Used for JNI
  // parameters (which should not be released).
  static JavaRef<jobject> CreateLeaky(JNIEnv* env, jobject obj) {
    return JavaRef<jobject>(env, obj);
  }

  // Allows calling methods as another type, or to pass to a function as
  // another type, but does not allow assignment to a ScopedJavaLocalRef.
  // There is an rvalue version of this in ScopedJavaLocalRef that enables
  // assignment.
  template <typename To>
    requires internal::IsJobject<To>
  const JavaRef<To>& As() const {
    return *reinterpret_cast<const JavaRef<To>*>(this);
  }

  // Convert this JavaRef<jobject> to the corresponding C++ type by
  // calling FromJniType.
  template <typename To>
  To ConvertTo(JNIEnv* env) const {
    return FromJniType<To>(env, *this);
  }

 protected:
  JavaRef(JavaRef&&) = default;
  JavaRef& operator=(JavaRef&&) = default;

#if JNI_ZERO_DCHECK_IS_ON()
  // Implementation contains a DCHECK; implement out-of-line when DCHECK_IS_ON.
  JavaRef(JNIEnv* env, jobject obj);
#else
  JavaRef(JNIEnv* env, jobject obj) : obj_(obj) {}
#endif

  // Used for move semantics. obj_ must have been released first if non-null.
  void steal(JavaRef&& other) {
    obj_ = other.obj_;
    other.obj_ = nullptr;
  }

  // The following are implementation detail convenience methods, for
  // use by the sub-classes.
  JNIEnv* SetNewLocalRef(JNIEnv* env, jobject obj);
  void SetNewGlobalRef(JNIEnv* env, jobject obj);
  void SetNewGlobalRefAndLeak(JNIEnv* env, jobject obj);
  void ResetLocalRef(JNIEnv* env);
  void ResetGlobalRef();

  jobject ReleaseInternal() {
    jobject obj = obj_;
    obj_ = nullptr;
    return obj;
  }

 private:
  jobject obj_ = nullptr;
};

template <typename T = jobject>
class ScopedJavaLocalRef;

// Generic base class for ScopedJavaLocalRef and ScopedJavaGlobalRef. Useful
// for allowing functions to accept a reference without having to mandate
// whether it is a local or global type.
template <typename T>
  requires internal::IsJobject<T>
class JNI_ZERO_TRIVIAL_ABI JavaRef : public JavaRef<jobject> {
 public:
  constexpr JavaRef() {}
  constexpr JavaRef(std::nullptr_t) {}

 protected:
  JavaRef(JavaRef&&) = default;
  JavaRef& operator=(JavaRef&&) = default;

  JavaRef(JNIEnv* env, jobject obj) : JavaRef<jobject>(env, obj) {}

 public:
  JavaRef(const JavaRef&) = delete;
  JavaRef& operator=(const JavaRef&) = delete;

  ~JavaRef() {}

  T obj() const { return static_cast<T>(JavaRef<jobject>::obj()); }

  // Define this only when the _jni.h header has been #included.
  const jni_zero_internal::_CalledByNatives<T>* operator->() const
    requires internal::HasCalledByNatives<T>
  {
    // CalledByNatives does the reverse reinterpret_cast<>.
    // This approach optimizes better than passing |this| as a parameter.
    return reinterpret_cast<const jni_zero_internal::_CalledByNatives<T>*>(
        this);
  }

  // Create a JavaRef that is not automatically released. Used for JNI
  // parameters (which should not be released).
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  static JavaRef<T> CreateLeaky(JNIEnv* env, U obj) {
    return JavaRef<T>(env, obj);
  }

  int32_t GetLength(JNIEnv* env) const
    requires std::is_convertible_v<T, jarray>
  {
    return env->GetArrayLength(this->obj());
  }

  size_t GetSize(JNIEnv* env) const
    requires std::is_convertible_v<T, jarray>
  {
    return static_cast<size_t>(GetLength(env));
  }

  ScopedJavaLocalRef<jobject> Get(JNIEnv* env, size_t index) const
    requires std::is_same_v<T, jobjectArray>;

  template <typename U>
  auto GetAs(JNIEnv* env, size_t index) const
    requires std::is_same_v<T, jobjectArray>
  {
    // Make GetAs<JFoo>() an alias of GetAs<ScopedJavaLocalRef<JFoo>>().
    if constexpr (internal::IsJobject<U>) {
      return Get(env, index).template As<U>();
    } else {
      return Get(env, index).template ConvertTo<U>(env);
    }
  }

  void Set(JNIEnv* env, size_t index, const JavaRef<jobject>& value) const
    requires std::is_same_v<T, jobjectArray>
  {
    env->SetObjectArrayElement(this->obj(), static_cast<int32_t>(index),
                               value.obj());
  }

  template <typename U>
  void CopyTo(JNIEnv* env, std::vector<ScopedJavaLocalRef<U>>* buf) const
    requires std::is_convertible_v<T, jobjectArray>;

  template <typename DestType>
    requires(internal::IsPrimitiveType<
                 typename internal::_JArrayElementType<T>::type> &&
             sizeof(DestType) ==
                 sizeof(typename internal::_JArrayElementType<T>::type))
  void CopyTo(JNIEnv* env, DestType* dest, size_t size) const {
    if (size == 0) {
      return;
    }
    using ElementType = typename internal::_JArrayElementType<T>::type;
    internal::_JniFuncMappings<ElementType>::GetArrayRegion(
        env, static_cast<JArray<ElementType>>(this->obj()), 0,
        static_cast<int32_t>(size), reinterpret_cast<ElementType*>(dest));
  }

  // The [[clang::lifetimebound]] is required because the lifetime of the
  // JArrayView cannot safely outlast the lifetime of |this|.
  auto CreateView(JNIEnv* env) const [[clang::lifetimebound]]
    requires std::is_same_v<T, jobjectArray>
  {
    return JArrayView<jobject>(env, this->obj());
  }

  template <
      typename ElementType = typename internal::_JArrayElementType<T>::type>
    requires(internal::IsPrimitiveType<ElementType> &&
             sizeof(ElementType) ==
                 sizeof(typename internal::_JArrayElementType<T>::type))
  auto CreateViewCritical(JNIEnv* env) const [[clang::lifetimebound]] {
    return JArrayViewCritical<ElementType>(
        env, reinterpret_cast<JArray<ElementType>>(this->obj()));
  }
};

// JavaRef specialization for JArray<T> where T is a jobject subclass.
template <typename T>
  requires internal::IsJobject<T>
class JNI_ZERO_TRIVIAL_ABI
    JavaRef<internal::_JObjectArray<T>*> : public JavaRef<jobjectArray> {
 public:
  constexpr JavaRef() = default;
  explicit constexpr JavaRef(std::nullptr_t) {}

 protected:
  JavaRef(JavaRef&&) = default;
  JavaRef& operator=(JavaRef&&) = default;

  JavaRef(JNIEnv* env, jobject obj) : JavaRef<jobjectArray>(env, obj) {}

 public:
  JavaRef(const JavaRef&) = delete;
  JavaRef& operator=(const JavaRef&) = delete;

  JArray<T> obj() const {
    return static_cast<JArray<T>>(JavaRef<jobject>::obj());
  }

  template <typename U>
    requires internal::IsConvertibleJObject<JArray<T>, U>
  static JavaRef<JArray<T>> CreateLeaky(JNIEnv* env, U obj) {
    return JavaRef<internal::_JObjectArray<T>*>(env, obj);
  }

  ScopedJavaLocalRef<T> Get(JNIEnv* env, size_t index) const;

  template <typename U>
  auto GetAs(JNIEnv* env, size_t index) const {
    // Make GetAs<JFoo>() an alias of GetAs<ScopedJavaLocalRef<JFoo>>().
    if constexpr (internal::IsJobject<U>) {
      return Get(env, index).template As<U>();
    } else {
      return Get(env, index).template ConvertTo<U>(env);
    }
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Set(JNIEnv* env, size_t index, const JavaRef<U>& value) const {
    env->SetObjectArrayElement(this->obj(), static_cast<int32_t>(index),
                               value.obj());
  }

  JArrayView<T> CreateView(JNIEnv* env) const [[clang::lifetimebound]] {
    return JArrayView<T>(env, obj());
  }
};

template <typename T>
JavaRef<T> CreateLeaky(JNIEnv* env, T obj) {
  return JavaRef<T>::CreateLeaky(env, obj);
}

// Holds a local reference to a Java object. The local reference is scoped
// to the lifetime of this object.
// Instances of this class may hold onto any JNIEnv passed into it until
// destroyed. Therefore, since a JNIEnv is only suitable for use on a single
// thread, objects of this class must be created, used, and destroyed, on a
// single thread.
// Therefore, this class should only be used as a stack-based object and from a
// single thread. If you wish to have the reference outlive the current
// callstack (e.g. as a class member) or you wish to pass it across threads,
// use a ScopedJavaGlobalRef instead.
template <typename T>
class JNI_ZERO_TRIVIAL_ABI ScopedJavaLocalRef : public JavaRef<T> {
 public:
  // Take ownership of a bare jobject. This does not create a new reference.
  // This should only be used by JNI helper functions, or in cases where code
  // must call JNIEnv methods directly.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  static ScopedJavaLocalRef Adopt(JNIEnv* env, U obj) {
    return ScopedJavaLocalRef(env, obj);
  }

  constexpr ScopedJavaLocalRef() {}
  constexpr ScopedJavaLocalRef(std::nullptr_t) {}

  // Copy constructor. This is required in addition to the copy conversion
  // constructor below.
  ScopedJavaLocalRef(const ScopedJavaLocalRef& other) : env_(other.env_) {
    JavaRef<jobject>::SetNewLocalRef(env_, other.obj());
  }

  // Copy conversion constructor.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef(const ScopedJavaLocalRef<U>& other) : env_(other.env_) {
    JavaRef<jobject>::SetNewLocalRef(env_, other.obj());
  }

  // Move constructor. This is required in addition to the move conversion
  // constructor below.
  ScopedJavaLocalRef(ScopedJavaLocalRef&& other) : env_(other.env_) {
    JavaRef<jobject>::steal(std::move(other));
  }

  // Move conversion constructor.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef(ScopedJavaLocalRef<U>&& other) : env_(other.env_) {
    JavaRef<jobject>::steal(std::move(other));
  }

  // Constructor for other JavaRef types.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  explicit ScopedJavaLocalRef(const JavaRef<U>& other) {
    Reset(other);
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef(JNIEnv* env, const JavaRef<U>& other) {
    Reset(other);
  }

  ~ScopedJavaLocalRef() { Reset(); }

  // Null assignment, for disambiguation.
  ScopedJavaLocalRef& operator=(std::nullptr_t) {
    Reset();
    return *this;
  }

  // Copy assignment.
  ScopedJavaLocalRef& operator=(const ScopedJavaLocalRef& other) {
    Reset(other);
    return *this;
  }

  // Copy conversion assignment.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef& operator=(const ScopedJavaLocalRef<U>& other) {
    Reset(other);
    return *this;
  }

  // Move assignment.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef& operator=(ScopedJavaLocalRef<U>&& other) {
    env_ = other.env_;
    Reset();
    JavaRef<T>::steal(std::move(other));
    return *this;
  }

  // Assignment for other JavaRef types.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaLocalRef& operator=(const JavaRef<U>& other) {
    Reset(other);
    return *this;
  }

  void Reset() { JavaRef<jobject>::ResetLocalRef(env_); }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(const ScopedJavaLocalRef<U>& other) {
    // We can copy over env_ here as |other| instance must be from the same
    // thread as |this| local ref. (See class comment for multi-threading
    // limitations, and alternatives).
    env_ = JavaRef<jobject>::SetNewLocalRef(other.env_, other.obj());
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(const JavaRef<U>& other) {
    // If |env_| was not yet set (is still null) it will be attached to the
    // current thread in SetNewLocalRef().
    env_ = JavaRef<jobject>::SetNewLocalRef(env_, other.obj());
  }

  // Releases the local reference to the caller. The caller *must* delete the
  // local reference when it is done with it. Note that calling a Java method
  // is *not* a transfer of ownership and Release() should not be used.
  T Release() { return static_cast<T>(JavaRef<T>::ReleaseInternal()); }

  // Alias for Release(). For use in templates when global refs are invalid.
  T ReleaseLocal() { return static_cast<T>(JavaRef<T>::ReleaseInternal()); }

  using JavaRef<T>::As;

  // Enables casting while assigning. E.g.:
  // ScopedJavaLocalRef<JFoo> foo = FuncThatReturnsJobject(env).As<JFoo>();
  template <typename To>
    requires internal::IsJobject<To>
  ScopedJavaLocalRef<To>&& As() && {
    return std::move(*reinterpret_cast<ScopedJavaLocalRef<To>*>(this));
  }

 private:
  ScopedJavaLocalRef(JNIEnv* env, jobject obj)
      : JavaRef<T>(env, obj), env_(env) {}

  // This class is only good for use on the thread it was created on so
  // it's safe to cache the non-threadsafe JNIEnv* inside this object.
  JNIEnv* env_ = nullptr;

  // Friend required to get env_ from conversions.
  template <typename U>
  friend class ScopedJavaLocalRef;
};

template <typename T>
ScopedJavaLocalRef<T> AdoptRef(JNIEnv* env, T obj) {
  return ScopedJavaLocalRef<T>::Adopt(env, obj);
}

// Holds a global reference to a Java object. The global reference is scoped
// to the lifetime of this object. This class does not hold onto any JNIEnv*
// passed to it, hence it is safe to use across threads (within the constraints
// imposed by the underlying Java object that it references).
template <typename T = jobject>
class JNI_ZERO_TRIVIAL_ABI ScopedJavaGlobalRef : public JavaRef<T> {
 public:
  constexpr ScopedJavaGlobalRef() {}
  constexpr ScopedJavaGlobalRef(std::nullptr_t) {}

  // Copy constructor. This is required in addition to the copy conversion
  // constructor below.
  ScopedJavaGlobalRef(const ScopedJavaGlobalRef& other) { Reset(other); }

  // Copy conversion constructor.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef(const ScopedJavaGlobalRef<U>& other) {
    Reset(other);
  }

  // Move constructor. This is required in addition to the move conversion
  // constructor below.
  ScopedJavaGlobalRef(ScopedJavaGlobalRef&& other) {
    JavaRef<T>::steal(std::move(other));
  }

  // Move conversion constructor.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef(ScopedJavaGlobalRef<U>&& other) {
    JavaRef<T>::steal(std::move(other));
  }

  // Conversion constructor for other JavaRef types.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  explicit ScopedJavaGlobalRef(const JavaRef<U>& other) {
    Reset(other);
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef(JNIEnv* env, const JavaRef<U>& other) {
    JavaRef<T>::SetNewGlobalRef(env, other.obj());
  }

  ~ScopedJavaGlobalRef() { Reset(); }

  // Null assignment, for disambiguation.
  ScopedJavaGlobalRef& operator=(std::nullptr_t) {
    Reset();
    return *this;
  }

  // Copy assignment.
  ScopedJavaGlobalRef& operator=(const ScopedJavaGlobalRef& other) {
    Reset(other);
    return *this;
  }

  // Copy conversion assignment.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef& operator=(const ScopedJavaGlobalRef<U>& other) {
    Reset(other);
    return *this;
  }

  // Move assignment.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef& operator=(ScopedJavaGlobalRef<U>&& other) {
    Reset();
    JavaRef<T>::steal(std::move(other));
    return *this;
  }

  // Assignment for other JavaRef types.
  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  ScopedJavaGlobalRef& operator=(const JavaRef<U>& other) {
    Reset(other);
    return *this;
  }

  void Reset() { JavaRef<T>::ResetGlobalRef(); }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(const ScopedJavaGlobalRef<U>& other) {
    Reset(nullptr, other);
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(const JavaRef<U>& other) {
    Reset(nullptr, other);
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(JNIEnv* env, const JavaRef<U>& other) {
    JavaRef<T>::SetNewGlobalRef(env, other.obj());
  }

  // Releases the global reference to the caller. The caller *must* delete the
  // global reference when it is done with it. Note that calling a Java method
  // is *not* a transfer of ownership and Release() should not be used.
  T Release() { return static_cast<T>(JavaRef<T>::ReleaseInternal()); }

  // Create a local reference.
  ScopedJavaLocalRef<T> AsLocalRef(JNIEnv* env) const {
    T j_obj = JavaRef<T>::obj();
    if (!j_obj) {
      return nullptr;
    }
    return jni_zero::AdoptRef(env, static_cast<T>(env->NewLocalRef(j_obj)));
  }

  using JavaRef<T>::As;

  // Enables casting while assigning. E.g.:
  // ScopedJavaGlobalRef<JFoo> foo = FuncThatReturnsJobject(env).As<JFoo>();
  template <typename To>
    requires internal::IsJobject<To>
  ScopedJavaGlobalRef<To>&& As() && {
    return std::move(*reinterpret_cast<ScopedJavaGlobalRef<To>*>(this));
  }
};

// Wrapper for working with weak references.
class JNI_ZERO_COMPONENT_BUILD_EXPORT ScopedJavaGlobalWeakRef {
 public:
  ScopedJavaGlobalWeakRef() = default;
  ScopedJavaGlobalWeakRef(const ScopedJavaGlobalWeakRef& orig);
  ScopedJavaGlobalWeakRef(ScopedJavaGlobalWeakRef&& orig) : obj_(orig.obj_) {
    orig.obj_ = nullptr;
  }
  ScopedJavaGlobalWeakRef(JNIEnv* env, const JavaRef<jobject>& obj);
  ~ScopedJavaGlobalWeakRef() { reset(); }

  void operator=(const ScopedJavaGlobalWeakRef& rhs) { Assign(rhs); }
  void operator=(ScopedJavaGlobalWeakRef&& rhs) { std::swap(obj_, rhs.obj_); }

  ScopedJavaLocalRef<jobject> get(JNIEnv* env) const;

  // Returns true if the weak reference has not been initialized to point at
  // an object (or ḣas had reset() called).
  // Do not call this to test if the object referred to still exists! The weak
  // reference remains initialized even if the target object has been collected.
  bool is_uninitialized() const { return obj_ == nullptr; }

  void reset();

 private:
  void Assign(const ScopedJavaGlobalWeakRef& rhs);

  jweak obj_ = nullptr;
};

// A global JavaRef that will never be released.
template <typename T = jobject>
class JNI_ZERO_COMPONENT_BUILD_EXPORT JNI_ZERO_TRIVIAL_ABI LeakedJavaGlobalRef
    : public JavaRef<T> {
 public:
  constexpr LeakedJavaGlobalRef() = default;
  constexpr LeakedJavaGlobalRef(std::nullptr_t) {}

  LeakedJavaGlobalRef(const LeakedJavaGlobalRef& other) = delete;
  LeakedJavaGlobalRef& operator=(const LeakedJavaGlobalRef&) = delete;
  LeakedJavaGlobalRef(LeakedJavaGlobalRef&& other) = default;
  LeakedJavaGlobalRef& operator=(LeakedJavaGlobalRef&&) = default;
  ~LeakedJavaGlobalRef() = default;

  void Reset() { JavaRef<T>::ResetGlobalRef(); }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void Reset(JNIEnv* env, const JavaRef<U>& j_object) {
    Reset();
    JavaRef<T>::SetNewGlobalRef(env, j_object.obj());
  }

  template <typename U>
    requires internal::IsConvertibleJObject<T, U>
  void ResetAndLeak(JNIEnv* env, const JavaRef<U>& j_object) {
    JavaRef<T>::SetNewGlobalRefAndLeak(env, j_object.obj());
  }

  // Create a local reference.
  ScopedJavaLocalRef<T> AsLocalRef(JNIEnv* env) const {
    T j_obj = JavaRef<T>::obj();
    if (!j_obj) {
      return nullptr;
    }
    return jni_zero::AdoptRef(env, static_cast<T>(env->NewLocalRef(j_obj)));
  }
};

template <typename T>
  requires internal::IsJobject<T>
inline ScopedJavaLocalRef<jobject> JavaRef<T>::Get(JNIEnv* env,
                                                   size_t index) const
  requires std::is_same_v<T, jobjectArray>
{
  jobject obj =
      env->GetObjectArrayElement(this->obj(), static_cast<int32_t>(index));
  return jni_zero::AdoptRef(env, obj);
}

template <typename T>
  requires internal::IsJobject<T>
template <typename U>
inline void JavaRef<T>::CopyTo(JNIEnv* env,
                               std::vector<ScopedJavaLocalRef<U>>* buf) const
  requires std::is_convertible_v<T, jobjectArray>
{
  jobjectArray arr = this->obj();
  int32_t length = this->GetLength(env);
  for (int32_t i = 0; i < length; i++) {
    jobject obj = env->GetObjectArrayElement(arr, i);
    buf->push_back(ScopedJavaLocalRef<U>::Adopt(env, static_cast<U>(obj)));
  }
}

template <typename T>
  requires internal::IsJobject<T>
inline ScopedJavaLocalRef<T> JavaRef<internal::_JObjectArray<T>*>::Get(
    JNIEnv* env,
    size_t index) const {
  jobject obj =
      env->GetObjectArrayElement(this->obj(), static_cast<int32_t>(index));
  return jni_zero::AdoptRef(env, static_cast<T>(obj));
}

// Identity overload for FromJniType when the destination type is already a
// JavaRef.
template <typename T>
  requires IsJavaRef<T>
inline T FromJniType(JNIEnv* env, const JavaRef<jobject>& obj) {
  return static_cast<T>(obj);
}

// Identity overload for ToJniType when the input type is already a JavaRef.
template <typename T>
  requires IsJavaRef<T>
inline decltype(auto) ToJniType(JNIEnv* env, T&& arg) {
  return std::forward<T>(arg);
}

}  // namespace jni_zero

#endif  // JNI_ZERO_JAVA_REFS_H_