/*
 * 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.
 */
/*!
 * \file tvm/ffi/reflection/registry.h
 * \brief Registry of reflection metadata.
 */
#ifndef TVM_FFI_REFLECTION_REGISTRY_H_
#define TVM_FFI_REFLECTION_REGISTRY_H_

#include <tvm/ffi/any.h>
#include <tvm/ffi/c_api.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/function_details.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <tvm/ffi/type_traits.h>

#include <iterator>
#include <optional>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

namespace tvm {
namespace ffi {
/*! \brief Reflection namespace */
namespace reflection {
/*!
 * \brief Types of temporary metadata hold in FieldInfoBuilder and MethodInfoBuilder,
 * before they are filled into final C metadata
 */
using _MetadataType = std::vector<std::pair<String, Any>>;  // NOLINT(bugprone-reserved-identifier)
/*!
 * \brief Builder for TVMFFIFieldInfo
 * \sa TVMFFIFieldInfo
 */
struct FieldInfoBuilder : public TVMFFIFieldInfo {
  /*! \brief Temporary metadata info to be filled into TVMFFIFieldInfo::metadata */
  _MetadataType metadata_;
};
/*!
 * \brief Builder for TVMFFIMethodInfo
 * \sa TVMFFIMethodInfo
 */
struct MethodInfoBuilder : public TVMFFIMethodInfo {
  /*! \brief Temporary metadata info to be filled into TVMFFIMethodInfo::metadata */
  _MetadataType metadata_;
};

/*!
 * \brief Trait that can be used to set information attached to a field or a method.
 * \sa DefaultValue, AttachFieldFlag
 */
struct InfoTrait {};

/*! \brief User-supplied metadata attached to a field or a method */
class Metadata : public InfoTrait {
 public:
  /*!
   * \brief Constructor
   * \param dict The initial dictionary
   */
  Metadata(std::initializer_list<std::pair<String, Any>> dict) : dict_(dict) {}
  /*!
   * \brief Move metadata into `FieldInfoBuilder`
   * \param info The field info builder.
   */
  inline void Apply(FieldInfoBuilder* info) const { this->Apply(&info->metadata_); }
  /*!
   * \brief Move metadata into `MethodInfoBuilder`
   * \param info The method info builder.
   */
  inline void Apply(MethodInfoBuilder* info) const { this->Apply(&info->metadata_); }

 private:
  friend class GlobalDef;
  template <typename T>
  friend class ObjectDef;
  template <typename T>
  friend class OverloadObjectDef;
  /*!
   * \brief Move metadata into a vector of key-value pairs.
   * \param out The output vector.
   */
  inline void Apply(_MetadataType* out) const {
    std::copy(std::make_move_iterator(dict_.begin()), std::make_move_iterator(dict_.end()),
              std::back_inserter(*out));
  }
  /*! \brief Convert the metadata to JSON string */
  static std::string ToJSON(const _MetadataType& metadata) {
    using ::tvm::ffi::details::StringObj;
    std::ostringstream os;
    os << "{";
    bool first = true;
    for (const auto& [key, value] : metadata) {
      if (!first) {
        os << ",";
      }
      os << "\"" << key << "\":";
      if (std::optional<int> v = value.as<int>()) {
        os << *v;
      } else if (std::optional<bool> v = value.as<bool>()) {
        os << (*v ? "true" : "false");
      } else if (std::optional<String> v = value.as<String>()) {
        String escaped = EscapeString(*v);
        os << escaped.c_str();
      } else {
        TVM_FFI_LOG_AND_THROW(TypeError) << "Metadata can be only int, bool or string, but on key `"
                                         << key << "`, the type is " << value.GetTypeKey();
      }
      first = false;
    }
    os << "}";
    return os.str();
  }

  std::vector<std::pair<String, Any>> dict_;
};
/*!
 * \brief Trait that can be used to set field default value
 */
class DefaultValue : public InfoTrait {
 public:
  /*!
   * \brief Constructor
   * \param value The value to be set
   */
  explicit DefaultValue(Any value) : value_(std::move(value)) {}

  /*!
   * \brief Apply the default value to the field info
   * \param info The field info.
   */
  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const {
    info->default_value_or_factory = AnyView(value_).CopyToTVMFFIAny();
    info->flags |= kTVMFFIFieldFlagBitMaskHasDefault;
  }

 private:
  Any value_;
};

/*!
 * \brief Trait that can be used to set field default factory.
 *
 * A default factory is a callable () -> Any that is invoked each time
 * a default value is needed, producing a fresh value. This is important
 * for mutable defaults (e.g., Array, Map) to avoid aliasing.
 */
class DefaultFactory : public InfoTrait {
 public:
  /*!
   * \brief Constructor
   * \param factory The factory function to be called to produce default values.
   */
  explicit DefaultFactory(Function factory) : factory_(std::move(factory)) {}

  /*!
   * \brief Apply the default factory to the field info
   * \param info The field info.
   */
  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const {
    info->default_value_or_factory = AnyView(factory_).CopyToTVMFFIAny();
    info->flags |= kTVMFFIFieldFlagBitMaskHasDefault;
    info->flags |= kTVMFFIFieldFlagBitMaskDefaultFromFactory;
  }

 private:
  Function factory_;
};

/*!
 * \brief Trait that can be used to attach field flag
 */
class AttachFieldFlag : public InfoTrait {
 public:
  /*!
   * \brief Attach a field flag to the field
   * \param flag The flag to be set
   */
  explicit AttachFieldFlag(int32_t flag) : flag_(flag) {}

  /*!
   * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashDef
   */
  TVM_FFI_INLINE static AttachFieldFlag SEqHashDef() {
    return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashDef);
  }
  /*!
   * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashIgnore
   */
  TVM_FFI_INLINE static AttachFieldFlag SEqHashIgnore() {
    return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashIgnore);
  }

  /*!
   * \brief Apply the field flag to the field info
   * \param info The field info.
   */
  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { info->flags |= flag_; }

 private:
  int32_t flag_;
};

/*!
 * \brief Trait that controls whether a field appears in repr output.
 *
 * By default, all fields appear in repr. Use `Repr(false)` to exclude a field.
 */
class Repr : public InfoTrait {
 public:
  /*!
   * \brief Constructor.
   * \param show Whether the field should appear in repr output.
   */
  explicit Repr(bool show) : show_(show) {}

  /*!
   * \brief Apply the repr flag to the field info.
   * \param info The field info.
   */
  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const {
    if (!show_) {
      info->flags |= kTVMFFIFieldFlagBitMaskReprOff;
    }
  }

 private:
  bool show_;
};

/*!
 * \brief Get the byte offset of a class member field.
 *
 * \tparam The original class.
 * \tparam T the field type.
 *
 * \param field_ptr A class member pointer
 * \returns The byteoffset
 */
template <typename Class, typename T>
TVM_FFI_INLINE int64_t GetFieldByteOffsetToObject(T Class::* field_ptr) {
  int64_t field_offset_to_class =
      reinterpret_cast<int64_t>(&(static_cast<Class*>(nullptr)->*field_ptr));
  return field_offset_to_class - details::ObjectUnsafe::GetObjectOffsetToSubclass<Class>();
}

/// \cond Doxygen_Suppress
class ReflectionDefBase {
 protected:
  template <typename T>
  static int FieldGetter(void* field, TVMFFIAny* result) {
    TVM_FFI_SAFE_CALL_BEGIN();
    *result = details::AnyUnsafe::MoveAnyToTVMFFIAny(Any(*reinterpret_cast<T*>(field)));
    TVM_FFI_SAFE_CALL_END();
  }

  template <typename T>
  static int FieldSetter(void* field, const TVMFFIAny* value) {
    TVM_FFI_SAFE_CALL_BEGIN();
    if constexpr (std::is_same_v<T, Any>) {
      *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value);
    } else {
      *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value).cast<T>();
    }
    TVM_FFI_SAFE_CALL_END();
  }

  template <typename T>
  static int ObjectCreatorDefault(TVMFFIObjectHandle* result) {
    TVM_FFI_SAFE_CALL_BEGIN();
    ObjectPtr<T> obj = make_object<T>();
    *result = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(obj));
    TVM_FFI_SAFE_CALL_END();
  }

  template <typename T>
  static int ObjectCreatorUnsafeInit(TVMFFIObjectHandle* result) {
    TVM_FFI_SAFE_CALL_BEGIN();
    ObjectPtr<T> obj = make_object<T>(UnsafeInit{});
    *result = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(obj));
    TVM_FFI_SAFE_CALL_END();
  }

  template <typename T>
  TVM_FFI_INLINE static void ApplyFieldInfoTrait(FieldInfoBuilder* info, const T& value) {
    if constexpr (std::is_base_of_v<InfoTrait, std::decay_t<T>>) {
      value.Apply(info);
    }
    if constexpr (std::is_same_v<std::decay_t<T>, char*>) {
      info->doc = TVMFFIByteArray{value, std::char_traits<char>::length(value)};
    }
  }

  template <typename T>
  TVM_FFI_INLINE static void ApplyMethodInfoTrait(MethodInfoBuilder* info, const T& value) {
    if constexpr (std::is_base_of_v<InfoTrait, std::decay_t<T>>) {
      value.Apply(info);
    }
    if constexpr (std::is_same_v<std::decay_t<T>, char*>) {
      info->doc = TVMFFIByteArray{value, std::char_traits<char>::length(value)};
    }
  }

  template <typename T>
  TVM_FFI_INLINE static void ApplyExtraInfoTrait(TVMFFITypeMetadata* info, const T& value) {
    if constexpr (std::is_same_v<std::decay_t<T>, char*>) {
      info->doc = TVMFFIByteArray{value, std::char_traits<char>::length(value)};
    }
  }

  template <typename Func>
  TVM_FFI_INLINE static Function GetMethod(std::string name, Func&& func) {
    return ffi::Function::FromTyped(WrapFunction(std::forward<Func>(func)), std::move(name));
  }

  template <typename Func>
  TVM_FFI_INLINE static Func&& WrapFunction(Func&& func) {
    return std::forward<Func>(func);
  }
  template <typename Class, typename R, typename... Args>
  TVM_FFI_INLINE static auto WrapFunction(R (Class::*func)(Args...)) {
    static_assert(std::is_base_of_v<ObjectRef, Class> || std::is_base_of_v<Object, Class>,
                  "Class must be derived from ObjectRef or Object");
    if constexpr (std::is_base_of_v<ObjectRef, Class>) {
      return [func](Class target, Args... params) -> R {
        // call method pointer
        return (target.*func)(std::forward<Args>(params)...);
      };
    }
    if constexpr (std::is_base_of_v<Object, Class>) {
      return [func](const Class* target, Args... params) -> R {
        // call method pointer
        return (const_cast<Class*>(target)->*func)(std::forward<Args>(params)...);
      };
    }
  }
  template <typename Class, typename R, typename... Args>
  TVM_FFI_INLINE static auto WrapFunction(R (Class::*func)(Args...) const) {
    static_assert(std::is_base_of_v<ObjectRef, Class> || std::is_base_of_v<Object, Class>,
                  "Class must be derived from ObjectRef or Object");
    if constexpr (std::is_base_of_v<ObjectRef, Class>) {
      return [func](const Class& target, Args... params) -> R {
        // call method pointer
        return (target.*func)(std::forward<Args>(params)...);
      };
    }
    if constexpr (std::is_base_of_v<Object, Class>) {
      return [func](const Class* target, Args... params) -> R {
        // call method pointer
        return (target->*func)(std::forward<Args>(params)...);
      };
    }
  }
};
/// \endcond

/*!
 * \brief GlobalDef helper to register a global function.
 *
 * \code{.cpp}
 * namespace refl = tvm::ffi::reflection;
 * refl::GlobalDef().def("my_ffi_extension.my_function", MyFunction);
 * \endcode
 */
class GlobalDef : public ReflectionDefBase {
 public:
  /*!
   * \brief Define a global function.
   *
   * \tparam Func The function type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the function.
   * \param func The function to be registered.
   * \param extra The extra arguments that can be docstring or subclass of InfoTrait.
   *
   * \return The reflection definition.
   */
  template <typename Func, typename... Extra>
  GlobalDef& def(const char* name, Func&& func, Extra&&... extra) {
    using FuncInfo = details::FunctionInfo<std::decay_t<Func>>;
    RegisterFunc(name, ffi::Function::FromTyped(std::forward<Func>(func), std::string(name)),
                 FuncInfo::TypeSchema(), std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Define a global function in ffi::PackedArgs format.
   *
   * \tparam Func The function type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the function.
   * \param func The function to be registered.
   * \param extra The extra arguments that can be docstring or subclass of InfoTrait.
   *
   * \return The reflection definition.
   */
  template <typename Func, typename... Extra>
  GlobalDef& def_packed(const char* name, Func func, Extra&&... extra) {
    RegisterFunc(name, ffi::Function::FromPacked(func), details::TypeSchemaImpl<Function>::v(),
                 std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Expose a class method as a global function.
   *
   * An argument will be added to the first position if the function is not static.
   *
   * \tparam Class The class type.
   * \tparam Func The function type.
   *
   * \param name The name of the method.
   * \param func The function to be registered.
   * \param extra The extra arguments that can be docstring.
   *
   * \return The reflection definition.
   */
  template <typename Func, typename... Extra>
  GlobalDef& def_method(const char* name, Func&& func, Extra&&... extra) {
    using FuncInfo = details::FunctionInfo<std::decay_t<Func>>;
    RegisterFunc(name, GetMethod(std::string(name), std::forward<Func>(func)),
                 FuncInfo::TypeSchema(), std::forward<Extra>(extra)...);
    return *this;
  }

 private:
  template <typename... Extra>  // NOLINTNEXTLINE(performance-unnecessary-value-param)
  void RegisterFunc(const char* name, ffi::Function func, String type_schema, Extra&&... extra) {
    MethodInfoBuilder info;
    info.name = TVMFFIByteArray{name, std::char_traits<char>::length(name)};
    info.doc = TVMFFIByteArray{nullptr, 0};
    info.flags = 0;
    info.method = AnyView(func).CopyToTVMFFIAny();
    info.metadata_.emplace_back("type_schema", type_schema);
    ((ApplyMethodInfoTrait(&info, std::forward<Extra>(extra)), ...));
    std::string metadata_str = Metadata::ToJSON(info.metadata_);
    info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()};
    TVM_FFI_CHECK_SAFE_CALL(TVMFFIFunctionSetGlobalFromMethodInfo(&info, 0));
  }
};

/*!
 * \brief Helper class to register a constructor method for object types.
 *
 * This helper is used with `ObjectDef::def()` to register an `__init__` method
 * that constructs an object instance with the specified argument types.
 *
 * \tparam Args The argument types for the constructor.
 *
 * Example usage:
 *
 * \code{.cpp}
 * class ExampleObject : public Object {
 *  public:
 *   int64_t v_i64;
 *   int32_t v_i32;
 *
 *   ExampleObject(int64_t v_i64, int32_t v_i32) : v_i64(v_i64), v_i32(v_i32) {}
 *   TVM_FFI_DECLARE_OBJECT_INFO("example.ExampleObject", ExampleObject, Object);
 * };
 *
 * // Register the constructor
 * refl::ObjectDef<ExampleObject>()
 *     .def(refl::init<int64_t, int32_t>());
 * \endcode
 *
 * \note The object type is automatically deduced from the `ObjectDef` context.
 */
template <typename... Args>
struct init {
  // Allow ObjectDef to access the execute function
  template <typename Class>
  friend class ObjectDef;
  template <typename T>
  friend class OverloadObjectDef;

  /*!
   * \brief Constructor
   */
  constexpr init() noexcept = default;

 private:
  /*!
   * \brief Execute the constructor
   * \tparam Class The class type.
   * \param args The arguments to be passed to the constructor.
   * \return The constructed object wrapped in an `ObjectRef`.
   */
  template <typename Class>
  static inline ObjectRef execute(Args&&... args) {
    return ObjectRef(ffi::make_object<Class>(std::forward<Args>(args)...));
  }
};

/*! \brief Well-known type attribute names used by the reflection system. */
namespace type_attr {
inline constexpr const char* kInit = "__ffi_init__";
inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__";
inline constexpr const char* kRepr = "__ffi_repr__";
}  // namespace type_attr

/*!
 * \brief Helper to register Object's reflection metadata.
 * \tparam Class The class type.
 *
 * \code{.cpp}
 * namespace refl = tvm::ffi::reflection;
 * refl::ObjectDef<MyClass>().def_ro("my_field", &MyClass::my_field);
 * \endcode
 */
template <typename Class>
class ObjectDef : public ReflectionDefBase {
 public:
  /*!
   * \brief Constructor
   * \tparam ExtraArgs The extra arguments.
   * \param extra_args The extra arguments.
   */
  template <typename... ExtraArgs>
  explicit ObjectDef(ExtraArgs&&... extra_args)
      : type_index_(Class::_GetOrAllocRuntimeTypeIndex()), type_key_(Class::_type_key) {
    RegisterExtraInfo(std::forward<ExtraArgs>(extra_args)...);
    AutoRegisterCopy();
  }

  /*!
   * \brief Define a readonly field.
   *
   * \tparam Class The class type.
   * \tparam T The field type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the field.
   * \param field_ptr The pointer to the field.
   * \param extra The extra arguments that can be docstring or default value.
   *
   * \return The reflection definition.
   */
  template <typename T, typename BaseClass, typename... Extra>
  TVM_FFI_INLINE ObjectDef& def_ro(const char* name, T BaseClass::* field_ptr, Extra&&... extra) {
    RegisterField(name, field_ptr, false, std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Define a read-write field.
   *
   * \tparam Class The class type.
   * \tparam T The field type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the field.
   * \param field_ptr The pointer to the field.
   * \param extra The extra arguments that can be docstring or default value.
   *
   * \return The reflection definition.
   */
  template <typename T, typename BaseClass, typename... Extra>
  TVM_FFI_INLINE ObjectDef& def_rw(const char* name, T BaseClass::* field_ptr, Extra&&... extra) {
    static_assert(Class::_type_mutable, "Only mutable classes are supported for writable fields");
    RegisterField(name, field_ptr, true, std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Define a method.
   *
   * \tparam Func The function type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the method.
   * \param func The function to be registered.
   * \param extra The extra arguments that can be docstring.
   *
   * \return The reflection definition.
   */
  template <typename Func, typename... Extra>
  TVM_FFI_INLINE ObjectDef& def(const char* name, Func&& func, Extra&&... extra) {
    RegisterMethod(name, false, std::forward<Func>(func), std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Define a static method.
   *
   * \tparam Func The function type.
   * \tparam Extra The extra arguments.
   *
   * \param name The name of the method.
   * \param func The function to be registered.
   * \param extra The extra arguments that can be docstring.
   *
   * \return The reflection definition.
   */
  template <typename Func, typename... Extra>
  TVM_FFI_INLINE ObjectDef& def_static(const char* name, Func&& func, Extra&&... extra) {
    RegisterMethod(name, true, std::forward<Func>(func), std::forward<Extra>(extra)...);
    return *this;
  }

  /*!
   * \brief Register a constructor for this object type.
   *
   * This method registers a static `__init__` method that constructs an instance
   * of the object with the specified argument types. The constructor can be invoked
   * from Python or other FFI bindings.
   *
   * \tparam Args The argument types for the constructor.
   * \tparam Extra Additional arguments (e.g., docstring).
   *
   * \param init_func An instance of `init<Args...>` specifying constructor signature.
   * \param extra Optional additional metadata such as docstring.
   *
   * \return Reference to this `ObjectDef` for method chaining.
   *
   * Example:
   *
   * \code{.cpp}
   * refl::ObjectDef<MyObject>()
   *     .def(refl::init<int64_t, std::string>(), "Constructor docstring");
   * \endcode
   */
  template <typename... Args, typename... Extra>
  TVM_FFI_INLINE ObjectDef& def([[maybe_unused]] init<Args...> init_func, Extra&&... extra) {
    RegisterMethod(kInitMethodName, true, &init<Args...>::template execute<Class>,
                   std::forward<Extra>(extra)...);
    return *this;
  }

 private:
  template <typename T>
  friend class OverloadObjectDef;

  /*! \brief Shallow-copy \p self via the C++ copy constructor. */
  static ObjectRef ShallowCopy(const Class* self) {
    return ObjectRef(ffi::make_object<Class>(*self));
  }

  void AutoRegisterCopy() {
    if constexpr (std::is_copy_constructible_v<Class>) {
      // Register __ffi_shallow_copy__ as an instance method
      RegisterMethod(type_attr::kShallowCopy, false, &ObjectDef::ShallowCopy);
      // Also register as a type attribute for generic deep copy lookup
      Function copy_fn = GetMethod(std::string(type_key_) + "." + type_attr::kShallowCopy,
                                   &ObjectDef::ShallowCopy);
      TVMFFIByteArray attr_name = {type_attr::kShallowCopy,
                                   std::char_traits<char>::length(type_attr::kShallowCopy)};
      TVMFFIAny attr_value = AnyView(copy_fn).CopyToTVMFFIAny();
      TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr_name, &attr_value));
    }
  }

  template <typename... ExtraArgs>
  void RegisterExtraInfo(ExtraArgs&&... extra_args) {
    TVMFFITypeMetadata info;
    info.total_size = sizeof(Class);
    info.structural_eq_hash_kind = Class::_type_s_eq_hash_kind;
    info.creator = nullptr;
    info.doc = TVMFFIByteArray{nullptr, 0};
    if constexpr (std::is_default_constructible_v<Class>) {
      info.creator = ObjectCreatorDefault<Class>;
    } else if constexpr (std::is_constructible_v<Class, UnsafeInit>) {
      info.creator = ObjectCreatorUnsafeInit<Class>;
    }
    // apply extra info traits
    ((ApplyExtraInfoTrait(&info, std::forward<ExtraArgs>(extra_args)), ...));
    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMetadata(type_index_, &info));
  }

  template <typename T, typename BaseClass, typename... ExtraArgs>
  void RegisterField(const char* name, T BaseClass::* field_ptr, bool writable,
                     ExtraArgs&&... extra_args) {
    static_assert(std::is_base_of_v<BaseClass, Class>, "BaseClass must be a base class of Class");
    FieldInfoBuilder info;
    info.name = TVMFFIByteArray{name, std::char_traits<char>::length(name)};
    info.field_static_type_index = TypeToFieldStaticTypeIndex<T>::value;
    // store byte offset and setter, getter
    // so the same setter can be reused for all the same type
    info.offset = GetFieldByteOffsetToObject<Class, T>(field_ptr);
    info.size = sizeof(T);
    info.alignment = alignof(T);
    info.flags = 0;
    if (writable) {
      info.flags |= kTVMFFIFieldFlagBitMaskWritable;
    }
    info.getter = FieldGetter<T>;
    info.setter = FieldSetter<T>;
    // initialize default value to nullptr
    info.default_value_or_factory = AnyView(nullptr).CopyToTVMFFIAny();
    info.doc = TVMFFIByteArray{nullptr, 0};
    info.metadata_.emplace_back("type_schema", details::TypeSchema<T>::v());
    // apply field info traits
    ((ApplyFieldInfoTrait(&info, std::forward<ExtraArgs>(extra_args)), ...));
    // call register
    std::string metadata_str = Metadata::ToJSON(info.metadata_);
    info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()};
    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterField(type_index_, &info));
  }

  // register a method
  template <typename Func, typename... Extra>
  void RegisterMethod(const char* name, bool is_static, Func&& func, Extra&&... extra) {
    using FuncInfo = details::FunctionInfo<std::decay_t<Func>>;
    MethodInfoBuilder info;
    info.name = TVMFFIByteArray{name, std::char_traits<char>::length(name)};
    info.doc = TVMFFIByteArray{nullptr, 0};
    info.flags = 0;
    if (is_static) {
      info.flags |= kTVMFFIFieldFlagBitMaskIsStaticMethod;
    }

    // obtain the method function
    Function method = GetMethod(std::string(type_key_) + "." + name, std::forward<Func>(func));
    info.method = AnyView(method).CopyToTVMFFIAny();
    info.metadata_.emplace_back("type_schema", FuncInfo::TypeSchema());
    // apply method info traits
    ((ApplyMethodInfoTrait(&info, std::forward<Extra>(extra)), ...));
    std::string metadata_str = Metadata::ToJSON(info.metadata_);
    info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()};
    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index_, &info));
  }

  int32_t type_index_;
  const char* type_key_;
  static constexpr const char* kInitMethodName = type_attr::kInit;
};

/*!
 * \brief Helper to register type attribute.
 * \tparam Class The class type.
 * \tparam ExtraArgs The extra arguments.
 *
 * \code{.cpp}
 * namespace refl = tvm::ffi::reflection;
 * refl::TypeAttrDef<MyClass>().def("func_attr", MyFunc);
 * \endcode
 */
template <typename Class, typename = std::enable_if_t<std::is_base_of_v<Object, Class>>>
class TypeAttrDef : public ReflectionDefBase {
 public:
  /*!
   * \brief Constructor
   * \tparam ExtraArgs The extra arguments.
   * \param extra_args The extra arguments.
   */
  template <typename... ExtraArgs>
  explicit TypeAttrDef(ExtraArgs&&... extra_args)
      : type_index_(Class::RuntimeTypeIndex()), type_key_(Class::_type_key) {}

  /*!
   * \brief Define a function-valued type attribute.
   *
   * \tparam Func The function type.
   *
   * \param name The name of the function.
   * \param func The function to be registered.
   *
   * \return The TypeAttrDef object.
   */
  template <typename Func>
  TypeAttrDef& def(const char* name, Func&& func) {
    TVMFFIByteArray name_array = {name, std::char_traits<char>::length(name)};
    ffi::Function ffi_func =
        GetMethod(std::string(type_key_) + "." + name, std::forward<Func>(func));
    TVMFFIAny value_any = AnyView(ffi_func).CopyToTVMFFIAny();
    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any));
    return *this;
  }

  /*!
   * \brief Define a constant-valued type attribute.
   *
   * \tparam T The type of the value.
   *
   * \param name The name of the attribute.
   * \param value The value of the attribute.
   *
   * \return The TypeAttrDef object.
   */
  template <typename T>
  TypeAttrDef& attr(const char* name, T value) {
    TVMFFIByteArray name_array = {name, std::char_traits<char>::length(name)};
    TVMFFIAny value_any = AnyView(value).CopyToTVMFFIAny();
    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any));
    return *this;
  }

 private:
  int32_t type_index_;
  const char* type_key_;
};

/*!
 * \brief Ensure the type attribute column is presented in the system.
 *
 * \param name The name of the type attribute.
 */
inline void EnsureTypeAttrColumn(std::string_view name) {
  TVMFFIByteArray name_array = {name.data(), name.size()};
  AnyView any_view(nullptr);
  TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_array,
                                                 reinterpret_cast<const TVMFFIAny*>(&any_view)));
}

}  // namespace reflection
}  // namespace ffi
}  // namespace tvm
#endif  // TVM_FFI_REFLECTION_REGISTRY_H_