/*
 * 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/reflection/creator.h>
#include <tvm/ffi/string.h>
#include <tvm/ffi/type_traits.h>

#include <iostream>
#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 = EscapeStringJSON(*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 Lowercase alias for DefaultValue, mirrors Python ``default``. */
class default_value : public DefaultValue {
  using DefaultValue::DefaultValue;
};

/*!
 * \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 Lowercase alias for DefaultFactory, mirrors Python ``default_factory``. */
class default_factory : public DefaultFactory {
  using DefaultFactory::DefaultFactory;
};

/*!
 * \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 InfoTrait to control whether a field participates in recursive comparison.
 *
 * Usage: `refl::compare(false)` marks a field to be excluded from
 * RecursiveEq/RecursiveLt/etc.
 */
class compare : public InfoTrait {
 public:
  /*!
   * \brief Constructor.
   * \param include Whether the field should participate in comparison.
   */
  explicit compare(bool include) : include_(include) {}

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

 private:
  bool include_;
};

/*!
 * \brief InfoTrait to control whether a field participates in recursive hashing.
 *
 * Usage: `refl::hash(false)` marks a field to be excluded from
 * RecursiveHash.
 */
class hash : public InfoTrait {
 public:
  /*!
   * \brief Constructor.
   * \param include Whether the field should participate in hashing.
   */
  explicit hash(bool include) : include_(include) {}

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

 private:
  bool include_;
};

// Forward-declare `init` so that the deduction guide can reference it.
template <typename... Args>
struct init;

/*!
 * \brief InfoTrait to mark a field as keyword-only in the auto-generated init.
 *
 * Usage: `refl::kw_only(true)` marks a field so it can only be provided via
 * the KWARGS calling convention in the auto-generated ``__ffi_init__``.
 */
class kw_only : public InfoTrait {
 public:
  /*!
   * \brief Constructor.
   * \param is_kw_only Whether the field is keyword-only.
   */
  explicit kw_only(bool is_kw_only) : kw_only_(is_kw_only) {}

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

 private:
  bool kw_only_;
};

/*!
 * \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 -
         ::tvm::ffi::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 = ::tvm::ffi::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 = ::tvm::ffi::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 = ::tvm::ffi::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 = ::tvm::ffi::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),
                 ::tvm::ffi::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 = ::tvm::ffi::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,
 * and (in its zero-argument specialization) an InfoTrait to control
 * whether a field participates in the auto-generated init.
 *
 * \tparam Args The argument types for the constructor.
 *
 * When used with one or more template arguments, ``init<Args...>`` is passed
 * to ``ObjectDef::def()`` to register an ``__init__`` method that constructs
 * an object with the specified argument types:
 *
 * \code{.cpp}
 * refl::ObjectDef<ExampleObject>()
 *     .def(refl::init<int64_t, int32_t>());
 * \endcode
 *
 * When used without template arguments (via CTAD), ``init(false)`` marks a
 * field to be excluded from the auto-generated ``__ffi_init__`` constructor,
 * or suppresses auto-init entirely when passed to ``ObjectDef``'s constructor:
 *
 * \code{.cpp}
 * // Per-field opt-out:
 * refl::ObjectDef<Foo>()
 *     .def_rw("hidden", &Foo::hidden, refl::init(false), refl::default_value(42));
 *
 * // Class-level opt-out:
 * refl::ObjectDef<Bar>(refl::init(false))
 *     .def_rw("x", &Bar::x);
 * \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 Specialization of ``init`` for zero template arguments.
 *
 * This specialization doubles as an ``InfoTrait`` that controls whether a field
 * participates in the auto-generated ``__ffi_init__``.
 *
 * - ``init(false)`` as a field trait: excludes the field from the auto-generated init.
 * - ``init(false)`` as an ``ObjectDef`` constructor argument: suppresses auto-init entirely.
 * - ``init<>()`` passed to ``ObjectDef::def()``: registers a zero-argument constructor.
 */
template <>
struct init<> : public InfoTrait {
  // Allow ObjectDef to access the execute function and include_ flag.
  template <typename Class>
  friend class ObjectDef;
  template <typename T>
  friend class OverloadObjectDef;

  /*!
   * \brief Default constructor (for zero-argument constructor registration).
   */
  constexpr init() noexcept = default;

  /*!
   * \brief Constructor for field/class-level init control.
   * \param include Whether the field should be included in the auto-generated init.
   */
  explicit init(bool include) : include_(include) {}

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

 private:
  bool include_ = true;

  /*!
   * \brief Execute the zero-argument constructor.
   * \tparam Class The class type.
   * \return The constructed object wrapped in an `ObjectRef`.
   */
  template <typename Class>
  static inline ObjectRef execute() {
    return ObjectRef(ffi::make_object<Class>());
  }
};

/*! \brief CTAD deduction guide: ``init(false)`` deduces to ``init<>``. */
#if !defined(TVM_FFI_DOXYGEN_MODE)
init(bool) -> init<>;
#endif

/*!
 * \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) {
    (MaybeSuppressAutoInit(extra_args), ...);
    RegisterExtraInfo(std::forward<ExtraArgs>(extra_args)...);
  }

  /*! \brief Non-copyable / non-movable. */
  ObjectDef(const ObjectDef&) = delete;
  ObjectDef& operator=(const ObjectDef&) = delete;
  ObjectDef(ObjectDef&&) = delete;
  ObjectDef& operator=(ObjectDef&&) = delete;

  /*!
   * \brief Destructor, which also potentially registers `__ffi_new__`, `__ffi_init__`,
   * `__ffi_shallow_copy__`.
   */
  ~ObjectDef() noexcept(false) {
    const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(type_index_);
    // Step 1. Register `__ffi_shallow_copy__` <== copy constructor (if it exists and is public)
    if constexpr (std::is_copy_constructible_v<Class>) {
      Function fn = Function::FromTyped(
          [](const Class* self) { return ObjectRef(ffi::make_object<Class>(*self)); },
          std::string(type_key_) + "." + type_attr::kShallowCopy);
      TVMFFIByteArray attr = AsByteArray(type_attr::kShallowCopy);
      TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny();
      TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any));
    }
    // Step 2. Register `__ffi_new__` <== info->metadata->creator
    // Also, `__ffi_init__` if no explicit init is defined.
    if (info->metadata != nullptr && info->metadata->creator != nullptr) {
      Function fn = Function::FromTyped(
          [creator = info->metadata->creator]() -> ObjectRef {
            TVMFFIObjectHandle handle;
            TVM_FFI_CHECK_SAFE_CALL(creator(&handle));
            return ObjectRef(::tvm::ffi::details::ObjectUnsafe::ObjectPtrFromOwned<Object>(
                static_cast<TVMFFIObject*>(handle)));
          },
          std::string(type_key_) + "." + type_attr::kNew);
      TVMFFIByteArray attr = AsByteArray(type_attr::kNew);
      TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny();
      TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any));
      // Step 3. Register `__ffi_init__` if no explicit init is defined.
      // Use Function::GetGlobal to look up the registration function, which lives in
      // dataclass.cc and may not be loaded yet for early (builtin) types.
      if (!has_explicit_init_) {
        if (std::optional<Function> opt = Function::GetGlobal("ffi._RegisterFFIInit")) {
          (*opt)(type_index_);
        }
      }
    }
  }

  /*!
   * \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) {
    has_explicit_init_ = true;
    // Register as TypeMethod (preserves type_schema metadata for Python stub generation).
    RegisterMethod(kInitMethodName, true, &init<Args...>::template execute<Class>,
                   std::forward<Extra>(extra)...);
    // Also mirror into __ffi_init__ TypeAttrColumn for runtime dispatch.
    const TVMFFITypeInfo* tinfo = TVMFFIGetTypeInfo(type_index_);
    constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit);
    for (int32_t i = 0; i < tinfo->num_methods; ++i) {
      if (tinfo->methods[i].name.size == attr_name.size &&
          std::strncmp(tinfo->methods[i].name.data, attr_name.data, attr_name.size) == 0) {
        TVM_FFI_CHECK_SAFE_CALL(
            TVMFFITypeRegisterAttr(type_index_, &attr_name, &tinfo->methods[i].method));
        break;
      }
    }
    return *this;
  }

 private:
  template <typename T>
  friend class OverloadObjectDef;

  /*! \brief If \p value is init(false), suppress auto-generated __ffi_init__. */
  template <typename T>
  void MaybeSuppressAutoInit(const T& value) {
    if constexpr (std::is_same_v<std::decay_t<T>, init<>>) {
      if (!value.include_) {
        has_explicit_init_ = true;
      }
    }
  }

  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 = reinterpret_cast<void*>(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", ::tvm::ffi::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 = ::tvm::ffi::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_;
  bool has_explicit_init_{false};
  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_
