# 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.
"""Cross-language enum types: named, frozen, ordinal-indexed singletons.

An ``Enum`` subclass has one of two usage modes, distinguished by whether
its ``type_key`` is already registered in the FFI type system:

* **Closed Python enum** — fresh ``type_key``, variants declared once in
  the class body.  Behavior matches ``enum.Enum``.
* **Cross-language registry** — ``type_key`` also registered in C++ (or
  another Python module).  Python and C++ both contribute variants to
  the same per-class registry, and consumers attach *extensible
  attributes* to variants from any module at any time.

See :class:`Enum` for declaration forms and :meth:`Enum.def_attr` for
extensible attributes.

Storage layout (mirrors ``include/tvm/ffi/enum.h``):

* ``__ffi_enum_entries__`` — ``Dict[str, Enum]``, name → variant.
* ``__ffi_enum_attrs__``   — ``Dict[str, List[Any]]``, extensible-attr
  name → column indexed by each variant's ordinal.
"""

from __future__ import annotations

import sys
import typing
from collections.abc import Callable, Iterator
from typing import Any, ClassVar

from typing_extensions import dataclass_transform

from .. import core
from ..container import Dict, List
from ..core import Object
from .c_class import c_class
from .field import Field, field
from .py_class import py_class

__all__ = [
    "ENUM_ATTRS_ATTR",
    "ENUM_ENTRIES_ATTR",
    "ENUM_VALUE_ENTRIES_ATTR",
    "Enum",
    "EnumAttrMap",
    "IntEnum",
    "StrEnum",
    "auto",
    "entry",
]

#: TypeAttr column storing ``Dict[str, Enum]`` (instance name → singleton).
ENUM_ENTRIES_ATTR = "__ffi_enum_entries__"

#: TypeAttr column storing ``Dict[str, List[Any]]`` of per-variant attrs.
ENUM_ATTRS_ATTR = "__ffi_enum_attrs__"

#: TypeAttr column storing ``Dict[Any, Enum]`` (payload → singleton) —
#: parallel to :data:`ENUM_ENTRIES_ATTR` but keyed by the user-visible
#: payload (``IntEnum.value`` / ``StrEnum.value``).  Populated by every
#: creator of a payload-carrying variant; consumed by FFI converters for
#: O(1) payload → variant resolution.
ENUM_VALUE_ENTRIES_ATTR = "__ffi_enum_value_entries__"


# ---------------------------------------------------------------------------
# entry() sentinel
# ---------------------------------------------------------------------------


class _EnumEntry:
    """Sentinel produced by :func:`entry`; consumed by ``Enum.__init_subclass__``.

    Holds the positional and keyword arguments forwarded to the subclass's
    ``__init__`` when the variant is materialized.  ``_value`` and ``_name``
    are auto-assigned (dense ordinal and class-body name) and must not
    appear in the captured arguments.
    """

    __slots__ = ("args", "kwargs")

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args: tuple[Any, ...] = args
        self.kwargs: dict[str, Any] = kwargs

    def __repr__(self) -> str:
        parts = [repr(a) for a in self.args]
        parts.extend(f"{k}={v!r}" for k, v in self.kwargs.items())
        return f"entry({', '.join(parts)})"


def entry(*args: Any, **kwargs: Any) -> Any:
    """Declare a new enum variant with values for its declared fields.

    ``entry(...)`` is a class-body sentinel; it never produces a real
    instance.  At class creation, :meth:`Enum.__init_subclass__` scans
    for these sentinels and, for each one, constructs a singleton variant
    by forwarding the captured positional and keyword arguments to the
    subclass's ``__init__``, together with an auto-assigned
    :attr:`~Enum._value` (dense ordinal) and :attr:`~Enum._name`
    (class-body name).

    Prefer :func:`auto` when a variant has no declared fields beyond the
    auto-assigned ordinal and name — it expresses intent without the
    empty-arg-list noise.

    When the enum's ``type_key`` is C++-backed (registered via
    ``refl::ObjectDef``), only keyword arguments are supported — field
    values are assigned via reflected setters keyed by name.  Passing
    positional arguments in that case raises :class:`TypeError`.
    ``entry(_value=...)`` and ``entry(_name=...)`` always raise
    :class:`TypeError` because those fields are auto-assigned.

    Examples
    --------
    Variant with declared fields:

    .. code-block:: python

        from typing import ClassVar


        class Activation(Enum, type_key="my.Activation"):
            output_zero: bool
            is_monotonic: bool

            relu: ClassVar[Activation] = entry(output_zero=True, is_monotonic=True)
            gelu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=False)

    Returns
    -------
    object
        An opaque sentinel.  The declared return type is ``Any`` so that
        ``ClassVar[Cls] = entry(...)`` type-checks even though the sentinel
        is not a real ``Cls``.

    """
    return _EnumEntry(*args, **kwargs)


def auto() -> Any:
    """Declare a new enum variant with no declared fields.

    Semantically equivalent to :func:`entry` called with no arguments but
    reads more clearly for the common case where a variant differs from
    its siblings only by name and ordinal.  The resulting singleton has
    only the auto-assigned :attr:`~Enum._value` and :attr:`~Enum._name`.

    ``auto()`` registers a *new* Python-side variant; it is not the right
    tool for binding to a pre-existing C++-registered entry (use a bare
    ``ClassVar[Cls]`` annotation for that — see :class:`Enum`).

    Examples
    --------
    .. code-block:: python

        class Status(Enum, type_key="my.Status"):
            ok = auto()
            err = auto()
            retry = auto()


        assert Status.ok._value == 0
        assert Status.err._name == "err"

    Returns
    -------
    object
        An opaque sentinel, the same kind returned by :func:`entry`.  The
        declared return type is ``Any`` so that both ``name = auto()`` and
        ``name: ClassVar[Cls] = auto()`` type-check.

    """
    return _EnumEntry()


# ---------------------------------------------------------------------------
# Class-level helpers
# ---------------------------------------------------------------------------


class _ClassProperty:
    """Read-only descriptor whose getter receives the owning class.

    Used for ``attr_dict`` so it works as a class-level attribute access
    (e.g., ``Op.attr_dict["has_side_effects"]``) without needing a metaclass.
    """

    __slots__ = ("_fget",)

    def __init__(self, fget: Callable[[type], Any]) -> None:
        self._fget = fget

    def __get__(self, instance: Any, owner: type | None = None) -> Any:
        cls = owner if owner is not None else type(instance)
        return self._fget(cls)


def _install_payload_enum_behaviors(cls: type, *, user_defined_repr: bool) -> None:
    """Install stdlib-like instance behavior for ``IntEnum`` / ``StrEnum``.

    Payload enums have two identities that plain ``Enum`` does not: the FFI
    singleton identity (``Priority.low.same_as(...)``) and a user-visible raw
    payload (``Priority.low.value == 10``).  These methods make the Python
    surface follow the payload where users expect stdlib-like behavior, while
    keeping the object identity available through ``same_as``.

    This runs after ``py_class`` / ``c_class`` because those decorators may
    install generated object dunders.  Installing here lets payload enums
    override generated repr/equality/hash defaults, but still preserves methods
    explicitly written in the enum subclass body.
    """

    def _payload_enum_eq(self: Enum, other: object) -> Any:
        """Compare payload enum variants by payload and accept raw payloads.

        Two variants of the same enum class compare like their ``value`` fields,
        and a variant can compare directly with its raw payload type
        (``Priority.low == 10``).  Unrelated objects return ``NotImplemented``
        so Python can try reflected comparison or fall back normally.
        """
        if isinstance(other, type(self)):
            return self.value == other.value  # type: ignore[attr-defined]
        payload_type = getattr(type(self), "__ffi_enum_payload_value_type__", None)
        if payload_type is not None and isinstance(other, payload_type):
            return self.value == other  # type: ignore[attr-defined]
        return NotImplemented

    def _payload_enum_ne(self: Enum, other: object) -> Any:
        """Negate payload equality while preserving ``NotImplemented`` semantics."""
        if (eq := _payload_enum_eq(self, other)) is not NotImplemented:
            return not eq
        return NotImplemented

    def _payload_enum_str(self: Enum) -> str:
        """Render as the raw payload, matching ``int`` / ``str`` enum expectations."""
        return str(self.value)  # type: ignore[attr-defined]

    def _payload_enum_hash(self: Enum) -> int:
        """Hash like the raw payload so equality and hashing stay consistent."""
        return hash(self.value)  # type: ignore[attr-defined]

    _PAYLOAD_ENUM_DUNDERS = {
        "__eq__": _payload_enum_eq,
        "__ne__": _payload_enum_ne,
        "__str__": _payload_enum_str,
        "__hash__": _payload_enum_hash,
    }

    def _payload_enum_name(self: Enum) -> str:
        """Expose the declaration name without making plain ``Enum`` grow ``name``."""
        return str(self._name)

    def _payload_enum_repr(self: Enum) -> str:
        """Render as ``Class.member``, matching Python payload enum conventions."""
        return f"{type(self).__name__}.{self.name}"  # ty: ignore[unresolved-attribute]

    if "name" not in cls.__dict__:
        cls.name = property(_payload_enum_name)  # ty: ignore[unresolved-attribute]
    if not user_defined_repr:
        cls.__repr__ = _payload_enum_repr  # type: ignore[attr-defined]
    for name, method in _PAYLOAD_ENUM_DUNDERS.items():
        if name not in cls.__dict__:
            setattr(cls, name, method)


# ---------------------------------------------------------------------------
# Enum base + EnumAttrMap
# ---------------------------------------------------------------------------


class _EnumMeta(type(Object)):
    """Metaclass for :class:`Enum` that wires class-level iteration and length.

    Kept as a dedicated subclass of ``type(Object)`` so ``iter(MyEnum)`` and
    ``len(MyEnum)`` only resolve on :class:`Enum` subclasses, rather than
    mutating the shared ``Object`` metaclass globally.
    """

    def __iter__(cls) -> Iterator[Any]:
        return iter(_ordered_entries(cls))  # ty: ignore[unresolved-attribute]

    def __len__(cls) -> int:
        return len(_ordered_entries(cls))  # ty: ignore[unresolved-attribute]

    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
        """Construct enum variants from existing variants, payloads, or names.

        Normal object construction is still delegated to ``Object`` / dataclass
        initialization.  The single-argument path is intercepted to match
        stdlib enum lookup behavior:

        * ``Priority(Priority.low)`` returns the same variant.
        * ``Priority(10)`` returns the member whose payload value is ``10``.
        * ``Priority("low")`` returns the member declared with that name.

        Payload lookup is only enabled for classes with
        ``__ffi_enum_payload_value_type__`` so plain enums do not accidentally
        interpret arbitrary class-body values as public payloads.
        """
        if not kwargs and len(args) == 1:
            value = args[0]
            if isinstance(value, cls):
                return value
            payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None)
            if payload_value_type is not None and isinstance(value, payload_value_type):
                value_entries = _value_entries_dict(cls)
                if value_entries is not None and value in value_entries:
                    return value_entries[value]
            if isinstance(value, str):
                entries = _entries_dict(cls)
                if entries is not None and value in entries:
                    return entries[value]
        return super().__call__(*args, **kwargs)


@dataclass_transform(
    eq_default=False,
    order_default=False,
    frozen_default=True,
    field_specifiers=(Field, field, entry, auto),
)
@c_class("ffi.Enum", init=False)
class Enum(Object, metaclass=_EnumMeta):
    """A named-singleton registry with cross-language identity.

    Subclasses declare variants: frozen, named, ordinal-indexed
    singletons — the familiar enum pattern.  Unlike ``enum.Enum``, an
    ``Enum`` subclass bound to an FFI-registered ``type_key`` has an
    **open variant set**: C++ translation units and other Python modules
    binding the same ``type_key`` can contribute variants to the shared
    registry.  Per-variant metadata can also be attached post-hoc via
    :meth:`def_attr` as an *extensible attribute*, outside the class
    definition.

    For **closed, Python-only enums**, use a fresh ``type_key`` with
    :func:`auto` / :func:`entry` — behavior matches ``enum.Enum``.

    Attributes
    ----------
    _value : int
        Dense ordinal assigned at registration (0-indexed per class).
    _name : str
        The variant's string name key (e.g., ``"Add"`` for ``Op.Add``).

    Closed Python enum
    ------------------
    Pick a fresh ``type_key`` and list variants with :func:`auto` or
    :func:`entry`.  The variant set is fixed at class-definition time.

    .. code-block:: python

        class Priority(Enum, type_key="my.Priority"):
            low = auto()
            medium = auto()
            high = auto()


        # Variants with declared fields — values supplied via entry(...).
        class Activation(Enum, type_key="my.Activation"):
            output_zero: bool
            is_monotonic: bool

            relu: ClassVar[Activation] = entry(output_zero=True, is_monotonic=True)
            gelu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=False)

    Cross-language registry
    -----------------------
    When ``type_key`` is already registered (typically by C++), the
    Python class *binds* to the existing type rather than creating a
    new one.  Bare ``ClassVar[Cls]`` annotations bind to variants
    already registered on the C++ side; :func:`entry` / :func:`auto`
    still register fresh Python variants whose ordinals extend past the
    C++ ones.  All variants — regardless of origin — land in the same
    per-class registry and are visible to every binder of the same
    ``type_key``.

    .. code-block:: python

        # Registered in C++ via refl::EnumDef<VariantObj>("Alpha")... .
        class Variant(Enum, type_key="testing.TestEnumVariant"):
            Alpha: ClassVar[Variant]  # binds to C++-registered "Alpha"
            Beta: ClassVar[Variant]  # binds to C++-registered "Beta"

    Declaration forms
    -----------------
    Four shapes are supported in the class body:

    1. ``name = auto()`` — new variant with no declared fields.
    2. ``name: ClassVar[Cls] = entry(**kwargs)`` — new variant; ``kwargs``
       populate declared fields.
    3. ``name = entry(**kwargs)`` — same as (2), without the ``ClassVar``
       annotation.
    4. ``name: ClassVar[Cls]`` — in cross-language mode, binds to an
       existing C++-registered variant (error if unknown); otherwise
       registers a new Python variant with only the auto-assigned
       :attr:`_value` and :attr:`_name` (equivalent to ``name = auto()``).

    Integer literals (``ok = 0``) are rejected: :attr:`_value` is
    auto-assigned, so a user-supplied ordinal would either silently
    duplicate or conflict.  ``entry(_value=...)`` and ``entry(_name=...)``
    raise :class:`TypeError` at class-body time.

    Differences from ``enum.Enum``
    ------------------------------
    * **Same**: hidden identity fields :attr:`_name`, :attr:`_value`,
      iteration, identity
      comparison; closed-set behavior when ``type_key`` is fresh.
    * **Extended**: ``entry(**kwargs)`` replaces the tuple-RHS idiom;
      ``dataclass_transform`` gives native type-checker support; open
      registry when ``type_key`` is shared across languages.
    * **Different**: :attr:`_value` is always the ordinal; :meth:`def_attr`
      adds extensible attributes outside the class schema.  Use
      :class:`IntEnum` / :class:`StrEnum` when you need a user-visible
      ``value`` field.
    * **Not provided**: ``Flag`` / ``IntFlag``, member aliasing,
      ``_missing_`` hook.

    Subclasses inherit :meth:`get`, :meth:`all_entries`, :meth:`def_attr`,
    and the ``attr_dict`` class-level view.

    """

    __slots__ = ()

    _value: int
    _name: str
    __ffi_enum_payload_value_type__: object = None

    def __init_subclass__(
        cls,
        *,
        type_key: str | None = None,
        frozen: bool = True,
        init: bool = True,
        repr: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init_subclass__(**kwargs)
        if type_key is None:
            return
        user_defined_repr = "__repr__" in cls.__dict__

        binders, python_entries = _collect_entry_declarations(cls)

        cxx_backed = core._object_type_key_to_index(type_key) is not None
        if cxx_backed:
            c_class(type_key, init=init, repr=repr)(cls)
        else:
            py_class(type_key, frozen=frozen)(cls)

        if getattr(cls, "__ffi_enum_payload_value_type__", None) is not None:
            _install_payload_enum_behaviors(cls, user_defined_repr=user_defined_repr)

        _resolve_entries(cls, binders, python_entries, type_key=type_key, cxx_backed=cxx_backed)

    @classmethod
    def get(cls, name: str) -> Enum:
        """Return the variant named *name*, or raise :class:`KeyError`."""
        entries = _entries_dict(cls)
        if entries is not None and name in entries:
            return entries[name]
        raise KeyError(f"{cls.__name__} has no variant named {name!r}")

    @classmethod
    def all_entries(cls) -> Iterator[Enum]:
        """Iterate over all variants, in ordinal (``_value``) order."""
        return iter(_ordered_entries(cls))

    @_ClassProperty
    def attr_dict(cls: type) -> Any:
        """Live ``Dict[str, List[Any]]`` backing every extensible attribute.

        The outer dict is keyed by extensible-attribute name; each value
        is a list indexed by variant ordinal.  Prefer :meth:`def_attr`
        for normal per-variant reads and writes; this property is for
        bulk inspection and for reading values written by C++
        ``EnumDef::set_attr``.
        """
        return _attrs_dict(cls) or Dict({})

    @classmethod
    def def_attr(
        cls,
        name: str,
        *,
        default: Any = core.MISSING,
    ) -> EnumAttrMap:
        """Declare an *extensible attribute* column on this enum.

        Extensible attributes let any consumer associate per-variant data
        outside the enum's class-body schema — a lowering function
        attached to an operator by a code generator, a cost model
        registered only on some targets, a documentation string added
        after the fact.  Writes are last-write-wins for the same
        ``(variant, name)`` pair and visible to every consumer that
        calls :meth:`def_attr` with the same name on the same enum,
        including C++ code writing via ``EnumDef::set_attr``.

        Extensible attributes differ from **declared fields**:

        ====================  ========================  ==========================
        Concept               Lives on                  Added by
        ====================  ========================  ==========================
        Declared field        The variant object        Enum author, in class body
        Extensible attribute  ``__ffi_enum_attrs__``    Any consumer, any time
        ====================  ========================  ==========================

        Rule of thumb: if the data is part of *what a variant is*,
        declare a field in the class body; if it's part of *what a
        consumer wants to know*, attach it with :meth:`def_attr`.

        Parameters
        ----------
        name
            The extensible-attribute name (e.g., ``"has_side_effects"``).
            Writes go to ``attr_dict[name]`` as a list indexed by each
            variant's ordinal.
        default
            Value returned by ``attr[variant]`` when nothing was
            registered for that variant.  Left as ``MISSING`` to raise
            :class:`KeyError` on unset variants instead.  The default
            is a property of *this* :class:`EnumAttrMap` view, not of
            the underlying column: calling :meth:`def_attr` twice with
            the same ``name`` but different defaults creates two views
            that share every explicit write but may disagree on unset
            variants — e.g., ``Op.def_attr("cost", default=0)`` and
            ``Op.def_attr("cost", default=-1)`` return ``0`` and ``-1``
            respectively for a variant that was never written to.

        Returns
        -------
        EnumAttrMap
            Mutable view over the column.  Use ``variant in attr`` to
            distinguish an explicit write from a default-hit.  ``None``
            is reserved as the "unset" sentinel (matching C++
            ``EnumDef::set_attr`` padding), so ``attr[variant] = None``
            raises :class:`TypeError` — store a typed wrapper (e.g. a
            ``0``/``False`` flag) when you need a falsy-but-present
            value.

        Notes
        -----
        :meth:`def_attr` is not a way to add fields to the enum's
        schema, subclass frozen variants, or bypass the frozen-instance
        invariant via ``setattr`` — for that, declare a field in the
        class body instead.

        """
        return EnumAttrMap(cls, name, default=default)


class EnumAttrMap:
    """Mutable per-variant view over an extensible-attribute column.

    Returned by :meth:`Enum.def_attr`.  Writes go to a ``List[Any]``
    column keyed by extensible-attribute name inside the per-class
    ``__ffi_enum_attrs__`` dict.  The list is indexed by each variant's
    ordinal (``variant._value``) and padded with ``None`` as new variants
    are registered.  The column is shared across every consumer —
    including C++ code writing via ``EnumDef::set_attr`` — and the data
    is not a field on the variant object.  See :meth:`Enum.def_attr` for
    full semantics.

    ``None`` is reserved as the column's "unset" sentinel (matching the
    C++ ``Any(nullptr)`` padding used by ``EnumDef::set_attr``), so
    :meth:`__setitem__` rejects ``None`` with :class:`TypeError` — an
    explicit ``attr[variant] = None`` would otherwise be
    indistinguishable from never-written and surprise ``variant in attr``
    / :meth:`__getitem__` readers.  To "clear" a previously written
    value, register a mutable container once and mutate it in place.
    """

    __slots__ = ("_default", "_enum_cls", "_name")

    def __init__(self, enum_cls: type, name: str, *, default: Any = core.MISSING) -> None:
        self._enum_cls = enum_cls
        self._name = name
        self._default = default

    def _ordinal_of(self, variant: object) -> int:
        if not isinstance(variant, self._enum_cls):
            raise TypeError(
                f"{self._enum_cls.__name__}.def_attr({self._name!r}) expects a "
                f"{self._enum_cls.__name__} variant, got {type(variant).__name__}"
            )
        return int(variant._value)  # type: ignore[attr-defined]

    def _column(self, *, create: bool) -> Any | None:
        """Return the ``List[Any]`` column for this attribute; create if missing.

        Returns ``None`` when ``create`` is false and the column doesn't exist.
        """
        attrs = _attrs_dict(self._enum_cls) if not create else _ensure_attrs_dict(self._enum_cls)
        if attrs is None:
            return None
        if self._name in attrs:
            return attrs[self._name]
        if not create:
            return None
        col = List([])
        attrs[self._name] = col
        return col

    def __setitem__(self, variant: object, value: Any) -> None:
        if value is None:
            raise TypeError(
                f"{self._enum_cls.__name__}.def_attr({self._name!r}): "
                f"None is reserved as the 'unset' sentinel for extensible "
                f"attributes and cannot be written explicitly."
            )
        ordinal = self._ordinal_of(variant)
        col = self._column(create=True)
        assert col is not None  # create=True always materialises the column.
        while len(col) <= ordinal:
            col.append(None)
        col[ordinal] = value

    def __getitem__(self, variant: object) -> Any:
        ordinal = self._ordinal_of(variant)
        col = self._column(create=False)
        if col is not None and ordinal < len(col):
            v = col[ordinal]
            if v is not None:
                return v
        if self._default is core.MISSING:
            raise KeyError(
                f"{self._enum_cls.__name__}.{variant._name} has no "  # type: ignore[attr-defined]
                f"extensible attribute {self._name!r} set"
            )
        return self._default

    def __contains__(self, variant: object) -> bool:
        if not isinstance(variant, self._enum_cls):
            return False
        try:
            ordinal = self._ordinal_of(variant)
        except TypeError:
            return False
        col = self._column(create=False)
        return col is not None and ordinal < len(col) and col[ordinal] is not None

    def get(self, variant: object, default: Any = None) -> Any:
        """Return the value for *variant*, or *default* if unset or foreign."""
        if not isinstance(variant, self._enum_cls):
            return default
        try:
            return self[variant]
        except KeyError:
            return default

    @property
    def name(self) -> str:
        """The extensible-attribute name passed to :meth:`Enum.def_attr`."""
        return self._name


class IntEnum(Enum):
    """Enum variant base with a user-visible ``value: int`` field."""

    __slots__ = ()
    value: int

    def __init_subclass__(
        cls,
        *,
        type_key: str | None = None,
        frozen: bool = True,
        init: bool = True,
        repr: bool = True,
        **kwargs: Any,
    ) -> None:
        _prepare_payload_enum_subclass(cls, value_type=int, base_name="IntEnum")
        super().__init_subclass__(
            type_key=type_key,
            frozen=frozen,
            init=init,
            repr=repr,
            **kwargs,
        )


class StrEnum(Enum):
    """Enum variant base with a user-visible ``value: str`` field."""

    __slots__ = ()
    value: str

    def __init_subclass__(
        cls,
        *,
        type_key: str | None = None,
        frozen: bool = True,
        init: bool = True,
        repr: bool = True,
        **kwargs: Any,
    ) -> None:
        _prepare_payload_enum_subclass(cls, value_type=str, base_name="StrEnum")
        super().__init_subclass__(
            type_key=type_key,
            frozen=frozen,
            init=init,
            repr=repr,
            **kwargs,
        )


# ---------------------------------------------------------------------------
# TypeAttr accessors
# ---------------------------------------------------------------------------


def _entries_dict(cls: type) -> Any:
    type_info = getattr(cls, "__tvm_ffi_type_info__", None)
    if type_info is None:
        return None
    return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)


def _attrs_dict(cls: type) -> Any:
    type_info = getattr(cls, "__tvm_ffi_type_info__", None)
    if type_info is None:
        return None
    return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)


def _value_entries_dict(cls: type) -> Any:
    """Return the live ``Dict[Any, Enum]`` payload-to-variant index, or None."""
    type_info = getattr(cls, "__tvm_ffi_type_info__", None)
    if type_info is None:
        return None
    return core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR)


def _ordered_entries(cls: type) -> list[Any]:
    """Return all variants ordered by ordinal."""
    entries = _entries_dict(cls)
    if entries is None:
        return []
    ordered: list[Any] = [None] * len(entries)
    for inst in entries.values():
        ordered[int(inst._value)] = inst
    return ordered


def _ensure_entries_dict(cls: type) -> Any:
    """Return the live ``__ffi_enum_entries__`` dict, registering it if absent."""
    type_info = cls.__tvm_ffi_type_info__  # ty: ignore[unresolved-attribute]
    entries = core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)
    if entries is not None:
        return entries
    entries = Dict({})
    core._register_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR, entries)
    # Re-read so mutations go through the ref owned by the registry.
    return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)


def _ensure_attrs_dict(cls: type) -> Any:
    """Return the live ``__ffi_enum_attrs__`` dict, registering it if absent."""
    type_info = cls.__tvm_ffi_type_info__  # ty: ignore[unresolved-attribute]
    attrs = core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)
    if attrs is not None:
        return attrs
    attrs = Dict({})
    core._register_type_attr(type_info.type_index, ENUM_ATTRS_ATTR, attrs)
    return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)


def _ensure_value_entries_dict(cls: type) -> Any:
    """Return the live ``__ffi_enum_value_entries__`` dict, registering it if absent."""
    type_info = cls.__tvm_ffi_type_info__  # ty: ignore[unresolved-attribute]
    entries = core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR)
    if entries is not None:
        return entries
    entries = Dict({})
    core._register_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR, entries)
    return core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR)


# ---------------------------------------------------------------------------
# Class-body scanning + entry materialisation
# ---------------------------------------------------------------------------


def _collect_entry_declarations(
    cls: type,
) -> tuple[list[str], dict[str, _EnumEntry]]:
    """Scan *cls.__dict__* for variant declarations.

    Returns ``(binders, python_entries)`` in declaration order:

    * *binders* — names annotated as ``ClassVar[Cls]`` with no assigned value.
      Each either binds to an existing C++-registered entry with the same
      name or registers a new blank Python entry.
    * *python_entries* — names assigned an ``entry(...)`` sentinel (with or
      without a ``ClassVar`` annotation).  Each registers a new Python entry
      using the captured args/kwargs.

    Matched assignments are removed from ``cls.__dict__`` so that
    ``@c_class`` / ``@py_class`` don't misinterpret them as field defaults or
    class constants.
    """
    annotations = _own_annotations(cls)
    dict_keys = set(cls.__dict__.keys())

    binders: list[str] = []
    ordinary_fields: set[str] = set()
    payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None)
    for name, ann in annotations.items():
        if name.startswith("_"):
            continue
        if _is_class_var(ann):
            if name not in dict_keys:
                binders.append(name)
        else:
            ordinary_fields.add(name)

    python_entries: dict[str, _EnumEntry] = {}
    for name, value in list(cls.__dict__.items()):
        if name.startswith("_"):
            continue
        if isinstance(value, _EnumEntry):
            python_entries[name] = value
            try:
                delattr(cls, name)
            except AttributeError:
                pass
        elif payload_value_type is not None and name not in ordinary_fields:
            if isinstance(value, (staticmethod, classmethod, property)) or callable(value):
                continue
            python_entries[name] = _EnumEntry(value=value)
            try:
                delattr(cls, name)
            except AttributeError:
                pass

    return binders, python_entries


def _resolve_entries(
    cls: type,
    binders: list[str],
    python_entries: dict[str, _EnumEntry],
    *,
    type_key: str,
    cxx_backed: bool,
) -> None:
    """Materialise *binders* and *python_entries* into class-attribute singletons.

    Processing order matches declaration order: ``binders`` first (because
    their annotations appear before any class-body assignments), then
    ``python_entries`` in their class-body order.  Each newly registered
    entry gets a dense ordinal equal to the current entries-dict size, so
    ordinals stay compact and stable across registrations.

    A cxx-backed enum (``type_key`` was already registered in the FFI type
    system before this Python subclass was created) supports mixing C++ and
    Python entries: bare ``ClassVar[Cls]`` binders must name an existing
    C++-registered entry, but ``entry(...)``/``auto()`` sentinels may add
    fresh Python-side entries whose ordinals extend past the C++ entries.
    """
    entries = _ensure_entries_dict(cls)
    payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None)

    for name in binders:
        if name in entries:
            # Already materialised — either C++-registered or previously bound.
            instance = entries[name]
            setattr(cls, name, instance)
            _index_payload(cls, instance, payload_value_type)
            continue
        if cxx_backed:
            raise _cxx_backed_unknown_binder_error(cls, name, type_key, entries)
        ordinal = len(entries)
        instance = _instantiate(cls, args=(), kwargs={}, ordinal=ordinal, name=name)
        entries[name] = instance
        setattr(cls, name, instance)
        _index_payload(cls, instance, payload_value_type)

    for name, e in python_entries.items():
        if name in entries:
            raise RuntimeError(
                f"Duplicate enum entry {name!r} for {cls.__name__}: already "
                f"registered as ordinal {int(entries[name]._value)}."
            )
        if "_value" in e.kwargs or "_name" in e.kwargs:
            raise TypeError(
                f"{cls.__name__}.{name}: `_value` and `_name` are auto-assigned "
                f"and must not appear in entry(...) arguments."
            )
        ordinal = len(entries)
        instance = _instantiate_entry(
            cls, entry=e, ordinal=ordinal, name=name, cxx_backed=cxx_backed
        )
        entries[name] = instance
        setattr(cls, name, instance)
        _index_payload(cls, instance, payload_value_type)


def _index_payload(cls: type, instance: Any, payload_value_type: type | None) -> None:
    """Record ``(instance.value → instance)`` in the value-entries column.

    No-op for non-payload enums.  For payload enums, the first writer of a
    given payload wins — matches the "first-match" semantics of the linear
    scan that FFI converters fall back to when this column is incomplete.
    """
    if payload_value_type is None:
        return
    payload = getattr(instance, "value", None)
    if payload is None:
        return
    value_entries = _ensure_value_entries_dict(cls)
    if payload not in value_entries:
        value_entries[payload] = instance


def _cxx_backed_unknown_binder_error(
    cls: type,
    name: str,
    type_key: str,
    entries: Any,
) -> RuntimeError:
    """Build a descriptive error for an unbindable bare ``ClassVar`` binder.

    A bare ``ClassVar[Cls]`` annotation on a cxx-backed enum means "bind to
    an existing C++ entry with this name" — if the C++ registry has no such
    entry, the declaration is almost always a typo.  For adding a *new*
    Python-side variant on a cxx-backed enum, use ``entry(...)`` or
    ``auto()`` instead.
    """
    known = list(entries.keys()) if entries is not None else []
    known_str = ", ".join(repr(k) for k in known) if known else "<none>"
    return RuntimeError(
        f"Cannot bind enum variant {name!r} on {cls.__name__}: the FFI "
        f"type {type_key!r} is already registered in C++ with entries "
        f"[{known_str}], but has no C++ entry named {name!r}. "
        f"Bare ``ClassVar[{cls.__name__}]`` binders on a C++-backed enum "
        f"must name an entry already registered in C++; they cannot "
        f"introduce new variants from Python. "
        f"If this was a typo, double-check the spelling against the known "
        f"entries above (`{name}: ClassVar[{cls.__name__}]`); if you meant "
        f"to add a new Python-side variant, use `{name} = auto()` or "
        f"`{name} = entry(...)` instead."
    )


def _instantiate(
    cls: type,
    *,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    ordinal: int,
    name: str,
) -> Any:
    """Construct a subclass instance with auto-assigned ``_value``/``_name``."""
    merged = dict(kwargs)
    merged["_value"] = ordinal
    merged["_name"] = name
    return cls(*args, **merged)


def _instantiate_entry(
    cls: type,
    *,
    entry: _EnumEntry,
    ordinal: int,
    name: str,
    cxx_backed: bool,
) -> Any:
    """Instantiate an enum entry and normalize construction errors."""
    try:
        if cxx_backed:
            return _instantiate_cxx_backed(
                cls, args=entry.args, kwargs=entry.kwargs, ordinal=ordinal, name=name
            )
        return _instantiate(cls, args=entry.args, kwargs=entry.kwargs, ordinal=ordinal, name=name)
    except Exception as err:
        raise TypeError(f"{cls.__name__}.{name}: invalid enum entry: {err}") from None


def _instantiate_cxx_backed(
    cls: type,
    *,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    ordinal: int,
    name: str,
) -> Any:
    """Construct a new variant of a cxx-backed enum without going through ``__init__``.

    C++-backed enums whose underlying type is registered with
    ``refl::init(false)`` (e.g. any subclass of ``EnumObj`` in C++) have no
    ``__ffi_init__``, so the usual ``cls(_value=..., _name=...)`` path is not
    available.  Mirror ``reflection::EnumDef`` by allocating a blank instance
    via ``__ffi_new__`` and populating fields through the frozen-setter
    escape hatch exposed on the reflected property descriptors.
    """
    if args:
        raise TypeError(
            f"{cls.__name__}.{name}: positional `entry(...)` args are not "
            f"supported when extending a C++-backed enum; use keyword "
            f"arguments naming reflected fields."
        )
    type_info = cls.__tvm_ffi_type_info__  # ty: ignore[unresolved-attribute]
    ffi_new = core._lookup_type_attr(type_info.type_index, "__ffi_new__")
    if ffi_new is None:
        raise RuntimeError(
            f"Cannot add Python enum variant {name!r} on {cls.__name__}: "
            f"its C++ type has no ``__ffi_new__`` allocator registered, so "
            f"blank instances cannot be created from Python."
        )
    instance = ffi_new()
    for key in ("_value", "_name", *kwargs.keys()):
        descriptor = getattr(cls, key, None)
        if descriptor is None or not hasattr(descriptor, "set"):
            raise TypeError(
                f"{cls.__name__}.{name}: cannot set field {key!r} on a "
                f"C++-backed enum — no reflected setter is available."
            )
    getattr(cls, "_value").set(instance, ordinal)
    getattr(cls, "_name").set(instance, name)
    for k, v in kwargs.items():
        getattr(cls, k).set(instance, v)
    return instance


# ---------------------------------------------------------------------------
# Annotation helpers
# ---------------------------------------------------------------------------


def _own_annotations(cls: type) -> dict[str, Any]:
    """Return *cls*'s own annotations dict (not inherited)."""
    if sys.version_info >= (3, 14):
        return dict(getattr(cls, "__annotations__", {}) or {})
    return dict(cls.__dict__.get("__annotations__", {}))


def _is_class_var(annotation: Any) -> bool:
    """Return True if *annotation* is ``ClassVar`` or ``ClassVar[...]``."""
    if annotation is ClassVar:
        return True
    if typing.get_origin(annotation) is ClassVar:
        return True
    if isinstance(annotation, str):
        stripped = annotation.replace(" ", "")
        return stripped.startswith("ClassVar") or stripped.startswith("typing.ClassVar")
    return False


def _annotation_matches_expected(annotation: Any, expected_type: type) -> bool:
    """Return True if *annotation* matches the required payload-field type."""
    if annotation is expected_type:
        return True
    if isinstance(annotation, str):
        stripped = annotation.replace(" ", "")
        expected = expected_type.__name__
        return stripped in {expected, f"builtins.{expected}"}
    return False


def _prepare_payload_enum_subclass(
    cls: type[Enum], *, value_type: type[Any], base_name: str
) -> None:
    """Inject and validate the user-visible ``value`` field for payload enums."""
    if "value" in cls.__dict__:
        raise TypeError(
            f"{base_name} reserves `value` as a declared field; use "
            "`entry(value=...)` or `<variant> = <payload>` to assign each "
            "variant's payload."
        )

    annotations = _own_annotations(cls)
    existing = annotations.get("value")
    if existing is not None:
        if _is_class_var(existing):
            raise TypeError(f"{base_name} reserves `value` as a declared field, not a ClassVar.")
        if not _annotation_matches_expected(existing, value_type):
            raise TypeError(
                f"{base_name} requires `value: {value_type.__name__}` on subclasses, "
                f"got {existing!r}."
            )
        return

    updated = _own_annotations(cls)
    updated["value"] = value_type
    cls.__annotations__ = updated
    cls.__ffi_enum_payload_value_type__ = value_type
