# 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.
"""Dunder method installation for ``@c_class`` and ``@py_class``."""

from __future__ import annotations

import inspect
from typing import TYPE_CHECKING, Any, Callable

from . import core
from .core import TypeInfo, object_repr

if TYPE_CHECKING:
    from .core import Function


def _make_init(
    type_cls: type,
    type_info: TypeInfo,
    ffi_init: Function,
    py_class_mode: bool = False,
) -> Callable[..., None]:
    """Build ``__init__`` that delegates to ``__ffi_init__``.

    Both ``@c_class`` and ``@py_class`` use the same constructor-call path
    (``self.__init_handle_by_constructor__(ffi_init, *args)``).  The only
    difference is how ``super().__init__()`` from a subclass is handled:

    * **c_class** — raises ``TypeError`` (subclass must define its own init).
    * **py_class** — silently skips when the C++ handle is already set, so
      ``super().__init__()`` is a harmless no-op.

    Parameters
    ----------
    type_cls
        The class to build an __init__ for.
    type_info
        The TypeInfo for *type_cls*.
    ffi_init
        The C++ ``__ffi_init__`` resolved at install time.
    py_class_mode
        If True, use a ``chandle`` guard instead of a ``TypeError`` guard.

    """
    sig = _make_init_signature(type_info)
    kwargs_obj = core.KWARGS
    missing = core.MISSING
    has_post_init = hasattr(type_cls, "__post_init__")

    if py_class_mode:

        def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
            if self.__chandle__() != 0:
                return
            ffi_args: list[Any] = list(args)
            if kwargs:
                ffi_args.append(kwargs_obj)
                for key, val in kwargs.items():
                    if val is not missing:
                        ffi_args.append(key)
                        ffi_args.append(val)
            self.__init_handle_by_constructor__(ffi_init, *ffi_args)
            if has_post_init:
                self.__post_init__()

    else:
        type_name = type_cls.__name__

        def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
            if type_info is not type(self).__tvm_ffi_type_info__:
                raise TypeError(
                    f"Calling `super().__init__()` is not supported for @c_class "
                    f"`{type_name}`. Use @py_class for inheritance with manual "
                    f"__init__, or define `{type(self).__name__}` with init=True."
                )
            ffi_args: list[Any] = list(args)
            if kwargs:
                ffi_args.append(kwargs_obj)
                for key, val in kwargs.items():
                    if val is not missing:
                        ffi_args.append(key)
                        ffi_args.append(val)
            self.__init_handle_by_constructor__(ffi_init, *ffi_args)
            if has_post_init:
                self.__post_init__()

    __init__.__signature__ = sig  # ty: ignore[invalid-assignment]
    __init__.__qualname__ = f"{type_cls.__qualname__}.__init__"
    __init__.__module__ = type_cls.__module__
    return __init__


def _make_init_signature(type_info: TypeInfo) -> inspect.Signature:
    """Build an ``inspect.Signature`` from reflection field metadata.

    Walks the parent chain (parent-first) to collect all ``init=True`` fields,
    reorders required-before-optional within each group, and returns a
    Signature for introspection.
    """
    positional: list[tuple[str, bool]] = []  # (name, has_default)
    kw_only: list[tuple[str, bool]] = []  # (name, has_default)

    # Walk the parent chain to collect all fields (parent-first order).
    all_fields: list[Any] = []
    ti: TypeInfo | None = type_info
    chain: list[TypeInfo] = []
    while ti is not None:
        chain.append(ti)
        ti = ti.parent_type_info
    for ancestor_info in reversed(chain):
        all_fields.extend(ancestor_info.fields)

    for field in all_fields:
        if not field.c_init:
            continue
        if field.c_kw_only:
            kw_only.append((field.name, field.c_has_default))
        else:
            positional.append((field.name, field.c_has_default))

    # Required params must come before optional ones within each group.
    pos_required = [(n, d) for n, d in positional if not d]
    pos_default = [(n, d) for n, d in positional if d]
    kw_required = [(n, d) for n, d in kw_only if not d]
    kw_default = [(n, d) for n, d in kw_only if d]

    params: list[inspect.Parameter] = []
    params.append(inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD))

    for name, _has_default in pos_required:
        params.append(inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD))

    for name, _has_default in pos_default:
        params.append(
            inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD, default=core.MISSING)
        )

    for name, _has_default in kw_required:
        params.append(inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY))

    for name, _has_default in kw_default:
        params.append(inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY, default=core.MISSING))

    return inspect.Signature(params)


# ---------------------------------------------------------------------------
# Copy / deepcopy / replace factories
# ---------------------------------------------------------------------------


def _make_copy(type_info: TypeInfo, shallow_copy_fn: Function | None) -> Callable[..., Any]:
    """Build ``__copy__`` with ``__ffi_shallow_copy__`` resolved at install time."""
    if shallow_copy_fn is not None:

        def __copy__(self: Any) -> Any:
            return shallow_copy_fn(self)

    else:

        def __copy__(self: Any) -> Any:
            raise TypeError(
                f"Type `{type(self).__name__}` does not support copy. "
                f"The underlying C++ type is not copy-constructible."
            )

    return __copy__


def _make_deepcopy(_type_info: TypeInfo) -> Callable[..., Any]:
    """Build ``__deepcopy__`` with ``DeepCopy`` resolved at install time."""
    from . import _ffi_api  # noqa: PLC0415

    deep_copy = _ffi_api.DeepCopy

    def __deepcopy__(self: Any, memo: Any = None) -> Any:
        return deep_copy(self)

    return __deepcopy__


def _make_replace(_type_info: TypeInfo) -> Callable[..., Any]:
    """Build ``__replace__`` with ``copy.copy`` resolved at install time."""
    import copy  # noqa: PLC0415

    copy_copy = copy.copy

    def __replace__(self: Any, **kwargs: Any) -> Any:
        obj = copy_copy(self)
        cls = type(obj)
        for key, value in kwargs.items():
            getattr(cls, key).set(obj, value)
        return obj

    return __replace__


# ---------------------------------------------------------------------------
# __init__ installation
# ---------------------------------------------------------------------------


def _set_match_args(cls: type, type_info: TypeInfo) -> None:
    """Set ``cls.__match_args__`` from reflected fields.

    Mirrors stdlib :func:`dataclasses.dataclass` semantics: the tuple
    contains the names of positional ``__init__`` fields (``init=True``
    and ``kw_only=False``), walking the parent chain in parent-first
    order.  If ``cls`` already defines ``__match_args__`` in its own
    ``__dict__``, it is left untouched.
    """
    if "__match_args__" in cls.__dict__:
        return
    chain: list[TypeInfo] = []
    ti: TypeInfo | None = type_info
    while ti is not None:
        chain.append(ti)
        ti = ti.parent_type_info
    names: list[str] = []
    for ancestor in reversed(chain):
        for tf in ancestor.fields or ():
            df = tf.dataclass_field
            if df is None:
                continue
            if df.init and not df.kw_only:
                names.append(tf.name)
    setattr(cls, "__match_args__", tuple(names))


def _install_dataclass_dunders(  # noqa: PLR0912, PLR0915
    cls: type,
    *,
    init: bool,
    repr: bool,
    eq: bool,
    order: bool,
    unsafe_hash: bool,
    match_args: bool = True,
    py_class_mode: bool = False,
) -> None:
    """Install structural dunder methods on *cls*.

    Each dunder delegates to the corresponding C++ recursive structural
    operation (``RecursiveEq``, ``RecursiveHash``, ``RecursiveLt``, etc.).
    If the user already defined a dunder in the class body
    (i.e. it exists in ``cls.__dict__``), it is left untouched.

    Parameters
    ----------
    cls
        The class to install dunders on.  Must have been processed by
        :func:`register_object` first (so ``__tvm_ffi_type_info__`` exists).
    init
        If True, install ``__init__`` from C++ reflection metadata.
    repr
        If True, install :func:`~tvm_ffi.core.object_repr` as ``__repr__``.
    eq
        If True, install ``__eq__`` and ``__ne__`` using ``RecursiveEq``.
        Returns ``NotImplemented`` for unrelated types so Python can
        fall back to identity comparison.
    order
        If True, install ``__lt__``, ``__le__``, ``__gt__``, ``__ge__``
        using ``RecursiveLt``/``Le``/``Gt``/``Ge``.  Returns
        ``NotImplemented`` for unrelated types.
    unsafe_hash
        If True, install ``__hash__`` using ``RecursiveHash``.
    match_args
        If True (default), set ``cls.__match_args__`` to the tuple of
        positional ``__init__`` field names for use with ``match``
        statements.  Skipped when the class already defines
        ``__match_args__`` in its body.
    py_class_mode
        If True, use a ``chandle`` guard for ``__init__`` so that
        ``super().__init__()`` is a no-op, and wrap user-defined
        ``__init__`` to allocate via ``__ffi_init__`` before user code.

    """
    from . import _ffi_api  # noqa: PLC0415

    type_info: TypeInfo = cls.__tvm_ffi_type_info__  # type: ignore[assignment]
    type_index: int = type_info.type_index
    # Look up __ffi_init__ from TypeMethod (preferred) or TypeAttrColumn (fallback).
    ffi_init: Function | None = None
    for method in type_info.methods:
        if method.name == "__ffi_init__":
            ffi_init = method.func
            break
    if ffi_init is None:
        ffi_init = core._lookup_type_attr(type_index, "__ffi_init__")
    ffi_new: Function | None = core._lookup_type_attr(type_index, "__ffi_new__")
    ffi_shallow_copy: Function | None = core._lookup_type_attr(type_index, "__ffi_shallow_copy__")

    # __init__
    # ┌─────────────────────┬──────────────────────┬──────────────────────┐
    # │                     │ user __init__        │ no user __init__     │
    # ├─────────────────────┼──────────────────────┼──────────────────────┤
    # │ c_class, init=True  │ keep as-is           │ _make_init           │
    # │ c_class, init=False │ keep as-is           │ TypeError guard      │
    # │ py_class, init=True │ wrap chandle guard   │ _make_init(py_class) │
    # │ py_class, init=False│ wrap chandle guard   │ TypeError guard      │
    # └─────────────────────┴──────────────────────┴──────────────────────┘
    if "__init__" not in cls.__dict__:
        if init and ffi_init is not None:
            cls.__init__ = _make_init(  # type: ignore[attr-defined]
                cls,
                type_info,
                ffi_init=ffi_init,
                py_class_mode=py_class_mode,
            )
        elif not init:
            # init=False, no user __init__: TypeError guard
            msg = (
                f"`{cls.__name__}` cannot be constructed directly. "
                f"Define a custom __init__ or use a factory method."
            )

            def __init___(self: Any, *args: Any, **kwargs: Any) -> None:
                raise TypeError(msg)

            __init___.__qualname__ = f"{cls.__qualname__}.__init__"
            __init___.__module__ = cls.__module__
            cls.__init__ = __init___  # type: ignore[attr-defined]
    elif py_class_mode and ffi_new is not None:
        # User-defined __init__: wrap with chandle guard so the C++ object
        # is allocated (via __ffi_new__) before the user's code runs.
        # We use __ffi_new__ (zero-arg allocator) rather than __ffi_init__
        # because the user's __init__ signature may not match the field
        # layout.  super().__init__() from within the user's code becomes
        # a no-op because chandle is already set.
        import functools  # noqa: PLC0415

        user_init = cls.__dict__["__init__"]

        @functools.wraps(user_init)
        def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
            if self.__chandle__() == 0:
                self.__init_handle_by_constructor__(ffi_new)
            user_init(self, *args, **kwargs)

        cls.__init__ = __init__  # type: ignore[attr-defined]

    # __repr__
    if repr and "__repr__" not in cls.__dict__:
        cls.__repr__ = object_repr  # type: ignore[attr-defined]

    def _is_comparable(self: Any, other: Any) -> bool:
        """Return True if *self* and *other* share a type hierarchy."""
        return isinstance(other, type(self)) or isinstance(self, type(other))

    # Semantic families (__eq__/__ne__, __lt__/__le__/__gt__/__ge__) are
    # treated as a unit: if the user defines any member, the whole family
    # is skipped so generated and user-defined methods don't disagree.
    if eq and not ({"__eq__", "__ne__"} & set(cls.__dict__)):
        recursive_eq = _ffi_api.RecursiveEq

        def __eq__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return recursive_eq(self, other)

        def __ne__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return not recursive_eq(self, other)

        cls.__eq__ = __eq__  # type: ignore[attr-defined]
        cls.__ne__ = __ne__  # type: ignore[attr-defined]

    if order and not ({"__lt__", "__le__", "__gt__", "__ge__"} & set(cls.__dict__)):
        recursive_lt = _ffi_api.RecursiveLt
        recursive_le = _ffi_api.RecursiveLe
        recursive_gt = _ffi_api.RecursiveGt
        recursive_ge = _ffi_api.RecursiveGe

        def __lt__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return recursive_lt(self, other)

        def __le__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return recursive_le(self, other)

        def __gt__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return recursive_gt(self, other)

        def __ge__(self: Any, other: Any) -> bool:
            if not _is_comparable(self, other):
                return NotImplemented
            return recursive_ge(self, other)

        cls.__lt__ = __lt__  # type: ignore[attr-defined]
        cls.__le__ = __le__  # type: ignore[attr-defined]
        cls.__gt__ = __gt__  # type: ignore[attr-defined]
        cls.__ge__ = __ge__  # type: ignore[attr-defined]

    if unsafe_hash and "__hash__" not in cls.__dict__:
        recursive_hash = _ffi_api.RecursiveHash

        def __hash__(self: Any) -> int:
            return recursive_hash(self)

        cls.__hash__ = __hash__  # type: ignore[attr-defined]

    # __copy__ / __deepcopy__ / __replace__
    if "__copy__" not in cls.__dict__:
        cls.__copy__ = _make_copy(type_info, ffi_shallow_copy)  # type: ignore[attr-defined]
    if "__deepcopy__" not in cls.__dict__:
        cls.__deepcopy__ = _make_deepcopy(type_info)  # type: ignore[attr-defined]
    if "__replace__" not in cls.__dict__:
        cls.__replace__ = _make_replace(type_info)  # type: ignore[attr-defined]

    if match_args:
        _set_match_args(cls, type_info)
