from __future__ import annotations

from functools import partial
from operator import methodcaller
from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol

from narwhals._compliant.any_namespace import (
    CatNamespace,
    DateTimeNamespace,
    ListNamespace,
    NameNamespace,
    StringNamespace,
    StructNamespace,
)
from narwhals._compliant.column import CompliantColumn
from narwhals._compliant.namespace import CompliantNamespace
from narwhals._compliant.typing import (
    AliasName,
    AliasNames,
    CompliantExprT_co,
    CompliantFrameT,
    CompliantLazyFrameT,
    CompliantSeriesOrNativeExprT_co,
    EagerDataFrameT,
    EagerExprT,
    EagerSeriesT,
    LazyExprT,
    NativeExprT,
)
from narwhals._utils import (
    _StoresCompliant,
    not_implemented,
    qualified_type_name,
    zip_strict,
)
from narwhals.dependencies import is_numpy_array, is_numpy_scalar
from narwhals.exceptions import MultiOutputExpressionError

if TYPE_CHECKING:
    from collections.abc import Sequence

    from typing_extensions import Self, TypeIs

    from narwhals._compliant.namespace import CompliantNamespace, EagerNamespace
    from narwhals._compliant.series import CompliantSeries
    from narwhals._compliant.typing import AliasNames, EvalNames, EvalSeries
    from narwhals._expression_parsing import ExprMetadata
    from narwhals._typing import NoDefault
    from narwhals._utils import Implementation, Version, _LimitedContext
    from narwhals.typing import (
        ClosedInterval,
        FillNullStrategy,
        IntoDType,
        ModeKeepStrategy,
        NonNestedLiteral,
        RankMethod,
        RollingInterpolationMethod,
        TimeUnit,
    )

__all__ = ["CompliantExpr", "DepthTrackingExpr", "EagerExpr", "LazyExpr", "NativeExpr"]


class NativeExpr(Protocol):
    """An `Expr`-like object from a package with [Lazy-only support](https://narwhals-dev.github.io/narwhals/extending/#levels-of-support).

    Protocol members are chosen *purely* for matching statically - as they
    are common to all currently supported packages.
    """

    def between(self, *args: Any, **kwds: Any) -> Any: ...

    # NOTE: None of these are annotated for `dx.Series`, but are added imperatively
    # Probably better to define a sub-protocol for `NativeSQLExpr`
    # - match `dx.Series` to `NativeExpr`
    # - match the others to `NativeSQLExpr`
    def __gt__(self, value: Any, /) -> Self: ...
    def __lt__(self, value: Any, /) -> Self: ...
    def __ge__(self, value: Any, /) -> Self: ...
    def __le__(self, value: Any, /) -> Self: ...
    def __eq__(self, value: Any, /) -> Self: ...  # type: ignore[override]
    def __ne__(self, value: Any, /) -> Self: ...  # type: ignore[override]


class CompliantExpr(
    CompliantColumn, Protocol[CompliantFrameT, CompliantSeriesOrNativeExprT_co]
):
    # NOTE: `narwhals`
    _implementation: Implementation
    _evaluate_output_names: EvalNames[CompliantFrameT]
    _alias_output_names: AliasNames | None
    # This should be set with extreme care, only in `_expression_parsing.py`,
    # and never from within any compliant class.
    _opt_metadata: ExprMetadata | None = None

    @property
    def _metadata(self) -> ExprMetadata:
        if self._opt_metadata is None:  # pragma: no cover
            msg = (
                "`_opt_metadata` is None. This is usually the result of trying to do "
                "some operation (such as `over`) which requires access to the metadata "
                "at the compliant level. You may want to consider rewriting your logic "
                "so that this operation is not necessary. Ideally you should avoid "
                "setting `_opt_metadata` manually."
            )
            raise AssertionError(msg)
        return self._opt_metadata

    def __call__(
        self, df: CompliantFrameT
    ) -> Sequence[CompliantSeriesOrNativeExprT_co]: ...
    def __narwhals_expr__(self) -> Self:  # pragma: no cover
        return self

    def __narwhals_namespace__(self) -> CompliantNamespace[CompliantFrameT, Self]: ...
    @classmethod
    def from_column_indices(
        cls, *column_indices: int, context: _LimitedContext
    ) -> Self: ...
    @classmethod
    def from_column_names(
        cls,
        evaluate_column_names: EvalNames[CompliantFrameT],
        /,
        *,
        context: _LimitedContext,
    ) -> Self: ...
    def broadcast(self) -> Self: ...

    # NOTE: `polars`
    def alias(self, name: str) -> Self: ...
    def all(self) -> Self: ...
    def any(self) -> Self: ...
    def any_value(self, *, ignore_nulls: bool) -> Self: ...
    def count(self) -> Self: ...
    def min(self) -> Self: ...
    def max(self) -> Self: ...
    def mean(self) -> Self: ...
    def sum(self) -> Self: ...
    def median(self) -> Self: ...
    def first(self, order_by: Sequence[str] = ()) -> Self: ...
    def last(self, order_by: Sequence[str] = ()) -> Self: ...
    def skew(self) -> Self: ...
    def kurtosis(self) -> Self: ...
    def std(self, *, ddof: int) -> Self: ...
    def var(self, *, ddof: int) -> Self: ...
    def n_unique(self) -> Self: ...
    def null_count(self) -> Self: ...
    def len(self) -> Self: ...
    def over(self, partition_by: Sequence[str], order_by: Sequence[str]) -> Self: ...
    def quantile(
        self, quantile: float, interpolation: RollingInterpolationMethod
    ) -> Self: ...
    def map_batches(
        self,
        function: Callable[[CompliantSeries[Any]], CompliantExpr[Any, Any]],
        return_dtype: IntoDType | None,
        *,
        returns_scalar: bool,
    ) -> Self: ...
    def replace_strict(
        self,
        default: Self | NoDefault,
        old: Sequence[Any],
        new: Sequence[Any],
        *,
        return_dtype: IntoDType | None,
    ) -> Self: ...
    @property
    def name(self) -> NameNamespace[Self]: ...


class ImplExpr(
    CompliantExpr[CompliantFrameT, CompliantSeriesOrNativeExprT_co],
    Protocol[CompliantFrameT, CompliantSeriesOrNativeExprT_co],
):
    @staticmethod
    def _eval_names_indices(indices: Sequence[int], /) -> EvalNames[CompliantFrameT]:
        def fn(df: CompliantFrameT) -> Sequence[str]:
            column_names = df.columns
            return [column_names[i] for i in indices]

        return fn

    def _evaluate_aliases(self, frame: CompliantFrameT, /) -> Sequence[str]:
        # NOTE: Ignore intermittent [False Negative]
        # Argument of type "CompliantFrameT@ImplExpr" cannot be assigned to parameter of type "CompliantFrameT@ImplExpr"
        #  Type "CompliantFrameT@ImplExpr" is not assignable to type "CompliantFrameT@ImplExpr"
        names = self._evaluate_output_names(frame)  # pyright: ignore[reportArgumentType]
        return alias(names) if (alias := self._alias_output_names) else names


class DepthTrackingExpr(
    ImplExpr[CompliantFrameT, CompliantSeriesOrNativeExprT_co],
    Protocol[CompliantFrameT, CompliantSeriesOrNativeExprT_co],
):
    # NOTE: pyright bug?
    # Method "from_column_names" overrides class "CompliantExpr" in an incompatible manner
    # Parameter 2 type mismatch: base parameter is type "EvalNames[CompliantFrameT@DepthTrackingExpr]", override parameter is type "EvalNames[CompliantFrameT@DepthTrackingExpr]"
    #   Type "EvalNames[CompliantFrameT@DepthTrackingExpr]" is not assignable to type "EvalNames[CompliantFrameT@DepthTrackingExpr]"
    #     Parameter 1: type "CompliantFrameT@DepthTrackingExpr" is incompatible with type "CompliantFrameT@DepthTrackingExpr"
    #       Type "CompliantFrameT@DepthTrackingExpr" is not assignable to type "CompliantFrameT@DepthTrackingExpr"
    @classmethod
    def from_column_names(  # pyright: ignore[reportIncompatibleMethodOverride]
        cls: type[Self],
        evaluate_column_names: EvalNames[CompliantFrameT],
        /,
        *,
        context: _LimitedContext,
    ) -> Self: ...

    def _is_elementary(self) -> bool:
        """Check if expr is elementary.

        Examples:
            - nw.col('a').mean()  # depth 1
            - nw.mean('a')  # depth 1
            - nw.len()  # depth 0

        as opposed to, say

            - nw.col('a').filter(nw.col('b')>nw.col('c')).max()

        Elementary expressions are the only ones supported properly in
        pandas, PyArrow, and Dask.
        """
        return len(list(self._metadata.op_nodes_reversed())) <= 2


class EagerExpr(
    DepthTrackingExpr[EagerDataFrameT, EagerSeriesT],
    Protocol[EagerDataFrameT, EagerSeriesT],
):
    _call: EvalSeries[EagerDataFrameT, EagerSeriesT]

    def __init__(
        self,
        call: EvalSeries[EagerDataFrameT, EagerSeriesT],
        *,
        evaluate_output_names: EvalNames[EagerDataFrameT],
        alias_output_names: AliasNames | None,
        implementation: Implementation,
        version: Version,
    ) -> None: ...

    def __call__(self, df: EagerDataFrameT) -> Sequence[EagerSeriesT]:
        return self._call(df)

    def __narwhals_namespace__(
        self,
    ) -> EagerNamespace[EagerDataFrameT, EagerSeriesT, Self, Any, Any]: ...
    @classmethod
    def _from_callable(
        cls,
        func: EvalSeries[EagerDataFrameT, EagerSeriesT],
        *,
        evaluate_output_names: EvalNames[EagerDataFrameT],
        alias_output_names: AliasNames | None,
        context: _LimitedContext,
    ) -> Self:
        return cls(
            func,
            evaluate_output_names=evaluate_output_names,
            alias_output_names=alias_output_names,
            implementation=context._implementation,
            version=context._version,
        )

    @classmethod
    def _from_series(cls, series: EagerSeriesT) -> Self:
        return cls(
            lambda _df: [series],
            evaluate_output_names=lambda _df: [series.name],
            alias_output_names=None,
            implementation=series._implementation,
            version=series._version,
        )

    def _with_alias_output_names(self, alias_name: AliasName | None, /) -> Self:
        current_alias_output_names = self._alias_output_names
        alias_output_names: AliasNames | None = (
            None
            if alias_name is None
            else (
                lambda output_names: [
                    alias_name(x) for x in current_alias_output_names(output_names)
                ]
            )
            if current_alias_output_names is not None
            else (lambda output_names: [alias_name(x) for x in output_names])
        )

        def func(df: EagerDataFrameT) -> list[EagerSeriesT]:
            if alias_output_names:
                return [
                    series.alias(name)
                    for series, name in zip_strict(
                        self(df), alias_output_names(self._evaluate_output_names(df))
                    )
                ]
            return [
                series.alias(name)
                for series, name in zip_strict(self(df), self._evaluate_output_names(df))
            ]

        return self.__class__(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=alias_output_names,
            implementation=self._implementation,
            version=self._version,
        )

    def _reuse_series(
        self, method_name: str, *, returns_scalar: bool = False, **kwargs: Any
    ) -> Self:
        """Reuse Series implementation for expression.

        If Series.foo is already defined, and we'd like Expr.foo to be the same, we can
        leverage this method to do that for us.

        Arguments:
            method_name: name of method.
            returns_scalar: whether the Series version returns a scalar. In this case,
                the expression version should return a 1-row Series.
            kwargs: keyword arguments to pass to function.
        """
        func = partial(
            self._reuse_series_inner,
            method_name=method_name,
            returns_scalar=returns_scalar,
            **kwargs,
        )
        return self._from_callable(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            context=self,
        )

    # For PyArrow.Series, we return Python Scalars (like Polars does) instead of PyArrow Scalars.
    # However, when working with expressions, we keep everything PyArrow-native.
    def _reuse_series_extra_kwargs(
        self, *, returns_scalar: bool = False
    ) -> dict[str, Any]:
        return {}

    @classmethod
    def _is_expr(cls, obj: Self | Any) -> TypeIs[Self]:
        return hasattr(obj, "__narwhals_expr__")

    def _reuse_series_inner(
        self,
        df: EagerDataFrameT,
        *,
        method_name: str,
        returns_scalar: bool,
        **kwargs: Any,
    ) -> Sequence[EagerSeriesT]:
        kwargs = {
            name: df._evaluate_single_output_expr(value)
            if self._is_expr(value)
            else value
            for name, value in kwargs.items()
        }
        method = methodcaller(
            method_name,
            **self._reuse_series_extra_kwargs(returns_scalar=returns_scalar),
            **kwargs,
        )
        out: Sequence[EagerSeriesT] = [
            series._from_scalar(method(series)) if returns_scalar else method(series)
            for series in self(df)
        ]
        aliases, names = self._evaluate_aliases(df), (s.name for s in out)
        if any(
            alias != name for alias, name in zip_strict(aliases, names)
        ):  # pragma: no cover
            msg = (
                f"Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues\n"
                f"Expression aliases: {aliases}\n"
            )
            raise AssertionError(msg)
        return out

    def _reuse_series_namespace(
        self,
        series_namespace: Literal["cat", "dt", "list", "name", "str", "struct"],
        method_name: str,
        **expressifiable_args: Any,
    ) -> Self:
        """Reuse Series implementation for expression.

        Just like `_reuse_series`, but for e.g. `Expr.dt.foo` instead
        of `Expr.foo`.

        Arguments:
            series_namespace: The Series namespace.
            method_name: name of method, within `series_namespace`.
            expressifiable_args: keyword arguments to pass to function, which may
                be expressifiable (e.g. `nw.col('a').str.replace('abc', nw.col('b')))`).
        """

        def inner(df: EagerDataFrameT) -> list[EagerSeriesT]:
            kwargs = {
                name: df._evaluate_single_output_expr(value)
                if self._is_expr(value)
                else value
                for name, value in expressifiable_args.items()
            }
            return [
                getattr(getattr(series, series_namespace), method_name)(**kwargs)
                for series in self(df)
            ]

        return self._from_callable(
            inner,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            context=self,
        )

    def broadcast(self) -> Self:
        # Mark the resulting Series with `_broadcast = True`.
        # Then, when extracting native objects, `extract_native` will
        # know what to do.
        def func(df: EagerDataFrameT) -> list[EagerSeriesT]:
            results = []
            for result in self(df):
                result._broadcast = True
                results.append(result)
            return results

        return type(self)(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            implementation=self._implementation,
            version=self._version,
        )

    def cast(self, dtype: IntoDType) -> Self:
        return self._reuse_series("cast", dtype=dtype)

    def _with_binary(self, operator: str, other: Self, /) -> Self:
        return self._reuse_series(operator, other=other)

    def _with_binary_right(self, operator: str, other: Self, /) -> Self:
        return self.alias("literal")._reuse_series(operator, other=other)

    def __eq__(self, other: Self) -> Self:  # type: ignore[override]
        return self._with_binary("__eq__", other)

    def __ne__(self, other: Self) -> Self:  # type: ignore[override]
        return self._with_binary("__ne__", other)

    def __ge__(self, other: Self) -> Self:
        return self._with_binary("__ge__", other)

    def __gt__(self, other: Self) -> Self:
        return self._with_binary("__gt__", other)

    def __le__(self, other: Self) -> Self:
        return self._with_binary("__le__", other)

    def __lt__(self, other: Self) -> Self:
        return self._with_binary("__lt__", other)

    def __and__(self, other: Self) -> Self:
        return self._with_binary("__and__", other)

    def __or__(self, other: Self) -> Self:
        return self._with_binary("__or__", other)

    def __add__(self, other: Self) -> Self:
        return self._with_binary("__add__", other)

    def __sub__(self, other: Self) -> Self:
        return self._with_binary("__sub__", other)

    def __rsub__(self, other: Self) -> Self:
        return self._with_binary_right("__rsub__", other)

    def __mul__(self, other: Self) -> Self:
        return self._with_binary("__mul__", other)

    def __truediv__(self, other: Self) -> Self:
        return self._with_binary("__truediv__", other)

    def __rtruediv__(self, other: Self) -> Self:
        return self._with_binary_right("__rtruediv__", other)

    def __floordiv__(self, other: Self) -> Self:
        return self._with_binary("__floordiv__", other)

    def __rfloordiv__(self, other: Self) -> Self:
        return self._with_binary_right("__rfloordiv__", other)

    def __pow__(self, other: Self) -> Self:
        return self._with_binary("__pow__", other)

    def __rpow__(self, other: Self) -> Self:
        return self._with_binary_right("__rpow__", other)

    def __mod__(self, other: Self) -> Self:
        return self._with_binary("__mod__", other)

    def __rmod__(self, other: Self) -> Self:
        return self._with_binary_right("__rmod__", other)

    # Unary
    def __invert__(self) -> Self:
        return self._reuse_series("__invert__")

    # Reductions
    def null_count(self) -> Self:
        return self._reuse_series("null_count", returns_scalar=True)

    def n_unique(self) -> Self:
        return self._reuse_series("n_unique", returns_scalar=True)

    def sum(self) -> Self:
        return self._reuse_series("sum", returns_scalar=True)

    def count(self) -> Self:
        return self._reuse_series("count", returns_scalar=True)

    def mean(self) -> Self:
        return self._reuse_series("mean", returns_scalar=True)

    def median(self) -> Self:
        return self._reuse_series("median", returns_scalar=True)

    def std(self, *, ddof: int) -> Self:
        return self._reuse_series("std", returns_scalar=True, ddof=ddof)

    def var(self, *, ddof: int) -> Self:
        return self._reuse_series("var", returns_scalar=True, ddof=ddof)

    def skew(self) -> Self:
        return self._reuse_series("skew", returns_scalar=True)

    def kurtosis(self) -> Self:
        return self._reuse_series("kurtosis", returns_scalar=True)

    def any(self) -> Self:
        return self._reuse_series("any", returns_scalar=True)

    def all(self) -> Self:
        return self._reuse_series("all", returns_scalar=True)

    def max(self) -> Self:
        return self._reuse_series("max", returns_scalar=True)

    def min(self) -> Self:
        return self._reuse_series("min", returns_scalar=True)

    def arg_min(self) -> Self:
        return self._reuse_series("arg_min", returns_scalar=True)

    def arg_max(self) -> Self:
        return self._reuse_series("arg_max", returns_scalar=True)

    # Other

    def clip(self, lower_bound: Self, upper_bound: Self) -> Self:
        return self._reuse_series(
            "clip", lower_bound=lower_bound, upper_bound=upper_bound
        )

    def clip_lower(self, lower_bound: Self) -> Self:
        return self._reuse_series("clip_lower", lower_bound=lower_bound)

    def clip_upper(self, upper_bound: Self) -> Self:
        return self._reuse_series("clip_upper", upper_bound=upper_bound)

    def is_null(self) -> Self:
        return self._reuse_series("is_null")

    def is_nan(self) -> Self:
        return self._reuse_series("is_nan")

    def fill_nan(self, value: float | None) -> Self:
        return self._reuse_series("fill_nan", value=value)

    def fill_null(
        self, value: Self | None, strategy: FillNullStrategy | None, limit: int | None
    ) -> Self:
        return self._reuse_series(
            "fill_null", value=value, strategy=strategy, limit=limit
        )

    def is_in(self, other: Any) -> Self:
        return self._reuse_series("is_in", other=other)

    def arg_true(self) -> Self:
        return self._reuse_series("arg_true")

    def filter(self, *predicates: Self) -> Self:
        plx = self.__narwhals_namespace__()
        predicate = plx.all_horizontal(*predicates, ignore_nulls=False)
        return self._reuse_series("filter", predicate=predicate)

    def drop_nulls(self) -> Self:
        return self._reuse_series("drop_nulls")

    def replace_strict(
        self,
        default: Self | NoDefault,
        old: Sequence[Any],
        new: Sequence[Any],
        *,
        return_dtype: IntoDType | None,
    ) -> Self:
        return self._reuse_series(
            "replace_strict", old=old, new=new, default=default, return_dtype=return_dtype
        )

    def sort(self, *, descending: bool, nulls_last: bool) -> Self:
        return self._reuse_series("sort", descending=descending, nulls_last=nulls_last)

    def abs(self) -> Self:
        return self._reuse_series("abs")

    def unique(self) -> Self:
        return self._reuse_series("unique", maintain_order=False)

    def diff(self) -> Self:
        return self._reuse_series("diff")

    def sample(
        self,
        n: int | None,
        *,
        fraction: float | None,
        with_replacement: bool,
        seed: int | None,
    ) -> Self:
        return self._reuse_series(
            "sample", n=n, fraction=fraction, with_replacement=with_replacement, seed=seed
        )

    def alias(self, name: str) -> Self:
        def alias_output_names(names: Sequence[str]) -> Sequence[str]:
            if len(names) != 1:
                msg = f"Expected function with single output, found output names: {names}"
                raise MultiOutputExpressionError(msg)
            return [name]

        return type(self)(
            lambda df: [series.alias(name) for series in self(df)],
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=alias_output_names,
            implementation=self._implementation,
            version=self._version,
        )

    def is_unique(self) -> Self:
        return self._reuse_series("is_unique")

    def is_first_distinct(self) -> Self:
        return self._reuse_series("is_first_distinct")

    def is_last_distinct(self) -> Self:
        return self._reuse_series("is_last_distinct")

    def quantile(
        self, quantile: float, interpolation: RollingInterpolationMethod
    ) -> Self:
        return self._reuse_series(
            "quantile",
            returns_scalar=True,
            quantile=quantile,
            interpolation=interpolation,
        )

    def head(self, n: int) -> Self:
        return self._reuse_series("head", n=n)

    def tail(self, n: int) -> Self:
        return self._reuse_series("tail", n=n)

    def round(self, decimals: int) -> Self:
        return self._reuse_series("round", decimals=decimals)

    def floor(self) -> Self:
        return self._reuse_series("floor")

    def ceil(self) -> Self:
        return self._reuse_series("ceil")

    def len(self) -> Self:
        return self._reuse_series("len", returns_scalar=True)

    def gather_every(self, n: int, offset: int) -> Self:
        return self._reuse_series("gather_every", n=n, offset=offset)

    def mode(self, *, keep: ModeKeepStrategy) -> Self:
        return self._reuse_series("mode", keep=keep)

    def is_finite(self) -> Self:
        return self._reuse_series("is_finite")

    def rolling_mean(self, window_size: int, *, min_samples: int, center: bool) -> Self:
        return self._reuse_series(
            "rolling_mean",
            window_size=window_size,
            min_samples=min_samples,
            center=center,
        )

    def rolling_std(
        self, window_size: int, *, min_samples: int, center: bool, ddof: int
    ) -> Self:
        return self._reuse_series(
            "rolling_std",
            window_size=window_size,
            min_samples=min_samples,
            center=center,
            ddof=ddof,
        )

    def rolling_sum(self, window_size: int, *, min_samples: int, center: bool) -> Self:
        return self._reuse_series(
            "rolling_sum", window_size=window_size, min_samples=min_samples, center=center
        )

    def rolling_var(
        self, window_size: int, *, min_samples: int, center: bool, ddof: int
    ) -> Self:
        return self._reuse_series(
            "rolling_var",
            window_size=window_size,
            min_samples=min_samples,
            center=center,
            ddof=ddof,
        )

    def map_batches(
        self,
        function: Callable[[Any], Any],
        return_dtype: IntoDType | None,
        *,
        returns_scalar: bool,
    ) -> Self:
        def func(df: EagerDataFrameT) -> Sequence[EagerSeriesT]:
            udf_series_in = self(df)
            output_names = (input_series.name for input_series in udf_series_in)
            udf_series_out = tuple(function(series) for series in udf_series_in)
            _first_in, _first_out = udf_series_in[0], udf_series_out[0]

            result: Sequence[EagerSeriesT]
            it = zip_strict(udf_series_out, output_names)
            if is_numpy_array(_first_out) or is_numpy_scalar(_first_out):
                from_numpy = partial(_first_in.from_numpy, context=self)
                result = tuple(from_numpy(arr).alias(out_name) for arr, out_name in it)
            elif isinstance(_first_out, _first_in.__class__):  # compliant series
                result = tuple(series.alias(out_name) for series, out_name in it)
            else:  # If everything else fails, assume scalar case
                from_scalar = _first_in._from_scalar
                result = tuple(from_scalar(val).alias(out_name) for val, out_name in it)

            if return_dtype is not None:
                result = tuple(series.cast(return_dtype) for series in result)

            is_scalar_result = tuple(len(r) == 1 for r in result)
            if (not returns_scalar) and any(is_scalar_result) and (len(df) > 1):
                _idx = is_scalar_result.index(True)  # Index of first result with length 1
                _type = type(udf_series_out[_idx])
                msg = (
                    "`map_batches` with `returns_scalar=False` must return a Series; "
                    f"found '{qualified_type_name(_type)}'.\n\nIf `returns_scalar` "
                    "is set to `True`, a returned value can be a scalar value."
                )
                raise TypeError(msg)
            return result

        return self._from_callable(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            context=self,
        )

    def shift(self, n: int) -> Self:
        return self._reuse_series("shift", n=n)

    def cum_sum(self, *, reverse: bool) -> Self:
        return self._reuse_series("cum_sum", reverse=reverse)

    def cum_count(self, *, reverse: bool) -> Self:
        return self._reuse_series("cum_count", reverse=reverse)

    def cum_min(self, *, reverse: bool) -> Self:
        return self._reuse_series("cum_min", reverse=reverse)

    def cum_max(self, *, reverse: bool) -> Self:
        return self._reuse_series("cum_max", reverse=reverse)

    def cum_prod(self, *, reverse: bool) -> Self:
        return self._reuse_series("cum_prod", reverse=reverse)

    def rank(self, method: RankMethod, *, descending: bool) -> Self:
        return self._reuse_series("rank", method=method, descending=descending)

    def log(self, base: float) -> Self:
        return self._reuse_series("log", base=base)

    def exp(self) -> Self:
        return self._reuse_series("exp")

    def sin(self) -> Self:
        return self._reuse_series("sin")

    def cos(self) -> Self:
        return self._reuse_series("cos")

    def sqrt(self) -> Self:
        return self._reuse_series("sqrt")

    def is_between(
        self, lower_bound: Self, upper_bound: Self, closed: ClosedInterval
    ) -> Self:
        return self._reuse_series(
            "is_between", lower_bound=lower_bound, upper_bound=upper_bound, closed=closed
        )

    def first(self, order_by: Sequence[str] = ()) -> Self:
        if not order_by:
            return self._reuse_series("first", returns_scalar=True)

        def func(df: EagerDataFrameT) -> Sequence[Any]:
            df = df.sort(*order_by, descending=False, nulls_last=False)
            return self.first()(df)

        return self._from_callable(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            context=self,
        )

    def last(self, order_by: Sequence[str] = ()) -> Self:
        if not order_by:
            return self._reuse_series("last", returns_scalar=True)

        def func(df: EagerDataFrameT) -> Sequence[Any]:
            df = df.sort(*order_by, descending=False, nulls_last=False)
            return self.last()(df)

        return self._from_callable(
            func,
            evaluate_output_names=self._evaluate_output_names,
            alias_output_names=self._alias_output_names,
            context=self,
        )

    def any_value(self, *, ignore_nulls: bool) -> Self:
        return self._reuse_series(
            "any_value", returns_scalar=True, ignore_nulls=ignore_nulls
        )

    @property
    def cat(self) -> EagerExprCatNamespace[Self]:
        return EagerExprCatNamespace(self)

    @property
    def dt(self) -> EagerExprDateTimeNamespace[Self]:
        return EagerExprDateTimeNamespace(self)

    @property
    def list(self) -> EagerExprListNamespace[Self]:
        return EagerExprListNamespace(self)

    @property
    def name(self) -> EagerExprNameNamespace[Self]:
        return EagerExprNameNamespace(self)

    @property
    def str(self) -> EagerExprStringNamespace[Self]:
        return EagerExprStringNamespace(self)

    @property
    def struct(self) -> EagerExprStructNamespace[Self]:
        return EagerExprStructNamespace(self)


# mypy thinks `NativeExprT` should be covariant, pyright thinks it should be invariant
class LazyExpr(  # type: ignore[misc]
    ImplExpr[CompliantLazyFrameT, NativeExprT], Protocol[CompliantLazyFrameT, NativeExprT]
):
    def _with_alias_output_names(self, func: AliasNames | None, /) -> Self: ...
    def alias(self, name: str) -> Self:
        def fn(names: Sequence[str]) -> Sequence[str]:
            if len(names) != 1:
                msg = f"Expected function with single output, found output names: {names}"
                raise ValueError(msg)
            return [name]

        return self._with_alias_output_names(fn)

    @property
    def name(self) -> LazyExprNameNamespace[Self]:
        return LazyExprNameNamespace(self)

    ewm_mean = not_implemented()  # type: ignore[misc]
    map_batches = not_implemented()  # type: ignore[misc]
    cat: not_implemented = not_implemented()  # type: ignore[assignment]


class _ExprNamespace(  # type: ignore[misc]
    _StoresCompliant[CompliantExprT_co], Protocol[CompliantExprT_co]
):
    _compliant_expr: CompliantExprT_co

    @property
    def compliant(self) -> CompliantExprT_co:
        return self._compliant_expr


class EagerExprNamespace(_ExprNamespace[EagerExprT], Generic[EagerExprT]):
    def __init__(self, expr: EagerExprT, /) -> None:
        self._compliant_expr = expr


class LazyExprNamespace(_ExprNamespace[LazyExprT], Generic[LazyExprT]):
    def __init__(self, expr: LazyExprT, /) -> None:
        self._compliant_expr = expr


class EagerExprCatNamespace(
    EagerExprNamespace[EagerExprT], CatNamespace[EagerExprT], Generic[EagerExprT]
):
    def get_categories(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("cat", "get_categories")


class EagerExprDateTimeNamespace(
    EagerExprNamespace[EagerExprT], DateTimeNamespace[EagerExprT], Generic[EagerExprT]
):
    def to_string(self, format: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "to_string", format=format)

    def replace_time_zone(self, time_zone: str | None) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "dt", "replace_time_zone", time_zone=time_zone
        )

    def convert_time_zone(self, time_zone: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "dt", "convert_time_zone", time_zone=time_zone
        )

    def timestamp(self, time_unit: TimeUnit) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "dt", "timestamp", time_unit=time_unit
        )

    def date(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "date")

    def year(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "year")

    def month(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "month")

    def day(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "day")

    def hour(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "hour")

    def minute(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "minute")

    def second(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "second")

    def millisecond(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "millisecond")

    def microsecond(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "microsecond")

    def nanosecond(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "nanosecond")

    def ordinal_day(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "ordinal_day")

    def weekday(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "weekday")

    def total_minutes(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "total_minutes")

    def total_seconds(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "total_seconds")

    def total_milliseconds(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "total_milliseconds")

    def total_microseconds(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "total_microseconds")

    def total_nanoseconds(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "total_nanoseconds")

    def truncate(self, every: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "truncate", every=every)

    def offset_by(self, by: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("dt", "offset_by", by=by)


class EagerExprListNamespace(
    EagerExprNamespace[EagerExprT], ListNamespace[EagerExprT], Generic[EagerExprT]
):
    def len(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "len")

    def unique(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "unique")

    def contains(self, item: NonNestedLiteral) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "contains", item=item)

    def get(self, index: int) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "get", index=index)

    def min(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "min")

    def max(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "max")

    def mean(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "mean")

    def median(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "median")

    def sum(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("list", "sum")

    def sort(self, *, descending: bool, nulls_last: bool) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "list", "sort", descending=descending, nulls_last=nulls_last
        )


class CompliantExprNameNamespace(  # type: ignore[misc]
    _ExprNamespace[CompliantExprT_co],
    NameNamespace[CompliantExprT_co],
    Protocol[CompliantExprT_co],
):
    def keep(self) -> CompliantExprT_co:
        return self._from_callable(None)

    def map(self, function: AliasName) -> CompliantExprT_co:
        return self._from_callable(function)

    def prefix(self, prefix: str) -> CompliantExprT_co:
        return self._from_callable(lambda name: f"{prefix}{name}")

    def suffix(self, suffix: str) -> CompliantExprT_co:
        return self._from_callable(lambda name: f"{name}{suffix}")

    def to_lowercase(self) -> CompliantExprT_co:
        return self._from_callable(str.lower)

    def to_uppercase(self) -> CompliantExprT_co:
        return self._from_callable(str.upper)

    @staticmethod
    def _alias_output_names(func: AliasName, /) -> AliasNames:
        def fn(output_names: Sequence[str], /) -> Sequence[str]:
            return [func(name) for name in output_names]

        return fn

    def _from_callable(self, func: AliasName | None, /) -> CompliantExprT_co: ...


class EagerExprNameNamespace(
    EagerExprNamespace[EagerExprT],
    CompliantExprNameNamespace[EagerExprT],
    Generic[EagerExprT],
):
    def _from_callable(self, func: AliasName | None) -> EagerExprT:
        expr = self.compliant
        return expr._with_alias_output_names(func)


class LazyExprNameNamespace(
    LazyExprNamespace[LazyExprT],
    CompliantExprNameNamespace[LazyExprT],
    Generic[LazyExprT],
):
    def _from_callable(self, func: AliasName | None) -> LazyExprT:
        expr = self.compliant
        output_names = self._alias_output_names(func) if func else None
        return expr._with_alias_output_names(output_names)


class EagerExprStringNamespace(
    EagerExprNamespace[EagerExprT], StringNamespace[EagerExprT], Generic[EagerExprT]
):
    def len_chars(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "len_chars")

    def replace(
        self, value: EagerExprT, pattern: str, *, literal: bool, n: int
    ) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "replace", pattern=pattern, value=value, literal=literal, n=n
        )

    def replace_all(
        self, value: EagerExprT, pattern: str, *, literal: bool
    ) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "replace_all", pattern=pattern, value=value, literal=literal
        )

    def strip_chars(self, characters: str | None) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "strip_chars", characters=characters
        )

    def starts_with(self, prefix: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "starts_with", prefix=prefix)

    def ends_with(self, suffix: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "ends_with", suffix=suffix)

    def contains(self, pattern: str, *, literal: bool) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "contains", pattern=pattern, literal=literal
        )

    def slice(self, offset: int, length: int | None) -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "slice", offset=offset, length=length
        )

    def split(self, by: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "split", by=by)

    def to_datetime(self, format: str | None) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "to_datetime", format=format)

    def to_date(self, format: str | None) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "to_date", format=format)

    def to_lowercase(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "to_lowercase")

    def to_uppercase(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "to_uppercase")

    def zfill(self, width: int) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "zfill", width=width)

    def to_titlecase(self) -> EagerExprT:
        return self.compliant._reuse_series_namespace("str", "to_titlecase")

    def pad_start(self, length: int, fill_char: str = " ") -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "pad_start", length=length, fill_char=fill_char
        )

    def pad_end(self, length: int, fill_char: str = " ") -> EagerExprT:
        return self.compliant._reuse_series_namespace(
            "str", "pad_end", length=length, fill_char=fill_char
        )


class EagerExprStructNamespace(
    EagerExprNamespace[EagerExprT], StructNamespace[EagerExprT], Generic[EagerExprT]
):
    def field(self, name: str) -> EagerExprT:
        return self.compliant._reuse_series_namespace("struct", "field", name=name).alias(
            name
        )
