# 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.

"""Provides the user-defined functions for evaluation of dataframes."""

from __future__ import annotations

import functools
import sys
from abc import ABCMeta, abstractmethod
from enum import Enum
from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, TypeVar, cast, overload

import pyarrow as pa

import datafusion._internal as df_internal
from datafusion import SessionContext
from datafusion.expr import Expr

# Imported at runtime so ``typing.get_type_hints`` resolves the capsule
# overloads; typing_extensions is a runtime dependency below 3.13.
if sys.version_info >= (3, 13):
    from types import CapsuleType as _PyCapsule
else:
    from typing_extensions import CapsuleType as _PyCapsule

if TYPE_CHECKING:
    _R = TypeVar("_R", bound=pa.Array)
    from collections.abc import Callable, Sequence


class Volatility(Enum):
    """Defines how stable or volatile a function is.

    When setting the volatility of a function, you can either pass this
    enumeration or a ``str``. The ``str`` equivalent is the lower case value of the
    name (`"immutable"`, `"stable"`, or `"volatile"`).
    """

    Immutable = 1
    """An immutable function will always return the same output when given the
    same input.

    DataFusion will attempt to inline immutable functions during planning.
    """

    Stable = 2
    """
    Returns the same value for a given input within a single queries.

    A stable function may return different values given the same input across
    different queries but must return the same value for a given input within a
    query. An example of this is the ``Now`` function. DataFusion will attempt to
    inline ``Stable`` functions during planning, when possible. For query
    ``select col1, now() from t1``, it might take a while to execute but ``now()``
    column will be the same for each output row, which is evaluated during
    planning.
    """

    Volatile = 3
    """A volatile function may change the return value from evaluation to
    evaluation.

    Multiple invocations of a volatile function may return different results
    when used in the same query. An example of this is the random() function.
    DataFusion can not evaluate such functions during planning. In the query
    ``select col1, random() from t1``, ``random()`` function will be evaluated
    for each output row, resulting in a unique random value for each row.
    """

    def __str__(self) -> str:
        """Returns the string equivalent."""
        return self.name.lower()


def data_type_or_field_to_field(value: pa.DataType | pa.Field, name: str) -> pa.Field:
    """Helper function to return a Field from either a Field or DataType."""
    if isinstance(value, pa.Field):
        return value
    return pa.field(name, type=value)


def data_types_or_fields_to_field_list(
    inputs: Sequence[pa.Field | pa.DataType] | pa.Field | pa.DataType,
) -> list[pa.Field]:
    """Helper function to return a list of Fields."""
    if isinstance(inputs, pa.DataType):
        return [pa.field("value", type=inputs)]
    if isinstance(inputs, pa.Field):
        return [inputs]

    return [
        data_type_or_field_to_field(v, f"value_{idx}") for (idx, v) in enumerate(inputs)
    ]


class ScalarUDFExportable(Protocol):
    """Type hint for object that has __datafusion_scalar_udf__ PyCapsule."""

    def __datafusion_scalar_udf__(self) -> object: ...  # noqa: D105


def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]:
    """Return ``True`` when ``value`` is a CPython ``PyCapsule``."""
    return value.__class__.__name__ == "PyCapsule"


class LogicalExtensionCodecExportable(Protocol):
    """Type hint for objects exposing ``__datafusion_logical_extension_codec__``.

    Implement the codec itself exactly as you would for a session that installs
    only yours. A session may hold several codecs, but each payload records the
    codec that wrote it and is only ever handed back to that codec, so there is
    no need to recognise or reject another library's payloads.

    A serialized plan records which codec wrote each payload, as a short id
    taken from your class's module and qualified name. An optional
    ``__datafusion_codec_id__`` attribute pins that id instead. It is not part
    of this protocol and is rarely needed: declare it when renaming your class
    must not stop older plans from decoding, or when one library installs two
    instances that own disjoint slices of the wire format.

    Args:
        session: A handle on the session this codec is being installed on. Take
            the task context provider from it rather than building a session of
            your own, so the decode callbacks resolve names against the session
            that runs the query.

            **Duck-type it; do not check its type.** This argument is the
            reference for every capsule getter in the protocol, and across them
            it is not always a session at all: the provider and catalog getters
            receive the host's logical codec as a bare ``PyCapsule``. Even when
            it *is* a session it is the PyO3 context from
            ``datafusion._internal``, not the
            :py:class:`~datafusion.context.SessionContext` wrapper, so it
            exposes every capsule getter and ``__datafusion_codec_id__`` —
            everything the protocol asks of it — but
            ``isinstance(session, SessionContext)`` is ``False`` even though its
            ``repr`` reads ``datafusion.SessionContext``. Pass it to
            ``ffi_logical_codec_from_pycapsule``, which handles both.

            The two bundle hooks are the exception:
            :py:class:`~datafusion.extensions.SessionComponentsExportable` and
            :py:class:`~datafusion.extensions.SessionPlannerExportable` are
            dispatched from Python and receive the wrapper. See
            :ref:`extension_getter_argument`.

    Returns:
        A ``datafusion_logical_extension_codec`` PyCapsule.

    Examples:
        A :py:class:`~datafusion.context.SessionContext` satisfies this
        protocol, which is how you obtain a real capsule to test against:

        >>> from datafusion import SessionContext
        >>> ctx = SessionContext()
        >>> type(ctx.__datafusion_logical_extension_codec__(ctx)).__name__
        'PyCapsule'
    """

    def __datafusion_logical_extension_codec__(  # noqa: D105
        self, session: Any
    ) -> object: ...


class PhysicalExtensionCodecExportable(Protocol):
    """Type hint for objects exposing ``__datafusion_physical_extension_codec__``.

    See :py:class:`LogicalExtensionCodecExportable` for ``session``, for why a
    codec need not recognise other libraries' payloads, and for
    ``__datafusion_codec_id__``.
    """

    def __datafusion_physical_extension_codec__(  # noqa: D105
        self, session: Any
    ) -> object: ...


class ScalarUDF:
    """Class for performing scalar user-defined functions (UDF).

    Scalar UDFs operate on a row by row basis. See also :py:class:`AggregateUDF` for
    operating on a group of rows.
    """

    def __init__(
        self,
        name: str,
        func: Callable[..., _R],
        input_fields: list[pa.Field],
        return_field: pa.Field,
        volatility: Volatility | str,
    ) -> None:
        """Instantiate a scalar user-defined function (UDF).

        See helper method :py:func:`udf` for argument details.
        """
        if hasattr(func, "__datafusion_scalar_udf__"):
            self._udf = df_internal.ScalarUDF.from_pycapsule(func)
            return
        if isinstance(input_fields, pa.DataType):
            input_fields = [input_fields]
        self._udf = df_internal.ScalarUDF(
            name, func, input_fields, return_field, str(volatility)
        )

    @classmethod
    def _from_internal(cls, internal: df_internal.ScalarUDF) -> ScalarUDF:
        """Wrap an already-constructed internal ``ScalarUDF`` handle.

        Used by :py:meth:`SessionContext.udf` and :py:meth:`from_pycapsule`
        to wrap a handle from the session's function registry or an FFI
        capsule without re-running :py:meth:`__init__`.
        """
        wrapper = cls.__new__(cls)
        wrapper._udf = internal
        return wrapper

    @property
    def name(self) -> str:
        """Return the registered name of this UDF.

        For UDFs imported via the FFI capsule protocol, this is the
        name the capsule itself reports — not the ``name`` argument
        passed to the constructor (which is ignored on the FFI path).

        Examples:
            >>> import pyarrow as pa
            >>> from datafusion import udf
            >>> double = udf(
            ...     lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
            ...     [pa.int64()],
            ...     pa.int64(),
            ...     volatility="immutable",
            ...     name="double",
            ... )
            >>> double.name
            'double'
        """
        return self._udf.name

    def __repr__(self) -> str:
        """Print a string representation of the Scalar UDF."""
        return self._udf.__repr__()

    def __call__(self, *args: Expr) -> Expr:
        """Execute the UDF.

        This function is not typically called by an end user. These calls will
        occur during the evaluation of the dataframe.
        """
        args_raw = [arg.expr for arg in args]
        return Expr(self._udf.__call__(*args_raw))

    @overload
    @staticmethod
    def udf(
        input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
        return_field: pa.DataType | pa.Field,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> Callable[..., ScalarUDF]: ...

    @overload
    @staticmethod
    def udf(
        func: Callable[..., _R],
        input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
        return_field: pa.DataType | pa.Field,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> ScalarUDF: ...

    @overload
    @staticmethod
    def udf(func: ScalarUDFExportable) -> ScalarUDF: ...

    @overload
    @staticmethod
    def udf(func: _PyCapsule) -> ScalarUDF: ...

    @staticmethod
    def udf(*args: Any, **kwargs: Any):  # noqa: D417, C901
        """Create a new User-Defined Function (UDF).

        This class can be used both as either a function or a decorator.

        Usage:
            - As a function: ``udf(func, input_fields, return_field,
              volatility, name)``.
            - As a decorator: ``@udf(input_fields, return_field, volatility, name)``.
              When used a decorator, do **not** pass ``func`` explicitly.

        In lieu of passing a PyArrow Field, you can pass a DataType for simplicity.
        When you do so, it will be assumed that the nullability of the inputs and
        output are True and that they have no metadata.

        Args:
            func (Callable, optional): Only needed when calling as a function.
                Skip this argument when using `udf` as a decorator. If you have a Rust
                backed ScalarUDF within a PyCapsule, you can pass this parameter
                and ignore the rest. They will be determined directly from the
                underlying function. See the online documentation for more information.
            input_fields (list[pa.Field | pa.DataType]): The data types or Fields
                of the arguments to ``func``. This list must be of the same length
                as the number of arguments.
            return_field (pa.DataType | pa.Field): The field of the return value
                from the function.
            volatility (Volatility | str): See `Volatility` for allowed values.
            name (Optional[str]): A descriptive name for the function.

        Returns:
            A user-defined function that can be used in SQL expressions,
            data aggregation, or window function calls.

        Examples:
            Using ``udf`` as a function:

            >>> import pyarrow.compute as pc
            >>> from datafusion.user_defined import ScalarUDF
            >>> def double_func(x):
            ...     return pc.multiply(x, 2)
            >>> double_udf = ScalarUDF.udf(
            ...     double_func, [pa.int64()], pa.int64(),
            ...     "volatile", "double_it")

            Using ``udf`` as a decorator:

            >>> @ScalarUDF.udf([pa.int64()], pa.int64(), "volatile")
            ... def decorator_double_udf(x):
            ...     return pc.multiply(x, 3)

            Apply to a dataframe:

            >>> ctx = dfn.SessionContext()
            >>> df = ctx.from_pydict({"x": [1, 2, 3]})
            >>> df.select(double_udf(col("x")).alias("result")).to_pydict()
            {'result': [2, 4, 6]}
            >>> df.select(decorator_double_udf(col("x")).alias("result")).to_pydict()
            {'result': [3, 6, 9]}
        """

        def _function(
            func: Callable[..., _R],
            input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
            return_field: pa.DataType | pa.Field,
            volatility: Volatility | str,
            name: str | None = None,
        ) -> ScalarUDF:
            if not callable(func):
                msg = "`func` argument must be callable"
                raise TypeError(msg)
            if name is None:
                if hasattr(func, "__qualname__"):
                    name = func.__qualname__.lower()
                else:
                    name = func.__class__.__name__.lower()
            input_fields = data_types_or_fields_to_field_list(input_fields)
            return_field = data_type_or_field_to_field(return_field, "value")
            return ScalarUDF(
                name=name,
                func=func,
                input_fields=input_fields,
                return_field=return_field,
                volatility=volatility,
            )

        def _decorator(
            input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
            return_field: pa.DataType | pa.Field,
            volatility: Volatility | str,
            name: str | None = None,
        ) -> Callable:
            def decorator(func: Callable) -> Callable:
                udf_caller = ScalarUDF.udf(
                    func, input_fields, return_field, volatility, name
                )

                @functools.wraps(func)
                def wrapper(*args: Any, **kwargs: Any) -> Callable:
                    return udf_caller(*args, **kwargs)

                return wrapper

            return decorator

        if not args and "func" in kwargs:
            args = (kwargs.pop("func"),)
        if args and (
            hasattr(args[0], "__datafusion_scalar_udf__") or _is_pycapsule(args[0])
        ):
            return ScalarUDF.from_pycapsule(args[0])

        if args and callable(args[0]):
            # Case 1: Used as a function, require the first parameter to be callable
            return _function(*args, **kwargs)
        # Case 2: Used as a decorator with parameters
        return _decorator(*args, **kwargs)

    @staticmethod
    def from_pycapsule(func: ScalarUDFExportable | _PyCapsule) -> ScalarUDF:
        """Create a Scalar UDF from ScalarUDF PyCapsule object.

        This function will instantiate a Scalar UDF that uses a DataFusion
        ScalarUDF that is exported via the FFI bindings.
        """
        if _is_pycapsule(func):
            return ScalarUDF._from_internal(df_internal.ScalarUDF.from_pycapsule(func))

        func = cast("ScalarUDFExportable", func)
        name = str(func.__class__)
        return ScalarUDF(
            name=name,
            func=func,
            input_fields=None,
            return_field=None,
            volatility=None,
        )


class Accumulator(metaclass=ABCMeta):
    """Defines how an :py:class:`AggregateUDF` accumulates values."""

    @abstractmethod
    def state(self) -> list[pa.Scalar]:
        """Return the current state.

        While this function template expects PyArrow Scalar values return type,
        you can return any value that can be converted into a Scalar. This
        includes basic Python data types such as integers and strings. In
        addition to primitive types, we currently support PyArrow, nanoarrow,
        and arro3 objects in addition to primitive data types. Other objects
        that support the Arrow FFI standard will be given a "best attempt" at
        conversion to scalar objects.
        """

    @abstractmethod
    def update(self, *values: pa.Array) -> None:
        """Evaluate an array of values and update state."""

    @abstractmethod
    def merge(self, states: list[pa.Array]) -> None:
        """Merge a set of states."""

    @abstractmethod
    def evaluate(self) -> pa.Scalar:
        """Return the resultant value.

        While this function template expects a PyArrow Scalar value return type,
        you can return any value that can be converted into a Scalar. This
        includes basic Python data types such as integers and strings. In
        addition to primitive types, we currently support PyArrow, nanoarrow,
        and arro3 objects in addition to primitive data types. Other objects
        that support the Arrow FFI standard will be given a "best attempt" at
        conversion to scalar objects.
        """


class AggregateUDFExportable(Protocol):
    """Type hint for object that has __datafusion_aggregate_udf__ PyCapsule."""

    def __datafusion_aggregate_udf__(self) -> object: ...  # noqa: D105


class AggregateUDF:
    """Class for performing scalar user-defined functions (UDF).

    Aggregate UDFs operate on a group of rows and return a single value. See
    also :py:class:`ScalarUDF` for operating on a row by row basis.
    """

    @overload
    def __init__(
        self,
        name: str,
        accumulator: Callable[[], Accumulator],
        input_types: list[pa.DataType],
        return_type: pa.DataType,
        state_type: list[pa.DataType],
        volatility: Volatility | str,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        accumulator: AggregateUDFExportable,
        input_types: None = ...,
        return_type: None = ...,
        state_type: None = ...,
        volatility: None = ...,
    ) -> None: ...

    def __init__(
        self,
        name: str,
        accumulator: Callable[[], Accumulator] | AggregateUDFExportable,
        input_types: list[pa.DataType] | None,
        return_type: pa.DataType | None,
        state_type: list[pa.DataType] | None,
        volatility: Volatility | str | None,
    ) -> None:
        """Instantiate a user-defined aggregate function (UDAF).

        See :py:func:`udaf` for a convenience function and argument
        descriptions.
        """
        if hasattr(accumulator, "__datafusion_aggregate_udf__"):
            self._udaf = df_internal.AggregateUDF.from_pycapsule(accumulator)
            return
        if (
            input_types is None
            or return_type is None
            or state_type is None
            or volatility is None
        ):
            msg = (
                "`input_types`, `return_type`, `state_type`, and `volatility` "
                "must be provided when `accumulator` is callable."
            )
            raise TypeError(msg)

        self._udaf = df_internal.AggregateUDF(
            name,
            accumulator,
            input_types,
            return_type,
            state_type,
            str(volatility),
        )

    @classmethod
    def _from_internal(cls, internal: df_internal.AggregateUDF) -> AggregateUDF:
        """Wrap an already-constructed internal ``AggregateUDF`` handle.

        Used by :py:meth:`SessionContext.udaf` and :py:meth:`from_pycapsule`
        to wrap a handle from the session's function registry or an FFI
        capsule without re-running :py:meth:`__init__`.
        """
        wrapper = cls.__new__(cls)
        wrapper._udaf = internal
        return wrapper

    @property
    def name(self) -> str:
        """Return the registered name of this UDAF.

        For UDAFs imported via the FFI capsule protocol, this is the
        name the capsule itself reports — not the ``name`` argument
        passed to the constructor (which is ignored on the FFI path).
        """
        return self._udaf.name

    def __repr__(self) -> str:
        """Print a string representation of the Aggregate UDF."""
        return self._udaf.__repr__()

    def __call__(self, *args: Expr) -> Expr:
        """Execute the UDAF.

        This function is not typically called by an end user. These calls will
        occur during the evaluation of the dataframe.
        """
        args_raw = [arg.expr for arg in args]
        return Expr(self._udaf.__call__(*args_raw))

    @overload
    @staticmethod
    def udaf(
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        state_type: list[pa.DataType],
        volatility: Volatility | str,
        name: str | None = None,
    ) -> Callable[..., AggregateUDF]: ...

    @overload
    @staticmethod
    def udaf(
        accum: Callable[[], Accumulator],
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        state_type: list[pa.DataType],
        volatility: Volatility | str,
        name: str | None = None,
    ) -> AggregateUDF: ...

    @overload
    @staticmethod
    def udaf(accum: AggregateUDFExportable) -> AggregateUDF: ...

    @overload
    @staticmethod
    def udaf(accum: _PyCapsule) -> AggregateUDF: ...

    @staticmethod
    def udaf(*args: Any, **kwargs: Any):  # noqa: D417, C901
        """Create a new User-Defined Aggregate Function (UDAF).

        This class allows you to define an aggregate function that can be used in
        data aggregation or window function calls.

        Usage:
            - As a function: ``udaf(accum, input_types, return_type, state_type, volatility, name)``.
            - As a decorator: ``@udaf(input_types, return_type, state_type, volatility, name)``.
              When using ``udaf`` as a decorator, do not pass ``accum`` explicitly.

        If your :py:class:`Accumulator` can be instantiated with no arguments, you
        can simply pass its type as ``accum``. If you need to pass additional
        arguments to its constructor, you can define a lambda or a factory method.
        During runtime the :py:class:`Accumulator` will be constructed for every
        instance in which this UDAF is used.

        Examples:
            >>> import pyarrow.compute as pc
            >>> from datafusion.user_defined import AggregateUDF, Accumulator, udaf
            >>> class Summarize(Accumulator):
            ...     def __init__(self, bias: float = 0.0):
            ...         self._sum = pa.scalar(bias)
            ...     def state(self):
            ...         return [self._sum]
            ...     def update(self, values):
            ...         self._sum = pa.scalar(
            ...             self._sum.as_py() + pc.sum(values).as_py())
            ...     def merge(self, states):
            ...         self._sum = pa.scalar(
            ...             self._sum.as_py() + pc.sum(states[0]).as_py())
            ...     def evaluate(self):
            ...         return self._sum

            Using ``udaf`` as a function:

            >>> udaf1 = AggregateUDF.udaf(
            ...     Summarize, pa.float64(), pa.float64(),
            ...     [pa.float64()], "immutable")

            Wrapping ``udaf`` with a function:

            >>> def sum_bias_10() -> Summarize:
            ...     return Summarize(10.0)
            >>> udaf2 = udaf(sum_bias_10, pa.float64(), pa.float64(), [pa.float64()],
            ...     "immutable")

            Using ``udaf`` with lambda:

            >>> udaf3 = udaf(lambda: Summarize(20.0), pa.float64(), pa.float64(),
            ...     [pa.float64()], "immutable")

            Using ``udaf`` as a decorator:

            >>> @AggregateUDF.udaf(
            ...     pa.float64(), pa.float64(),
            ...     [pa.float64()], "immutable")
            ... def udaf4():
            ...     return Summarize(10.0)

            Apply to a dataframe:

            >>> ctx = dfn.SessionContext()
            >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})
            >>> df.aggregate([], [udaf1(col("a")).alias("total")]).collect_column(
            ...     "total")[0].as_py()
            6.0
            >>> df.aggregate([], [udaf2(col("a")).alias("total")]).collect_column(
            ...     "total")[0].as_py()
            16.0
            >>> df.aggregate([], [udaf3(col("a")).alias("total")]).collect_column(
            ...     "total")[0].as_py()
            26.0
            >>> df.aggregate([], [udaf4(col("a")).alias("total")]).collect_column(
            ...     "total")[0].as_py()
            16.0

        Args:
            accum: The accumulator python function. Only needed when calling as a
                function. Skip this argument when using ``udaf`` as a decorator.
                If you have a Rust backed AggregateUDF within a PyCapsule, you can
                pass this parameter and ignore the rest. They will be determined
                directly from the underlying function. See the online documentation
                for more information.
            input_types: The data types of the arguments to ``accum``.
            return_type: The data type of the return value.
            state_type: The data types of the intermediate accumulation.
            volatility: See :py:class:`Volatility` for allowed values.
            name: A descriptive name for the function.

        Returns:
            A user-defined aggregate function, which can be used in either data
            aggregation or window function calls.
        """  # noqa: E501 W505

        def _function(
            accum: Callable[[], Accumulator],
            input_types: pa.DataType | list[pa.DataType],
            return_type: pa.DataType,
            state_type: list[pa.DataType],
            volatility: Volatility | str,
            name: str | None = None,
        ) -> AggregateUDF:
            if not callable(accum):
                msg = "`func` must be callable."
                raise TypeError(msg)
            if not isinstance(accum(), Accumulator):
                msg = "Accumulator must implement the abstract base class Accumulator"
                raise TypeError(msg)
            if name is None:
                name = accum().__class__.__qualname__.lower()
            if isinstance(input_types, pa.DataType):
                input_types = [input_types]
            return AggregateUDF(
                name=name,
                accumulator=accum,
                input_types=input_types,
                return_type=return_type,
                state_type=state_type,
                volatility=volatility,
            )

        def _decorator(
            input_types: pa.DataType | list[pa.DataType],
            return_type: pa.DataType,
            state_type: list[pa.DataType],
            volatility: Volatility | str,
            name: str | None = None,
        ) -> Callable[..., Callable[..., Expr]]:
            def decorator(accum: Callable[[], Accumulator]) -> Callable[..., Expr]:
                udaf_caller = AggregateUDF.udaf(
                    accum, input_types, return_type, state_type, volatility, name
                )

                @functools.wraps(accum)
                def wrapper(*args: Any, **kwargs: Any) -> Expr:
                    return udaf_caller(*args, **kwargs)

                return wrapper

            return decorator

        if not args and "accum" in kwargs:
            args = (kwargs.pop("accum"),)
        if args and (
            hasattr(args[0], "__datafusion_aggregate_udf__") or _is_pycapsule(args[0])
        ):
            return AggregateUDF.from_pycapsule(args[0])

        if args and callable(args[0]):
            # Case 1: Used as a function, require the first parameter to be callable
            return _function(*args, **kwargs)
        # Case 2: Used as a decorator with parameters
        return _decorator(*args, **kwargs)

    @staticmethod
    def from_pycapsule(func: AggregateUDFExportable | _PyCapsule) -> AggregateUDF:
        """Create an Aggregate UDF from AggregateUDF PyCapsule object.

        This function will instantiate a Aggregate UDF that uses a DataFusion
        AggregateUDF that is exported via the FFI bindings.
        """
        if _is_pycapsule(func):
            return AggregateUDF._from_internal(
                df_internal.AggregateUDF.from_pycapsule(func)
            )

        capsule = cast("AggregateUDFExportable", func)
        name = str(capsule.__class__)
        return AggregateUDF(
            name=name,
            accumulator=capsule,
            input_types=None,
            return_type=None,
            state_type=None,
            volatility=None,
        )


class WindowEvaluator:
    """Evaluator class for user-defined window functions (UDWF).

    It is up to the user to decide which evaluate function is appropriate.

    +------------------------+--------------------------------+------------------+---------------------------+
    | ``uses_window_frame``  | ``supports_bounded_execution`` | ``include_rank`` | function_to_implement     |
    +========================+================================+==================+===========================+
    | False (default)        | False (default)                | False (default)  | ``evaluate_all``          |
    +------------------------+--------------------------------+------------------+---------------------------+
    | False                  | True                           | False            | ``evaluate``              |
    +------------------------+--------------------------------+------------------+---------------------------+
    | False                  | True/False                     | True             | ``evaluate_all_with_rank``|
    +------------------------+--------------------------------+------------------+---------------------------+
    | True                   | True/False                     | True/False       | ``evaluate``              |
    +------------------------+--------------------------------+------------------+---------------------------+
    """  # noqa: W505, E501

    def memoize(self) -> None:
        """Perform a memoize operation to improve performance.

        When the window frame has a fixed beginning (e.g UNBOUNDED
        PRECEDING), some functions such as FIRST_VALUE and
        NTH_VALUE do not need the (unbounded) input once they have
        seen a certain amount of input.

        `memoize` is called after each input batch is processed, and
        such functions can save whatever they need
        """

    def get_range(self, idx: int, num_rows: int) -> tuple[int, int]:  # noqa: ARG002
        """Return the range for the window function.

        If `uses_window_frame` flag is `false`. This method is used to
        calculate required range for the window function during
        stateful execution.

        Generally there is no required range, hence by default this
        returns smallest range(current row). e.g seeing current row is
        enough to calculate window result (such as row_number, rank,
        etc)

        Args:
            idx:: Current index
            num_rows: Number of rows.
        """
        return (idx, idx + 1)

    def is_causal(self) -> bool:
        """Get whether evaluator needs future data for its result."""
        return False

    def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
        """Evaluate a window function on an entire input partition.

        This function is called once per input *partition* for window functions that
        *do not use* values from the window frame, such as
        :py:func:`~datafusion.functions.row_number`,
        :py:func:`~datafusion.functions.rank`,
        :py:func:`~datafusion.functions.dense_rank`,
        :py:func:`~datafusion.functions.percent_rank`,
        :py:func:`~datafusion.functions.cume_dist`,
        :py:func:`~datafusion.functions.lead`,
        and :py:func:`~datafusion.functions.lag`.

        It produces the result of all rows in a single pass. It
        expects to receive the entire partition as the ``value`` and
        must produce an output column with one output row for every
        input row.

        ``num_rows`` is required to correctly compute the output in case
        ``len(values) == 0``

        Implementing this function is an optimization. Certain window
        functions are not affected by the window frame definition or
        the query doesn't have a frame, and ``evaluate`` skips the
        (costly) window frame boundary calculation and the overhead of
        calling ``evaluate`` for each output row.

        For example, the `LAG` built in window function does not use
        the values of its window frame (it can be computed in one shot
        on the entire partition with ``Self::evaluate_all`` regardless of the
        window defined in the ``OVER`` clause)

        .. code-block:: text

            lag(x, 1) OVER (ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)

        However, ``avg()`` computes the average in the window and thus
        does use its window frame.

        .. code-block:: text

            avg(x) OVER (PARTITION BY y ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)
        """  # noqa: W505, E501

    def evaluate(
        self, values: list[pa.Array], eval_range: tuple[int, int]
    ) -> pa.Scalar:
        """Evaluate window function on a range of rows in an input partition.

        This is the simplest and most general function to implement
        but also the least performant as it creates output one row at
        a time. It is typically much faster to implement stateful
        evaluation using one of the other specialized methods on this
        trait.

        Returns a [`ScalarValue`] that is the value of the window
        function within `range` for the entire partition. Argument
        `values` contains the evaluation result of function arguments
        and evaluation results of ORDER BY expressions. If function has a
        single argument, `values[1..]` will contain ORDER BY expression results.
        """

    def evaluate_all_with_rank(
        self, num_rows: int, ranks_in_partition: list[tuple[int, int]]
    ) -> pa.Array:
        """Called for window functions that only need the rank of a row.

        Evaluate the partition evaluator against the partition using
        the row ranks. For example, ``rank(col("a"))`` produces

        .. code-block:: text

            a | rank
            - + ----
            A | 1
            A | 1
            C | 3
            D | 4
            D | 4

        For this case, `num_rows` would be `5` and the
        `ranks_in_partition` would be called with

        .. code-block:: text

            [
                (0,1),
                (2,2),
                (3,4),
            ]

        The user must implement this method if ``include_rank`` returns True.
        """

    def supports_bounded_execution(self) -> bool:
        """Can the window function be incrementally computed using bounded memory?"""
        return False

    def uses_window_frame(self) -> bool:
        """Does the window function use the values from the window frame?"""
        return False

    def include_rank(self) -> bool:
        """Can this function be evaluated with (only) rank?"""
        return False


class WindowUDFExportable(Protocol):
    """Type hint for object that has __datafusion_window_udf__ PyCapsule."""

    def __datafusion_window_udf__(self) -> object: ...  # noqa: D105


class WindowUDF:
    """Class for performing window user-defined functions (UDF).

    Window UDFs operate on a partition of rows. See
    also :py:class:`ScalarUDF` for operating on a row by row basis.
    """

    def __init__(
        self,
        name: str,
        func: Callable[[], WindowEvaluator],
        input_types: list[pa.DataType],
        return_type: pa.DataType,
        volatility: Volatility | str,
    ) -> None:
        """Instantiate a user-defined window function (UDWF).

        See :py:func:`udwf` for a convenience function and argument
        descriptions.
        """
        if hasattr(func, "__datafusion_window_udf__"):
            self._udwf = df_internal.WindowUDF.from_pycapsule(func)
            return
        self._udwf = df_internal.WindowUDF(
            name, func, input_types, return_type, str(volatility)
        )

    @classmethod
    def _from_internal(cls, internal: df_internal.WindowUDF) -> WindowUDF:
        """Wrap an already-constructed internal ``WindowUDF`` handle.

        Used by :py:meth:`SessionContext.udwf` and :py:meth:`from_pycapsule`
        to wrap a handle from the session's function registry or an FFI
        capsule without re-running :py:meth:`__init__`.
        """
        wrapper = cls.__new__(cls)
        wrapper._udwf = internal
        return wrapper

    @property
    def name(self) -> str:
        """Return the registered name of this UDWF.

        For UDWFs imported via the FFI capsule protocol, this is the
        name the capsule itself reports — not the ``name`` argument
        passed to the constructor (which is ignored on the FFI path).
        """
        return self._udwf.name

    def __repr__(self) -> str:
        """Print a string representation of the Window UDF."""
        return self._udwf.__repr__()

    def __call__(self, *args: Expr) -> Expr:
        """Execute the UDWF.

        This function is not typically called by an end user. These calls will
        occur during the evaluation of the dataframe.
        """
        args_raw = [arg.expr for arg in args]
        return Expr(self._udwf.__call__(*args_raw))

    @overload
    @staticmethod
    def udwf(
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> Callable[..., WindowUDF]: ...

    @overload
    @staticmethod
    def udwf(
        func: Callable[[], WindowEvaluator],
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> WindowUDF: ...

    @overload
    @staticmethod
    def udwf(func: WindowUDFExportable) -> WindowUDF: ...

    @overload
    @staticmethod
    def udwf(func: _PyCapsule) -> WindowUDF: ...

    @staticmethod
    def udwf(*args: Any, **kwargs: Any):  # noqa: D417
        """Create a new User-Defined Window Function (UDWF).

        This class can be used both as either a function or a decorator.

        Usage:
            - As a function: ``udwf(func, input_types, return_type, volatility, name)``.
            - As a decorator: ``@udwf(input_types, return_type, volatility, name)``.
              When using ``udwf`` as a decorator, do not pass ``func`` explicitly.

        Examples:
            >>> from datafusion.user_defined import WindowUDF, WindowEvaluator, udwf
            >>> class BiasedNumbers(WindowEvaluator):
            ...     def __init__(self, start: int = 0):
            ...         self.start = start
            ...     def evaluate_all(self, values, num_rows):
            ...         return pa.array(
            ...             [self.start + i for i in range(num_rows)])

            Using ``udwf`` as a function:

            >>> udwf1 = WindowUDF.udwf(
            ...     BiasedNumbers, pa.int64(), pa.int64(), "immutable")
            >>> def bias_10() -> BiasedNumbers:
            ...    return BiasedNumbers(10)
            >>> udwf2 = udwf(bias_10, pa.int64(), pa.int64(), "immutable")
            >>> udwf3 = udwf(
            ...     lambda: BiasedNumbers(20), pa.int64(), pa.int64(), "immutable"
            ... )

            Using ``udwf`` as a decorator:

            >>> @WindowUDF.udwf(pa.int64(), pa.int64(), "immutable")
            ... def biased_numbers():
            ...     return BiasedNumbers(10)

            Apply to a dataframe:

            >>> ctx = dfn.SessionContext()
            >>> df = ctx.from_pydict({"a": [10, 20, 30]})
            >>> df.select(udwf1(col("a")).alias("result")).to_pydict()
            {'result': [0, 1, 2]}
            >>> df.select(udwf2(col("a")).alias("result")).to_pydict()
            {'result': [10, 11, 12]}
            >>> df.select(udwf3(col("a")).alias("result")).to_pydict()
            {'result': [20, 21, 22]}
            >>> df.select(biased_numbers(col("a")).alias("result")).to_pydict()
            {'result': [10, 11, 12]}

        Args:
            func: Only needed when calling as a function. Skip this argument when
                using ``udwf`` as a decorator. If you have a Rust backed WindowUDF
                within a PyCapsule, you can pass this parameter and ignore the rest.
                They will be determined directly from the underlying function. See
                the online documentation for more information.
            input_types: The data types of the arguments.
            return_type: The data type of the return value.
            volatility: See :py:class:`Volatility` for allowed values.
            name: A descriptive name for the function.

        Returns:
            A user-defined window function that can be used in window function calls.
        """
        if not args and "func" in kwargs:
            args = (kwargs.pop("func"),)
        if args and (
            hasattr(args[0], "__datafusion_window_udf__") or _is_pycapsule(args[0])
        ):
            return WindowUDF.from_pycapsule(args[0])

        if args and callable(args[0]):
            # Case 1: Used as a function, require the first parameter to be callable
            return WindowUDF._create_window_udf(*args, **kwargs)
        # Case 2: Used as a decorator with parameters
        return WindowUDF._create_window_udf_decorator(*args, **kwargs)

    @staticmethod
    def _create_window_udf(
        func: Callable[[], WindowEvaluator],
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> WindowUDF:
        """Create a WindowUDF instance from function arguments."""
        if not callable(func):
            msg = "`func` must be callable."
            raise TypeError(msg)
        if not isinstance(func(), WindowEvaluator):
            msg = "`func` must implement the abstract base class WindowEvaluator"
            raise TypeError(msg)

        name = name or func.__qualname__.lower()
        input_types = (
            [input_types] if isinstance(input_types, pa.DataType) else input_types
        )

        return WindowUDF(name, func, input_types, return_type, volatility)

    @staticmethod
    def _get_default_name(func: Callable) -> str:
        """Get the default name for a function based on its attributes."""
        if hasattr(func, "__qualname__"):
            return func.__qualname__.lower()
        return func.__class__.__name__.lower()

    @staticmethod
    def _normalize_input_types(
        input_types: pa.DataType | list[pa.DataType],
    ) -> list[pa.DataType]:
        """Convert a single DataType to a list if needed."""
        if isinstance(input_types, pa.DataType):
            return [input_types]
        return input_types

    @staticmethod
    def _create_window_udf_decorator(
        input_types: pa.DataType | list[pa.DataType],
        return_type: pa.DataType,
        volatility: Volatility | str,
        name: str | None = None,
    ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]:
        """Create a decorator for a WindowUDF."""

        def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]:
            udwf_caller = WindowUDF._create_window_udf(
                func, input_types, return_type, volatility, name
            )

            @functools.wraps(func)
            def wrapper(*args: Any, **kwargs: Any) -> Expr:
                return udwf_caller(*args, **kwargs)

            return wrapper

        return decorator

    @staticmethod
    def from_pycapsule(func: WindowUDFExportable | _PyCapsule) -> WindowUDF:
        """Create a Window UDF from WindowUDF PyCapsule object.

        This function will instantiate a Window UDF that uses a DataFusion
        WindowUDF that is exported via the FFI bindings.
        """
        if _is_pycapsule(func):
            return WindowUDF._from_internal(df_internal.WindowUDF.from_pycapsule(func))

        func = cast("WindowUDFExportable", func)
        name = str(func.__class__)
        return WindowUDF(
            name=name,
            func=func,
            input_types=None,
            return_type=None,
            volatility=None,
        )


def _wrap_session_kwarg_for_udtf(func: Callable[..., Any]) -> Callable[..., Any]:
    """Adapt the raw internal session pyo3 object back to a Python wrapper.

    The Rust call site forwards a ``datafusion._internal.SessionContext``,
    but UDTF authors expect to interact with the public
    :class:`datafusion.SessionContext` wrapper. This closure wraps the
    internal object once per call before delegating to ``func``.
    """

    @functools.wraps(func, updated=())
    def adapter(*args: Any, session: Any, **kwargs: Any) -> Any:
        wrapped = SessionContext.__new__(SessionContext)
        wrapped.ctx = session
        return func(*args, session=wrapped, **kwargs)

    return adapter


class TableFunction:
    """Class for performing user-defined table functions (UDTF).

    Table functions generate new table providers based on the
    input expressions.
    """

    def __init__(
        self,
        name: str,
        func: Callable[..., Any],
        ctx: SessionContext | None = None,
        *,
        with_session: bool = False,
    ) -> None:
        """Instantiate a user-defined table function (UDTF).

        Set ``with_session=True`` to have the calling
        :class:`SessionContext` passed as a ``session`` keyword argument
        on each invocation. Use it inside the callback to look up
        registered tables, UDFs, or session configuration. When
        ``with_session`` is ``False`` (the default), ``func`` is invoked
        with the positional expression arguments only.

        ``with_session=True`` is only supported for pure-Python callables.
        Passing it together with an FFI-exported table function (one
        exposing ``__datafusion_table_function__``) raises
        :class:`TypeError`.

        Registry mutations performed through the injected session (such
        as registering tables or UDFs) propagate to the caller's
        :class:`SessionContext` because the registries are shared.
        Configuration changes do **not** propagate; the wrapper holds
        its own clone of the session config.

        See :py:func:`udtf` for a convenience function and argument
        descriptions.
        """
        if with_session and hasattr(func, "__datafusion_table_function__"):
            msg = (
                "`with_session=True` is not supported for FFI-exported table "
                "functions; session injection requires a pure-Python callable."
            )
            raise TypeError(msg)
        registered = _wrap_session_kwarg_for_udtf(func) if with_session else func
        self._udtf = df_internal.TableFunction(name, registered, ctx, with_session)

    def __call__(self, *args: Expr) -> Any:
        """Execute the UDTF and return a table provider."""
        args_raw = [arg.expr for arg in args]
        return self._udtf.__call__(*args_raw)

    @overload
    @staticmethod
    def udtf(
        name: str,
        *,
        with_session: bool = False,
    ) -> Callable[..., Any]: ...

    @overload
    @staticmethod
    def udtf(
        func: Callable[..., Any],
        name: str,
        *,
        with_session: bool = False,
    ) -> TableFunction: ...

    @staticmethod
    def udtf(*args: Any, with_session: bool = False, **kwargs: Any):
        """Create a new User-Defined Table Function (UDTF).

        Pass ``with_session=True`` to have the calling
        :class:`SessionContext` injected as a ``session`` keyword
        argument on each invocation.
        """
        if not args and "func" in kwargs:
            args = (kwargs.pop("func"),)
        if args and callable(args[0]):
            # Case 1: Used as a function, require the first parameter to be callable
            return TableFunction._create_table_udf(
                *args, with_session=with_session, **kwargs
            )
        if args and hasattr(args[0], "__datafusion_table_function__"):
            # Case 2: We have a datafusion FFI provided function
            if with_session:
                msg = (
                    "`with_session=True` is not supported for FFI-exported "
                    "table functions; session injection requires a "
                    "pure-Python callable."
                )
                raise TypeError(msg)
            return TableFunction(args[1], args[0])
        # Case 3: Used as a decorator with parameters
        return TableFunction._create_table_udf_decorator(
            *args, with_session=with_session, **kwargs
        )

    @staticmethod
    def _create_table_udf(
        func: Callable[..., Any],
        name: str,
        *,
        with_session: bool = False,
    ) -> TableFunction:
        """Create a TableFunction instance from function arguments."""
        if not callable(func):
            msg = "`func` must be callable."
            raise TypeError(msg)

        return TableFunction(name, func, with_session=with_session)

    @staticmethod
    def _create_table_udf_decorator(
        name: str | None = None,
        *,
        with_session: bool = False,
    ) -> Callable[[Callable[..., Any]], TableFunction]:
        """Create a decorator for a TableFunction."""

        def decorator(func: Callable[..., Any]) -> TableFunction:
            return TableFunction._create_table_udf(
                func, name, with_session=with_session
            )

        return decorator

    def __repr__(self) -> str:
        """User printable representation."""
        return self._udtf.__repr__()


# Convenience exports so we can import instead of treating as
# variables at the package root
udf = ScalarUDF.udf
udaf = AggregateUDF.udaf
udwf = WindowUDF.udwf
udtf = TableFunction.udtf