# 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.
"""The ``py_class`` decorator: Python-defined FFI classes with dataclass semantics."""

from __future__ import annotations

import sys
import typing
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, ClassVar, TypeVar

from typing_extensions import dataclass_transform

from .. import core
from .._dunder import _install_dataclass_dunders
from ..core import MISSING, TypeSchema
from ..registry import _add_class_attrs
from .field import KW_ONLY, Field, field

_T = TypeVar("_T", bound=type)


# ---------------------------------------------------------------------------
# Module-level state
# ---------------------------------------------------------------------------
#
# Registration happens in two phases:
#
#   Phase 1 (_register_type_without_fields)
#       Allocates a C-level type index and inserts the class into the
#       global type registry.  This must happen early so that self-
#       referential and mutually-referential annotations can resolve
#       the class via ``TypeSchema.from_annotation()``.  Phase 1 always
#       succeeds (or raises immediately for non-Object parents).
#
#   Phase 2 (_register_fields_into_type)
#       Resolves string annotations via ``typing.get_type_hints``,
#       converts them to ``TypeSchema`` / ``Field`` objects, validates
#       field ordering, registers fields with the Cython layer, and
#       installs ``__init__``, ``__repr__``, ``__eq__``, etc.
#
#       If ``get_type_hints`` raises ``NameError`` (forward reference
#       not yet defined), the class is added to ``_PENDING_CLASSES``
#       and retried after each successful phase-2.  If phase-2 fails
#       for any other reason, ``_rollback_registration`` undoes phase-1
#       so the type key can be reused.
# ---------------------------------------------------------------------------


@dataclass
class _PendingClass:
    """Bookkeeping for a class whose annotations couldn't be resolved yet."""

    cls: type
    type_info: Any  # core.TypeInfo
    globalns: dict[str, Any]
    params: dict[str, Any]


#: Classes whose phase-2 (field registration) was deferred because
#: ``typing.get_type_hints`` raised ``NameError`` on an unresolved
#: forward reference.  Retried after each successful phase-2 via
#: :func:`_flush_pending`.
_PENDING_CLASSES: list[_PendingClass] = []

#: Per-module mapping of ``class.__name__ → class`` for every
#: ``@py_class``-decorated type.  Used as *localns* when resolving
#: annotations so that mutual references between classes in the same
#: module work even before the second class is assigned to the module
#: variable by Python.
_PY_CLASS_BY_MODULE: dict[str, dict[str, type]] = {}

# ---------------------------------------------------------------------------
# Phase 1: type registration
# ---------------------------------------------------------------------------


def _register_type_without_fields(cls: type, type_key: str | None) -> Any:
    """Phase 1: allocate type index and register the type (always succeeds)."""
    parent_info: core.TypeInfo | None = None
    for base in cls.__bases__:
        parent_info = core._type_cls_to_type_info(base)
        if parent_info is None:
            parent_info = getattr(base, "__tvm_ffi_type_info__", None)
        if parent_info is not None:
            break
    if parent_info is None:
        raise TypeError(
            f"{cls.__name__} must inherit from a registered FFI Object type (e.g. tvm_ffi.Object)"
        )
    if type_key is None:
        type_key = f"{cls.__module__}.{cls.__qualname__}"
    info = core._register_py_class(parent_info, type_key, cls)
    setattr(cls, "__tvm_ffi_type_info__", info)
    # Register in resolution namespace so sibling classes can find us
    _PY_CLASS_BY_MODULE.setdefault(cls.__module__, {})[cls.__name__] = cls
    return info


def _rollback_registration(cls: type, type_info: Any) -> None:
    """Undo :func:`_register_type_without_fields` after a phase-2 failure.

    The C-level type index is permanently consumed (cannot be reclaimed),
    but the Python-level registry dicts are cleaned up so a retry with
    the same type key does not hit "already registered".
    """
    # Remove from the Cython-level registry dicts (TYPE_KEY_TO_INFO,
    # TYPE_CLS_TO_INFO, TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS).
    core._rollback_py_class(type_info)  # ty: ignore[unresolved-attribute]
    # Remove from our own module-level resolution namespace.
    _PY_CLASS_BY_MODULE.get(cls.__module__, {}).pop(cls.__name__, None)
    for attr in ("__tvm_ffi_type_info__", "__tvm_ffi_is_dataclass__"):
        try:
            delattr(cls, attr)
        except AttributeError:
            pass


# ---------------------------------------------------------------------------
# Phase 2: annotation resolution, field registration, dunder installation
# ---------------------------------------------------------------------------


def _collect_own_fields(  # noqa: PLR0912
    cls: type,
    hints: dict[str, Any],
    decorator_kw_only: bool,
    decorator_frozen: bool,
) -> list[Field]:
    """Parse own annotations into :class:`Field` objects.

    - Skips ``ClassVar`` annotations.
    - Handles ``KW_ONLY`` sentinel.
    - Extracts ``Field`` metadata from class attributes (set via :func:`field`).
    - Handles bare defaults (non-``Field`` values).
    - Converts resolved types to ``TypeSchema``.
    - Resolves ``hash=None`` to follow ``compare``.
    """
    fields: list[Field] = []
    kw_only_active = decorator_kw_only
    # Python 3.14+ (PEP 749): annotations are lazily evaluated via
    # __annotate__ and no longer stored directly in __dict__.  getattr()
    # triggers evaluation and returns per-class annotations correctly.
    # On Python < 3.14, getattr() follows MRO and returns *parent*
    # annotations when the child has none — use __dict__ to avoid that.
    if sys.version_info >= (3, 14):
        own_annotations: dict[str, str] = getattr(cls, "__annotations__", {})
    else:
        own_annotations = cls.__dict__.get("__annotations__", {})

    for name in own_annotations:
        resolved_type = hints.get(name)
        # Skip ClassVar
        if (
            resolved_type is None
            or resolved_type is ClassVar
            or typing.get_origin(resolved_type) is ClassVar
        ):
            continue

        # KW_ONLY sentinel
        if resolved_type is KW_ONLY:
            kw_only_active = True
            if name in cls.__dict__:
                try:
                    delattr(cls, name)
                except AttributeError:
                    pass
            continue

        # Extract Field from class dict (inline of _pop_field_from_class)
        class_val = cls.__dict__.get(name, MISSING)
        if isinstance(class_val, Field):
            f = class_val
        elif class_val is not MISSING:
            f = field(default=class_val)
        else:
            f = field()
        if class_val is not MISSING:
            try:
                delattr(cls, name)
            except AttributeError:
                pass

        # Fill in name, schema, and resolved type (set by the decorator, not the user)
        f.name = name
        f._ty_schema = TypeSchema.from_annotation(resolved_type)
        f.type = resolved_type

        # Resolve kw_only: None means "inherit from decorator"
        if f.kw_only is None:
            f.kw_only = kw_only_active

        # Apply class-level frozen when the field doesn't explicitly set it
        if decorator_frozen and not f.frozen:
            f.frozen = True

        # Resolve hash=None → follow compare (native dataclass semantics)
        if f.hash is None:
            f.hash = f.compare

        fields.append(f)

    return fields


def method(fn: Any) -> Any:
    """Mark a ``@py_class`` method for FFI TypeMethod registration.

    Decorate any staticmethod or plain instance method on a ``@py_class``
    body to have it land in the C-level ``TVMFFITypeInfo.methods[]``
    table. Once registered, the method is resolvable by name from any
    FFI consumer — Python-side reflection via ``TypeInfo.methods``,
    C++, Rust — through the same path already used by C++-defined
    methods declared via ``refl::ObjectDef<T>().def(...)``.

    Example::

        from tvm_ffi import Object, method
        from tvm_ffi.dataclasses import py_class


        @py_class("example.Node")
        class Node(Object):
            x: int

            @method
            def label(self) -> str:
                return f"N({self.x})"


        # The method is now in ``TypeInfo.methods`` and FFI-callable:
        info = Node.__tvm_ffi_type_info__
        fn = next(m.func for m in info.methods if m.name == "label")
        fn(Node(x=7))  # -> "N(7)"

    ``staticmethod`` is supported: the marker is written onto the
    underlying function and unwrapped at registration time. Plain
    functions are also accepted — the marker lives on the function
    object directly. ``classmethod`` is rejected at decoration time
    because its ``cls``-first dispatch does not match the
    packed-call convention.
    """
    if isinstance(fn, staticmethod):
        fn.__func__.__ffi_method__ = True
        return fn
    if isinstance(fn, classmethod):
        raise TypeError(
            "@tvm_ffi.method: @classmethod is not supported for FFI "
            "TypeMethod registration — the classmethod's ``cls`` first "
            "arg does not match the packed-call convention. Use "
            "@staticmethod or a plain instance method instead.",
        )
    if not callable(fn):
        raise TypeError(
            f"@tvm_ffi.method: expected a callable, got {type(fn).__name__}.",
        )
    fn.__ffi_method__ = True
    return fn


def _is_method_marked(value: Any) -> bool:
    """Return True when ``value`` is a callable marked by :func:`method`."""
    if isinstance(value, (staticmethod, classmethod)):
        return getattr(value.__func__, "__ffi_method__", False) is True
    if callable(value):
        return getattr(value, "__ffi_method__", False) is True
    return False


def _validate_method_name(cls: type, name: str) -> None:
    """Reject ``@method``-marked names that collide with reserved namespaces.

    Names in :data:`_FFI_TYPE_ATTR_NAMES` and Python-protocol dunders
    are not allowed for ``@method`` — they are routed through the
    TypeAttrColumn / Python-protocol paths instead.
    """
    if name in _FFI_TYPE_ATTR_NAMES:
        raise NameError(
            f"@py_class({cls.__name__!r}): {name!r} is a TypeAttrColumn "
            "name — define it directly on the class body without "
            "``@method``; the FFI system routes it to ``TVMFFITypeRegisterAttr``.",
        )
    if name.startswith("__ffi_"):
        raise NameError(
            f"@py_class({cls.__name__!r}): {name!r} starts with the "
            "reserved ``__ffi_`` prefix. Pick a different name for your "
            "``@method``-decorated method.",
        )
    if name.startswith("__") and name.endswith("__"):
        raise NameError(
            f"@py_class({cls.__name__!r}): {name!r} is a Python protocol "
            "dunder — these are reserved for Python semantics and cannot "
            "be registered as FFI TypeMethods.",
        )


def _collect_py_methods(cls: type) -> list[tuple[str, Any, bool]] | None:
    """Extract FFI-registered entries from a ``@py_class`` body.

    Two sources are collected:

    1. **TypeAttrColumn dunders** — names in :data:`_FFI_RECOGNIZED_METHODS`
       that appear in ``cls.__dict__``. Both callables (e.g.
       ``__ffi_repr__``) and non-callable values flow here; the Cython
       layer routes them to ``TVMFFITypeRegisterAttr`` based on name.
    2. **User TypeMethods** — every callable in ``cls.__dict__`` marked
       with :func:`method`. Registered via ``TVMFFITypeRegisterMethod``
       so the method is resolvable by name from any FFI consumer
       (introspection through ``TypeInfo.methods``, name-based lookup
       from C++ / Rust, etc.). The decorator pattern keeps the
       per-class declaration co-located with the method body; no
       separate allowlist.

    Validation runs at registration time — reserved ``__ffi_*`` names
    and Python protocol dunders cannot be ``@method``-decorated; those
    are reserved by the TypeAttrColumn and Python semantics respectively.

    Returns the ``(name, value, is_static)`` list, or :data:`None` when
    no entries were found.
    """
    methods: list[tuple[str, Any, bool]] = []
    for name, value in cls.__dict__.items():
        marked = _is_method_marked(value)
        if name not in _FFI_RECOGNIZED_METHODS and not marked:
            continue
        if marked:
            _validate_method_name(cls, name)
        # In every case, registering a classmethod as a TypeMethod is
        # wrong: the packed-call convention places ``self`` (an instance)
        # in slot 0, but classmethod's descriptor binds slot 0 to the
        # class.
        if isinstance(value, classmethod):
            raise TypeError(
                f"@py_class({cls.__name__!r}): {name!r} is wrapped by "
                "@classmethod, which is incompatible with FFI "
                "registration — the cls-first arg breaks the packed-call "
                "convention. Use @staticmethod or a plain instance "
                "method. If you wrote ``@classmethod @method``, swap to "
                "``@staticmethod @method`` (or drop @classmethod).",
            )
        is_static = isinstance(value, staticmethod)
        func = value.__func__ if is_static else value
        methods.append((name, func, is_static))
    return methods if methods else None


def _build_localns(cls: type, *, cross_module: bool = False) -> dict[str, Any]:
    """Build the localns dict for resolving ``cls``'s annotations.

    By default, includes only classes from ``cls.__module__``, preserving
    standard Python name resolution semantics.  When ``cross_module=True``,
    also includes classes from all other registered modules as a fallback
    — this is needed when ``cls`` has a forward reference to a class in
    another module that can't appear in ``cls.__module__``'s globals due
    to a circular import (e.g. the target is imported only under
    ``if TYPE_CHECKING:``).

    Cross-module entries are added with ``setdefault`` so same-module
    classes and the class itself always take precedence over foreign
    classes with the same ``__name__``.
    """
    localns = dict(_PY_CLASS_BY_MODULE.get(cls.__module__, {}))
    localns[cls.__name__] = cls
    if cross_module:
        for mod_name, mod_classes in list(_PY_CLASS_BY_MODULE.items()):
            if mod_name == cls.__module__:
                continue
            for name, klass in mod_classes.items():
                localns.setdefault(name, klass)
    return localns


def _register_fields_into_type(
    cls: type,
    type_info: Any,
    globalns: dict[str, Any],
    params: dict[str, Any],
) -> bool:
    """Phase 2: resolve annotations, register fields, install dunders.

    Returns True on success, False if forward references are unresolved.
    """
    # Resolve string annotations to types; return False (defer) on NameError.
    #
    # First try with module-scoped localns (standard Python name resolution).
    # On NameError, retry with a cross-module localns that includes classes
    # from every registered module — this handles circular imports where the
    # target of a forward reference is imported only under TYPE_CHECKING and
    # therefore never enters the declaring module's globals.
    kwargs: dict[str, Any] = {"globalns": globalns, "localns": _build_localns(cls)}
    if sys.version_info >= (3, 11):
        kwargs["include_extras"] = True
    try:
        hints = typing.get_type_hints(cls, **kwargs)
    except (NameError, AttributeError):
        kwargs["localns"] = _build_localns(cls, cross_module=True)
        try:
            hints = typing.get_type_hints(cls, **kwargs)
        except (NameError, AttributeError):
            return False

    own_fields = _collect_own_fields(cls, hints, params["kw_only"], params["frozen"])
    py_methods = _collect_py_methods(cls)

    # Register fields and type-level structural eq/hash kind with the C layer.
    structure_kind = _STRUCTURE_KIND_MAP.get(params.get("structural_eq"))
    type_info._register_fields(own_fields, structure_kind)
    # Attach the user's Field sentinel to each TypeField so the
    # ``tvm_ffi.dataclasses.fields()`` compat layer can recover defaults
    # and default_factory values.  _register_fields preserves order, so
    # own_fields and type_info.fields line up 1:1.
    for py_field, type_field in zip(own_fields, type_info.fields):
        type_field.dataclass_field = py_field
    # Register user-defined dunder methods and read back system-generated ones.
    # Non-callable entries whose names are in _FFI_TYPE_ATTR_NAMES are routed
    # to TVMFFITypeRegisterAttr by the Cython layer.
    type_info._register_py_methods(py_methods, type_attr_names=_FFI_TYPE_ATTR_NAMES)
    _add_class_attrs(cls, type_info)

    # Remove deferred __init__ and restore user-defined __init__ if saved
    if "_py_class_deferred_init" in cls.__dict__:
        # Always remove the deferred wrapper
        if "__init__" in cls.__dict__:
            delattr(cls, "__init__")
        try:
            delattr(cls, "_py_class_deferred_init")
        except AttributeError:
            pass
        # Restore user-defined __init__ if it was saved
        user_init = cls.__dict__.get("_py_class_user_init")
        if user_init is not None:
            cls.__init__ = user_init
            delattr(cls, "_py_class_user_init")

    _install_dataclass_dunders(
        cls,
        init=params["init"],
        repr=params["repr"],
        eq=params["eq"],
        order=params["order"],
        unsafe_hash=params["unsafe_hash"],
        match_args=params["match_args"],
        py_class_mode=True,
    )
    return True


# ---------------------------------------------------------------------------
# Deferred resolution (when phase 2 cannot run at decoration time)
# ---------------------------------------------------------------------------


def _flush_pending() -> None:
    """Retry all pending classes.  Called after each successful phase 2."""
    changed = True
    while changed:
        changed = False
        remaining: list[_PendingClass] = []
        for entry in _PENDING_CLASSES:
            if _register_fields_into_type(entry.cls, entry.type_info, entry.globalns, entry.params):
                changed = True
            else:
                remaining.append(entry)
        _PENDING_CLASSES[:] = remaining


def _raise_unresolved_forward_reference(cls: type, globalns: dict[str, Any]) -> None:
    """Raise :class:`TypeError` listing the annotations that cannot be resolved."""
    localns = _build_localns(cls, cross_module=True)
    unresolved: list[str] = []
    for name, ann_str in getattr(cls, "__annotations__", {}).items():
        if isinstance(ann_str, str):
            try:
                eval(ann_str, globalns, localns)
            except NameError:
                unresolved.append(f"{name}: {ann_str}")
    raise TypeError(
        f"Cannot instantiate {cls.__name__}: unresolved forward references: {unresolved}"
    )


def _make_temporary_init(
    cls: type, type_info: Any, globalns: dict[str, Any], params: dict[str, Any]
) -> Callable[[...], None]:
    def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
        if type_info.fields is None:
            try:
                if not _register_fields_into_type(cls, type_info, globalns, params):
                    _raise_unresolved_forward_reference(cls, globalns)
                # cls stays in _PENDING_CLASSES after phase-2 succeeds; drop it
                # before _flush_pending so the loop doesn't hit the Cython-level
                # "_register_fields already called" assertion on a second pass.
                _PENDING_CLASSES[:] = [p for p in _PENDING_CLASSES if p.cls is not cls]
                _flush_pending()
            except Exception:
                # Remove from pending list and roll back so the type key can be reused.
                _PENDING_CLASSES[:] = [p for p in _PENDING_CLASSES if p.cls is not cls]
                _rollback_registration(cls, type_info)
                raise
        # cls.__init__ has been replaced by the real init (or restored user init)
        cls.__init__(self, *args, **kwargs)

    __init__.__qualname__ = f"{cls.__qualname__}.__init__"
    __init__.__module__ = cls.__module__
    return __init__


def _install_deferred_init(
    cls: type,
    type_info: Any,
    globalns: dict[str, Any],
    params: dict[str, Any],
) -> None:
    """Install a temporary ``__init__`` that completes registration on first call.

    Preserves a user-defined ``__init__`` if present in *cls.__dict__*;
    it is restored by :func:`_register_fields_into_type` after registration
    completes so that ``_install_dataclass_dunders`` sees it and skips
    auto-generation.
    """
    # Save user-defined __init__ before overwriting
    user_init = cls.__dict__.get("__init__")
    if user_init is not None:
        cls._py_class_user_init = user_init  # type: ignore[attr-defined]

    cls.__init__ = _make_temporary_init(cls, type_info, globalns, params)  # type: ignore[assignment]
    cls._py_class_deferred_init = True  # type: ignore[attr-defined]


# ---------------------------------------------------------------------------
# Main decorator
# ---------------------------------------------------------------------------


#: Mapping from Python string names to C-level ``TVMFFISEqHashKind`` enum values.
_STRUCTURE_KIND_MAP: dict[str | None, int] = {
    None: 0,  # kTVMFFISEqHashKindUnsupported (default; no metadata registered)
    "tree": 1,  # kTVMFFISEqHashKindTreeNode
    "var": 2,  # kTVMFFISEqHashKindFreeVar
    "dag": 3,  # kTVMFFISEqHashKindDAGNode
    "const-tree": 4,  # kTVMFFISEqHashKindConstTreeNode
    "singleton": 5,  # kTVMFFISEqHashKindUniqueInstance
}

#: Names that should be registered as TypeAttrColumn entries (for C++
#: dispatch via ``TypeAttrColumn``), NOT as TypeMethod.
#:
#: See ``reflection::type_attr`` in ``accessor.h`` for the C++ constants.
_FFI_TYPE_ATTR_NAMES: frozenset[str] = frozenset(
    {
        "__ffi_repr__",
        "__ffi_hash__",
        "__ffi_eq__",
        "__ffi_compare__",
        "__ffi_convert__",
        "__any_hash__",
        "__any_equal__",
        "__s_equal__",
        "__s_hash__",
        "__data_to_json__",
        "__data_from_json__",
    }
)

#: Allowlist of dunder names that ``_collect_py_methods`` collects from
#: the class body.  Names in ``_FFI_TYPE_ATTR_NAMES`` are registered as
#: TypeAttrColumn entries; all other names are registered as TypeMethod.
#:
#: System-managed names (``__ffi_new__``, ``__ffi_init__``,
#: ``__ffi_shallow_copy__``) are intentionally
#: absent because the C++ runtime generates them.
_FFI_RECOGNIZED_METHODS: frozenset[str] = _FFI_TYPE_ATTR_NAMES


@dataclass_transform(
    eq_default=False,
    order_default=False,
    field_specifiers=(field, Field),
)
def py_class(  # noqa: PLR0913
    cls_or_type_key: type | str | None = None,
    /,
    *,
    type_key: str | None = None,
    frozen: bool = False,
    init: bool = True,
    repr: bool = True,
    eq: bool = False,
    order: bool = False,
    unsafe_hash: bool = False,
    match_args: bool = True,
    kw_only: bool = False,
    structural_eq: str | None = None,
    slots: bool = True,
) -> Callable[[_T], _T] | _T:
    """Register a Python-defined FFI class with dataclass-style semantics.

    Can be used as:

    .. code-block:: python

        @py_class  # bare decorator
        class Point(Object):
            x: float
            y: float


        @py_class("my.Point")  # with explicit type_key
        class Point(Object): ...


        @py_class(eq=True, order=True)  # with options
        class Point(Object): ...


        @py_class("my.Point", eq=True)  # both
        class Point(Object): ...


        @py_class(structural_eq="tree")  # structural eq/hash kind
        class MyNode(Object):
            value: int
            span: Object = field(structural_eq="ignore")

    Parameters
    ----------
    cls_or_type_key
        When a string, used as the FFI type key.  When a type (bare
        decorator usage), the class to decorate.
    type_key
        Explicit FFI type key.  Auto-generated from
        ``{module}.{qualname}`` when omitted.
    frozen
        If True, all fields are read-only after ``__init__`` by default.
        Individual fields can still be marked ``field(frozen=True)`` on a
        non-frozen class.  Use ``type(obj).field_name.set(obj, value)``
        as an escape hatch when mutation is necessary.
    init
        If True (default), generate ``__init__`` from field annotations.
    repr
        If True (default), generate ``__repr__``.
    eq
        If True, generate ``__eq__`` and ``__ne__`` using recursive
        field-wise content equality.  Default False, in which case the
        class inherits the pointer-based ``__eq__`` from ``Object``
        (``a == b`` is equivalent to ``a.same_as(b)``).  If the class
        body defines ``__eq__`` or ``__ne__``, the generator is skipped
        and the user definition is preserved.
    order
        If True, generate ``__lt__``, ``__le__``, ``__gt__``, ``__ge__``.
        Requires ``eq=True``.
    unsafe_hash
        If True, generate ``__hash__`` using recursive field-wise
        content hashing (unsafe for mutable objects).  Default False,
        in which case the class inherits the handle-address ``__hash__``
        from ``Object``.  A user-defined ``__hash__`` in the class body
        is preserved.
    match_args
        If True (default), set ``__match_args__`` to a tuple of the
        positional ``__init__`` field names (``init=True`` and not
        ``kw_only``), enabling ``match`` statements.  Ignored when the
        class body already defines ``__match_args__``.
    kw_only
        If True, all fields are keyword-only in ``__init__`` by default.
    structural_eq
        Structural equality/hashing kind for this type.  Controls how
        instances participate in ``structural_equal`` and
        ``structural_hash``.  Valid values are:

        - ``None`` (default): structural comparison is not supported.
        - ``"tree"``: content-based comparison, the safe default for
          most IR nodes.
        - ``"var"``: compared by binding position, for variable types.
        - ``"dag"``: content + sharing-aware comparison, for dataflow
          graph nodes.
        - ``"const-tree"``: like ``"tree"`` with a pointer-equality
          fast path (only safe for types with no transitive ``"var"``
          children).
        - ``"singleton"``: pointer equality only, for singleton types.

        This parameter is **independent** of ``eq`` / ``unsafe_hash``:
        it only configures how ``structural_equal`` / ``structural_hash``
        walk the object in C++ and never installs or alters Python-level
        ``__eq__`` / ``__hash__``.  See Notes below.
    slots
        Accepted for ``dataclass_transform`` compatibility.  Object
        subclasses always use ``__slots__ = ()`` via the metaclass.

    Returns
    -------
    Callable | type
        A class decorator, or the decorated class (bare usage).

    Notes
    -----
    Three orthogonal equality/hashing mechanisms coexist on a
    ``@py_class`` type, each controlled by an independent knob:

    - ``a == b`` / ``hash(a)`` — selected by ``eq`` / ``unsafe_hash``
      params (or user-defined ``__eq__`` / ``__hash__`` in the class
      body).  Default: pointer-based ``same_as`` and handle-address
      hash, inherited from ``Object``.
    - ``structural_equal(a, b)`` / ``structural_hash(a)`` — selected
      by the ``structural_eq`` param.  Default (``None``): structural
      comparison is unsupported for this type.
    - ``a.same_as(b)`` — always available; always pointer comparison.

    The typical pattern is to leave ``eq`` / ``unsafe_hash`` at their
    defaults so ``==`` and ``hash()`` stay cheap and pointer-based
    (ideal for pass-internal bookkeeping such as visited-set tracking),
    and call ``structural_equal`` / ``structural_hash`` explicitly at
    the points that require the heavy semantic check.

    Combining ``eq=True`` (or ``unsafe_hash=True``) with a
    ``structural_eq`` kind is legal but gives the type two different
    recursive equalities — a Python-level one for ``==`` and a C++
    structural one for ``structural_equal`` — which rarely coincide.
    Prefer setting only one.

    """
    if order and not eq:
        raise ValueError("order=True requires eq=True")
    if structural_eq not in _STRUCTURE_KIND_MAP:
        raise ValueError(
            f"structural_eq must be one of "
            f"{sorted(k for k in _STRUCTURE_KIND_MAP if k is not None)}"
            f" or None, got {structural_eq!r}"
        )

    effective_type_key = type_key
    params: dict[str, Any] = {
        "frozen": frozen,
        "init": init,
        "repr": repr,
        "eq": eq,
        "order": order,
        "unsafe_hash": unsafe_hash,
        "match_args": match_args,
        "kw_only": kw_only,
        "structural_eq": structural_eq,
    }

    def decorator(cls: _T) -> _T:
        nonlocal effective_type_key
        globalns = getattr(sys.modules.get(cls.__module__, None), "__dict__", {})

        info = _register_type_without_fields(cls, effective_type_key)

        try:
            if _register_fields_into_type(cls, info, globalns, params):
                _flush_pending()
            else:
                _PENDING_CLASSES.append(_PendingClass(cls, info, globalns, params))
                _install_deferred_init(cls, info, globalns, params)
        except Exception:
            # Phase-2 failed (bad annotation, field ordering, etc.).
            # Roll back phase-1 so the type key can be reused after
            # the user fixes the error.
            _rollback_registration(cls, info)
            raise

        # Marker: distinguishes @c_class / @py_class types from FFI containers
        # (Array, List, Map, Dict) that also have __tvm_ffi_type_info__ but are
        # not dataclasses.  Used by is_dataclass() in common.py.
        setattr(cls, "__tvm_ffi_is_dataclass__", True)
        return cls

    # Handle different calling conventions:
    #   @py_class                → cls_or_type_key is the class
    #   @py_class("key")         → cls_or_type_key is a string
    #   @py_class()              → cls_or_type_key is None
    #   @py_class(eq=True)       → cls_or_type_key is None
    if cls_or_type_key is None:
        return decorator
    if isinstance(cls_or_type_key, str):
        effective_type_key = cls_or_type_key
        return decorator
    if isinstance(cls_or_type_key, type):
        return decorator(cls_or_type_key)
    raise TypeError(f"py_class: expected str or type, got {type(cls_or_type_key)}")
