# Utilities for expression parsing
# Useful for backends which don't have any concept of expressions, such
# and pandas or PyArrow.
# ! Any change to this module will trigger the pyspark and pyspark-connect tests in CI
from __future__ import annotations

from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Callable, Literal, cast

from narwhals._utils import zip_strict
from narwhals.dependencies import is_numpy_array_1d
from narwhals.exceptions import (
    InvalidIntoExprError,
    InvalidOperationError,
    MultiOutputExpressionError,
)

if TYPE_CHECKING:
    from collections.abc import Iterator, Sequence

    from typing_extensions import Never, TypeIs

    from narwhals._compliant import CompliantExpr, CompliantFrameT
    from narwhals._compliant.typing import (
        AliasNames,
        CompliantExprAny,
        CompliantFrameAny,
        CompliantNamespaceAny,
        EvalNames,
    )
    from narwhals.expr import Expr
    from narwhals.series import Series
    from narwhals.typing import IntoExpr, NonNestedLiteral, _1DArray


def is_expr(obj: Any) -> TypeIs[Expr]:
    """Check whether `obj` is a Narwhals Expr."""
    from narwhals.expr import Expr

    return isinstance(obj, Expr)


def is_series(obj: Any) -> TypeIs[Series[Any]]:
    """Check whether `obj` is a Narwhals Expr."""
    from narwhals.series import Series

    return isinstance(obj, Series)


def combine_evaluate_output_names(
    *exprs: CompliantExpr[CompliantFrameT, Any],
) -> EvalNames[CompliantFrameT]:
    # Follow left-hand-rule for naming. E.g. `nw.sum_horizontal(expr1, expr2)` takes the
    # first name of `expr1`.
    def evaluate_output_names(df: CompliantFrameT) -> Sequence[str]:
        return exprs[0]._evaluate_output_names(df)[:1]

    return evaluate_output_names


def combine_alias_output_names(*exprs: CompliantExprAny) -> AliasNames | None:
    # Follow left-hand-rule for naming. E.g. `nw.sum_horizontal(expr1.alias(alias), expr2)` takes the
    # aliasing function of `expr1` and apply it to the first output name of `expr1`.
    if exprs[0]._alias_output_names is None:
        return None

    def alias_output_names(names: Sequence[str]) -> Sequence[str]:
        return exprs[0]._alias_output_names(names)[:1]  # type: ignore[misc]

    return alias_output_names


def evaluate_output_names_and_aliases(
    expr: CompliantExprAny, df: CompliantFrameAny, exclude: Sequence[str]
) -> tuple[Sequence[str], Sequence[str]]:
    output_names = expr._evaluate_output_names(df)
    aliases = (
        output_names
        if expr._alias_output_names is None
        else expr._alias_output_names(output_names)
    )
    if exclude and expr._metadata.expansion_kind.is_multi_unnamed():
        output_names, aliases = zip_strict(
            *[
                (x, alias)
                for x, alias in zip_strict(output_names, aliases)
                if x not in exclude
            ]
        )
    return output_names, aliases


class ExprKind(Enum):
    """Describe which kind of expression we are dealing with."""

    LITERAL = auto()
    """e.g. `nw.lit(1)`"""

    AGGREGATION = auto()
    """Reduces to a single value, not affected by row order, e.g. `nw.col('a').mean()`"""

    ORDERABLE_AGGREGATION = auto()
    """Reduces to a single value, affected by row order, e.g. `nw.col('a').arg_max()`"""

    ELEMENTWISE = auto()
    """Preserves length, can operate without context for surrounding rows, e.g. `nw.col('a').abs()`."""

    ORDERABLE_WINDOW = auto()
    """Depends on the rows around it and on their order, e.g. `diff`."""

    WINDOW = auto()
    """Depends on the rows around it and possibly their order, e.g. `rank`."""

    FILTRATION = auto()
    """Changes length, not affected by row order, e.g. `drop_nulls`."""

    ORDERABLE_FILTRATION = auto()
    """Changes length, affected by row order, e.g. `tail`."""

    OVER = auto()
    """Results from calling `.over` on expression."""

    COL = auto()
    """Results from calling `nw.col`."""

    NTH = auto()
    """Results from calling `nw.nth`."""

    EXCLUDE = auto()
    """Results from calling `nw.exclude`."""

    ALL = auto()
    """Results from calling `nw.all`."""

    SELECTOR = auto()
    """Results from creating an expression with a selector."""

    WHEN_THEN = auto()
    """Results from `when/then expression`, possibly followed by `otherwise`."""

    SERIES = auto()
    """Results from converting a Series to Expr."""

    @property
    def is_orderable(self) -> bool:
        # Any operation which may be affected by `order_by`, such as `cum_sum`,
        # `diff`, `rank`, `arg_max`, ...
        return self in {
            ExprKind.ORDERABLE_WINDOW,
            ExprKind.WINDOW,
            ExprKind.ORDERABLE_AGGREGATION,
            ExprKind.ORDERABLE_FILTRATION,
        }

    @property
    def is_elementwise(self) -> bool:
        # Any operation which can operate on each row independently
        # of the rows around it, e.g. `abs(), __add__, sum_horizontal, ...`
        return self in {
            ExprKind.ALL,
            ExprKind.COL,
            ExprKind.ELEMENTWISE,
            ExprKind.EXCLUDE,
            ExprKind.LITERAL,
            ExprKind.NTH,
            ExprKind.SELECTOR,
            ExprKind.SERIES,
            ExprKind.WHEN_THEN,
        }

    @property
    def is_scalar_like(self) -> bool:
        return self in {
            ExprKind.AGGREGATION,
            ExprKind.LITERAL,
            ExprKind.ORDERABLE_AGGREGATION,
        }


def is_scalar_like(obj: CompliantExprAny) -> bool:
    return obj._metadata.is_scalar_like


class ExpansionKind(Enum):
    """Describe what kind of expansion the expression performs."""

    SINGLE = auto()
    """e.g. `nw.col('a'), nw.sum_horizontal(nw.all())`"""

    MULTI_NAMED = auto()
    """e.g. `nw.col('a', 'b')`"""

    MULTI_UNNAMED = auto()
    """e.g. `nw.all()`, nw.nth(0, 1)"""

    def is_multi_unnamed(self) -> bool:
        return self is ExpansionKind.MULTI_UNNAMED

    def is_multi_output(self) -> bool:
        return self in {ExpansionKind.MULTI_NAMED, ExpansionKind.MULTI_UNNAMED}

    def __and__(self, other: ExpansionKind) -> Literal[ExpansionKind.MULTI_UNNAMED]:
        if self is ExpansionKind.MULTI_UNNAMED and other is ExpansionKind.MULTI_UNNAMED:
            # e.g. nw.selectors.all() - nw.selectors.numeric().
            return ExpansionKind.MULTI_UNNAMED
        # Don't attempt anything more complex, keep it simple and raise in the face of ambiguity.
        msg = f"Unsupported ExpansionKind combination, got {self} and {other}, please report a bug."  # pragma: no cover
        raise AssertionError(msg)  # pragma: no cover


class ExprNode:
    """An operation to create or modify an expression.

    Parameters:
        kind: ExprKind of operation.
        name: Name of function, as defined in the compliant protocols.
        exprs: Expressifiable arguments to function.
        str_as_lit: Whether to interpret strings as literals when they
            are present in `exprs`.
        allow_multi_output: Whether to allow any of `exprs` to be multi-output.
        kwargs: Other (non-expressifiable) arguments to function.
    """

    def __init__(
        self,
        kind: ExprKind,
        name: str,
        /,
        *exprs: IntoExpr | NonNestedLiteral,
        str_as_lit: bool = False,
        allow_multi_output: bool = False,
        **kwargs: Any,
    ) -> None:
        self.kind: ExprKind = kind
        self.name: str = name
        self.exprs: Sequence[IntoExpr | NonNestedLiteral] = exprs
        self.kwargs: dict[str, Any] = kwargs
        self.str_as_lit: bool = str_as_lit
        self.allow_multi_output: bool = allow_multi_output

        # Cached methods.
        self._is_orderable_cached: bool | None = None
        self._is_elementwise_cached: bool | None = None

    def __repr__(self) -> str:
        if self.name == "col":
            names = ", ".join(str(x) for x in self.kwargs["names"])
            return f"col({names})"
        arg_str = []
        expr_repr = ", ".join(str(x) for x in self.exprs)
        kwargs_repr = ", ".join(f"{key}={value}" for key, value in self.kwargs.items())
        if self.exprs:
            arg_str.append(expr_repr)
        if self.kwargs:
            arg_str.append(kwargs_repr)
        return f"{self.name}({', '.join(arg_str)})"

    def as_dict(self) -> dict[str, Any]:  # pragma: no cover
        # Just for debugging.
        return {
            "kind": self.kind,
            "name": self.name,
            "exprs": self.exprs,
            "kwargs": self.kwargs,
            "str_as_lit": self.str_as_lit,
            "allow_multi_output": self.allow_multi_output,
        }

    def _with_kwargs(self, **kwargs: Any) -> ExprNode:
        return self.__class__(
            self.kind, self.name, *self.exprs, str_as_lit=self.str_as_lit, **kwargs
        )

    def _push_down_over_node_in_place(
        self, over_node: ExprNode, over_node_without_order_by: ExprNode
    ) -> None:
        exprs: list[IntoExpr | NonNestedLiteral] = []
        # Note: please keep this as a for-loop (rather than a list-comprehension)
        # so that pytest-cov highlights any uncovered branches.
        over_node_order_by = over_node.kwargs["order_by"]
        over_node_partition_by = over_node.kwargs["partition_by"]
        for expr in self.exprs:
            if not is_expr(expr):
                exprs.append(expr)
            elif over_node_order_by and any(
                expr_node.is_orderable() for expr_node in expr._nodes
            ):
                exprs.append(expr._with_over_node(over_node))
            elif over_node_partition_by and not all(
                expr_node.is_elementwise() for expr_node in expr._nodes
            ):
                exprs.append(expr._with_over_node(over_node_without_order_by))
            else:
                # If there's no `partition_by`, then `over_node_without_order_by` is a no-op.
                exprs.append(expr)
        self.exprs = exprs

    def is_orderable(self) -> bool:
        if self._is_orderable_cached is None:
            # Note: don't combine these if/then statements so that pytest-cov shows if
            # anything is uncovered.
            if self.kind.is_orderable:  # noqa: SIM114
                self._is_orderable_cached = True
            elif any(
                any(node.is_orderable() for node in expr._nodes)
                for expr in self.exprs
                if is_expr(expr)
            ):
                self._is_orderable_cached = True
            else:
                self._is_orderable_cached = False
        return self._is_orderable_cached

    def is_elementwise(self) -> bool:
        if self._is_elementwise_cached is None:
            # Note: don't combine these if/then statements so that pytest-cov shows if
            # anything is uncovered.
            if not self.kind.is_elementwise:  # noqa: SIM114
                self._is_elementwise_cached = False
            elif any(
                any(not node.is_elementwise() for node in expr._nodes)
                for expr in self.exprs
                if is_expr(expr)
            ):
                self._is_elementwise_cached = False
            else:
                self._is_elementwise_cached = True
        return self._is_elementwise_cached


class ExprMetadata:
    """Expression metadata.

    Parameters:
        expansion_kind: What kind of expansion the expression performs.
        has_windows: Whether it already contains window functions.
        is_elementwise: Whether it can operate row-by-row without context
            of the other rows around it.
        is_literal: Whether it is just a literal wrapped in an expression.
        is_scalar_like: Whether it is a literal or an aggregation.
        n_orderable_ops: The number of order-dependent operations. In the
            lazy case, this number must be `0` by the time the expression
            is evaluated.
        preserves_length: Whether the expression preserves the input length.
        current_node: The current ExprNode in the linked list.
        prev: Reference to the previous ExprMetadata in the linked list (None for root).
    """

    __slots__ = (
        "current_node",
        "expansion_kind",
        "has_windows",
        "is_elementwise",
        "is_literal",
        "is_scalar_like",
        "n_orderable_ops",
        "preserves_length",
        "prev",
    )

    def __init__(
        self,
        expansion_kind: ExpansionKind,
        *,
        has_windows: bool = False,
        n_orderable_ops: int = 0,
        preserves_length: bool = True,
        is_elementwise: bool = True,
        is_scalar_like: bool = False,
        is_literal: bool = False,
        current_node: ExprNode,
        prev: ExprMetadata | None = None,
    ) -> None:
        if is_literal:
            assert is_scalar_like  # noqa: S101  # debug assertion
        self.expansion_kind: ExpansionKind = expansion_kind
        self.has_windows: bool = has_windows
        self.n_orderable_ops: int = n_orderable_ops
        self.is_elementwise: bool = is_elementwise
        self.preserves_length: bool = preserves_length
        self.is_scalar_like: bool = is_scalar_like
        self.is_literal: bool = is_literal
        self.current_node: ExprNode = current_node
        self.prev: ExprMetadata | None = prev

    def __init_subclass__(cls, /, *args: Any, **kwds: Any) -> Never:  # pragma: no cover
        msg = f"Cannot subclass {cls.__name__!r}"
        raise TypeError(msg)

    def __repr__(self) -> str:  # pragma: no cover
        nodes = tuple(reversed(tuple(self.iter_nodes_reversed())))
        return (
            f"ExprMetadata(\n"
            f"  expansion_kind: {self.expansion_kind},\n"
            f"  has_windows: {self.has_windows},\n"
            f"  n_orderable_ops: {self.n_orderable_ops},\n"
            f"  is_elementwise: {self.is_elementwise},\n"
            f"  preserves_length: {self.preserves_length},\n"
            f"  is_scalar_like: {self.is_scalar_like},\n"
            f"  is_literal: {self.is_literal},\n"
            f"  nodes: {nodes},\n"
            ")"
        )

    def iter_nodes_reversed(self) -> Iterator[ExprNode]:
        """Iterate through all nodes from current to root."""
        current: ExprMetadata | None = self
        while current is not None:
            yield current.current_node
            current = current.prev

    @classmethod
    def from_node(
        cls, node: ExprNode, *compliant_exprs: CompliantExprAny
    ) -> ExprMetadata:
        return KIND_TO_METADATA_CONSTRUCTOR[node.kind](node, *compliant_exprs)

    def with_node(
        self,
        node: ExprNode,
        compliant_expr: CompliantExprAny,
        *compliant_expr_args: CompliantExprAny,
    ) -> ExprMetadata:
        return KIND_TO_METADATA_UPDATER[node.kind](
            self, node, compliant_expr, *compliant_expr_args
        )

    @classmethod
    def from_aggregation(cls, node: ExprNode) -> ExprMetadata:
        return cls(
            ExpansionKind.SINGLE,
            is_elementwise=False,
            preserves_length=False,
            is_scalar_like=True,
            current_node=node,
            prev=None,
        )

    @classmethod
    def from_literal(cls, node: ExprNode) -> ExprMetadata:
        return cls(
            ExpansionKind.SINGLE,
            is_elementwise=True,
            preserves_length=False,
            is_literal=True,
            is_scalar_like=True,
            current_node=node,
            prev=None,
        )

    @classmethod
    def from_series(cls, node: ExprNode) -> ExprMetadata:
        return cls(ExpansionKind.SINGLE, current_node=node, prev=None)

    @classmethod
    def from_col(cls, node: ExprNode) -> ExprMetadata:
        # e.g. `nw.col('a')`, `nw.nth(0)`
        return (
            cls(ExpansionKind.SINGLE, current_node=node, prev=None)
            if len(node.kwargs["names"]) == 1
            else cls.from_selector_multi_named(node)
        )

    @classmethod
    def from_nth(cls, node: ExprNode) -> ExprMetadata:
        return (
            cls(ExpansionKind.SINGLE, current_node=node, prev=None)
            if len(node.kwargs["indices"]) == 1
            else cls.from_selector_multi_named(node)
        )

    @classmethod
    def from_selector_multi_named(cls, node: ExprNode) -> ExprMetadata:
        # e.g. `nw.col('a', 'b')`
        return cls(ExpansionKind.MULTI_NAMED, current_node=node, prev=None)

    @classmethod
    def from_selector_multi_unnamed(cls, node: ExprNode) -> ExprMetadata:
        # e.g. `nw.all()`
        return cls(ExpansionKind.MULTI_UNNAMED, current_node=node, prev=None)

    @classmethod
    def from_elementwise(
        cls, node: ExprNode, *compliant_exprs: CompliantExprAny
    ) -> ExprMetadata:
        return combine_metadata(
            *compliant_exprs, to_single_output=True, current_node=node, prev=None
        )

    @property
    def is_filtration(self) -> bool:
        return not self.preserves_length and not self.is_scalar_like

    def with_aggregation(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata:
        if self.is_scalar_like:
            msg = "Can't apply aggregations to scalar-like expressions."
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=self.has_windows,
            n_orderable_ops=self.n_orderable_ops,
            preserves_length=False,
            is_elementwise=False,
            is_scalar_like=True,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_orderable_aggregation(
        self, node: ExprNode, _ce: CompliantExprAny
    ) -> ExprMetadata:
        # Deprecated, used only in stable.v1.
        if self.is_scalar_like:  # pragma: no cover
            msg = "Can't apply aggregations to scalar-like expressions."
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=self.has_windows,
            n_orderable_ops=self.n_orderable_ops + 1,
            preserves_length=False,
            is_elementwise=False,
            is_scalar_like=True,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_elementwise(
        self,
        node: ExprNode,
        compliant_expr: CompliantExprAny,
        *compliant_expr_args: CompliantExprAny,
    ) -> ExprMetadata:
        return combine_metadata(
            compliant_expr,
            *compliant_expr_args,
            to_single_output=False,
            current_node=node,
            prev=compliant_expr._metadata,
        )

    def with_window(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata:
        # Window function which may (but doesn't have to) be used with `over(order_by=...)`.
        if self.is_scalar_like:
            msg = "Can't apply window (e.g. `rank`) to scalar-like expression."
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=self.has_windows,
            # The function isn't order-dependent (but, users can still use `order_by` if they wish!),
            # so we don't increment `n_orderable_ops`.
            n_orderable_ops=self.n_orderable_ops,
            preserves_length=self.preserves_length,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_orderable_window(
        self, node: ExprNode, _ce: CompliantExprAny
    ) -> ExprMetadata:
        # Window function which must be used with `over(order_by=...)`.
        if self.is_scalar_like:
            msg = "Can't apply orderable window (e.g. `diff`, `shift`) to scalar-like expression."
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=self.has_windows,
            n_orderable_ops=self.n_orderable_ops + 1,
            preserves_length=self.preserves_length,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_ordered_over(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata:
        if self.has_windows:
            msg = "Cannot nest `over` statements."
            raise InvalidOperationError(msg)
        if self.is_elementwise or self.is_filtration:
            msg = (
                "Cannot use `over` on expressions which are elementwise\n"
                "(e.g. `abs`) or which change length (e.g. `drop_nulls`)."
            )
            raise InvalidOperationError(msg)
        n_orderable_ops = self.n_orderable_ops
        if (
            not n_orderable_ops
            and next(self.op_nodes_reversed()).kind is not ExprKind.WINDOW
        ):
            msg = (
                "Cannot use `order_by` in `over` on expression which isn't orderable.\n"
                "If your expression is orderable, then make sure that `over(order_by=...)`\n"
                "comes immediately after the order-dependent expression.\n\n"
                "Hint: instead of\n"
                "  - `(nw.col('price').diff() + 1).over(order_by='date')`\n"
                "write:\n"
                "  + `nw.col('price').diff().over(order_by='date') + 1`\n"
            )
            raise InvalidOperationError(msg)
        if next(self.op_nodes_reversed()).kind.is_orderable and n_orderable_ops > 0:
            n_orderable_ops -= 1
        return ExprMetadata(
            self.expansion_kind,
            has_windows=True,
            n_orderable_ops=n_orderable_ops,
            preserves_length=True,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_partitioned_over(
        self, node: ExprNode, _ce: CompliantExprAny
    ) -> ExprMetadata:
        if self.has_windows:
            msg = "Cannot nest `over` statements."
            raise InvalidOperationError(msg)
        if self.is_elementwise or self.is_filtration:
            msg = (
                "Cannot use `over` on expressions which are elementwise\n"
                "(e.g. `abs`) or which change length (e.g. `drop_nulls`)."
            )
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=True,
            n_orderable_ops=self.n_orderable_ops,
            preserves_length=True,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_over(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata:
        if node.kwargs["order_by"]:
            return self.with_ordered_over(node, _ce)
        if not node.kwargs["partition_by"]:  # pragma: no cover
            msg = "At least one of `partition_by` or `order_by` must be specified."
            raise InvalidOperationError(msg)
        return self.with_partitioned_over(node, _ce)

    def with_filtration(
        self, node: ExprNode, *compliant_exprs: CompliantExprAny
    ) -> ExprMetadata:
        if self.is_scalar_like:
            msg = "Can't apply filtration (e.g. `drop_nulls`) to scalar-like expression."
            raise InvalidOperationError(msg)
        result_has_windows = any(x._metadata.has_windows for x in compliant_exprs)
        result_n_orderable_ops = sum(x._metadata.n_orderable_ops for x in compliant_exprs)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=result_has_windows,
            n_orderable_ops=result_n_orderable_ops,
            preserves_length=False,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def with_orderable_filtration(
        self, node: ExprNode, _ce: CompliantExprAny
    ) -> ExprMetadata:
        if self.is_scalar_like:
            msg = "Can't apply filtration (e.g. `drop_nulls`) to scalar-like expression."
            raise InvalidOperationError(msg)
        return ExprMetadata(
            self.expansion_kind,
            has_windows=self.has_windows,
            n_orderable_ops=self.n_orderable_ops + 1,
            preserves_length=False,
            is_elementwise=False,
            is_scalar_like=False,
            is_literal=False,
            current_node=node,
            prev=self,
        )

    def op_nodes_reversed(self) -> Iterator[ExprNode]:
        for node in self.iter_nodes_reversed():
            if node.name.startswith(("name.", "alias")):
                # Skip nodes which only do aliasing.
                continue
            yield node


KIND_TO_METADATA_CONSTRUCTOR: dict[ExprKind, Callable[[ExprNode], ExprMetadata]] = {
    ExprKind.AGGREGATION: ExprMetadata.from_aggregation,
    ExprKind.ALL: ExprMetadata.from_selector_multi_unnamed,
    ExprKind.ELEMENTWISE: ExprMetadata.from_elementwise,
    ExprKind.EXCLUDE: ExprMetadata.from_selector_multi_unnamed,
    ExprKind.SERIES: ExprMetadata.from_series,
    ExprKind.COL: ExprMetadata.from_col,
    ExprKind.LITERAL: ExprMetadata.from_literal,
    ExprKind.NTH: ExprMetadata.from_nth,
    ExprKind.SELECTOR: ExprMetadata.from_selector_multi_unnamed,
}

KIND_TO_METADATA_UPDATER: dict[ExprKind, Callable[..., ExprMetadata]] = {
    ExprKind.AGGREGATION: ExprMetadata.with_aggregation,
    ExprKind.ELEMENTWISE: ExprMetadata.with_elementwise,
    ExprKind.FILTRATION: ExprMetadata.with_filtration,
    ExprKind.ORDERABLE_AGGREGATION: ExprMetadata.with_orderable_aggregation,
    ExprKind.ORDERABLE_FILTRATION: ExprMetadata.with_orderable_filtration,
    ExprKind.OVER: ExprMetadata.with_over,
    ExprKind.ORDERABLE_WINDOW: ExprMetadata.with_orderable_window,
    ExprKind.WINDOW: ExprMetadata.with_window,
}


def combine_metadata(
    *compliant_exprs: CompliantExprAny,
    to_single_output: bool,
    current_node: ExprNode,
    prev: ExprMetadata | None,
) -> ExprMetadata:
    """Combine metadata from `args`.

    Arguments:
        compliant_exprs: Expression arguments.
        to_single_output: Whether the result is always single-output, regardless
            of the inputs (e.g. `nw.sum_horizontal`).
        current_node: The current node being added.
        prev: ExprMetadata of previous node.
    """
    n_filtrations = 0
    result_expansion_kind = ExpansionKind.SINGLE
    result_has_windows = False
    result_n_orderable_ops = 0
    # result preserves length if at least one input does
    result_preserves_length = False
    # result is elementwise if all inputs are elementwise
    result_is_elementwise = True
    # result is scalar-like if all inputs are scalar-like
    result_is_scalar_like = True
    # result is literal if all inputs are literal
    result_is_literal = True

    for i, ce in enumerate(compliant_exprs):
        metadata = ce._metadata
        assert metadata is not None  # noqa: S101
        if metadata.expansion_kind.is_multi_output():
            expansion_kind = metadata.expansion_kind
            if not to_single_output:
                result_expansion_kind = (
                    result_expansion_kind & expansion_kind if i > 0 else expansion_kind
                )

        result_has_windows |= metadata.has_windows
        result_n_orderable_ops += metadata.n_orderable_ops
        result_preserves_length |= metadata.preserves_length
        result_is_elementwise &= metadata.is_elementwise
        result_is_scalar_like &= metadata.is_scalar_like
        result_is_literal &= metadata.is_literal
        n_filtrations += int(metadata.is_filtration)
    if n_filtrations > 1:
        msg = "Length-changing expressions can only be used in isolation, or followed by an aggregation"
        raise InvalidOperationError(msg)
    if result_preserves_length and n_filtrations:
        msg = "Cannot combine length-changing expressions with length-preserving ones or aggregations"
        raise InvalidOperationError(msg)
    return ExprMetadata(
        result_expansion_kind,
        has_windows=result_has_windows,
        n_orderable_ops=result_n_orderable_ops,
        preserves_length=result_preserves_length,
        is_elementwise=result_is_elementwise,
        is_scalar_like=result_is_scalar_like,
        is_literal=result_is_literal,
        current_node=current_node,
        prev=prev,
    )


def check_expressions_preserve_length(
    *args: CompliantExprAny, function_name: str
) -> None:
    # Raise if any argument in `args` isn't length-preserving.
    # For Series input, we don't raise (yet), we let such checks happen later,
    # as this function works lazily and so can't evaluate lengths.

    if not all(x._metadata.preserves_length for x in args):
        msg = f"Expressions which aggregate or change length cannot be passed to '{function_name}'."
        raise InvalidOperationError(msg)


def _parse_into_expr(
    arg: IntoExpr | NonNestedLiteral | _1DArray,
    *,
    str_as_lit: bool = False,
    backend: Any = None,
    allow_literal: bool = True,
) -> Expr:
    from narwhals.functions import col, lit, new_series

    if isinstance(arg, str) and not str_as_lit:
        return col(arg)
    if is_numpy_array_1d(arg):
        return new_series("", arg, backend=backend)._to_expr()
    if is_series(arg):
        return arg._to_expr()
    if is_expr(arg):
        return arg
    if not allow_literal:
        raise InvalidIntoExprError.from_invalid_type(type(arg))
    return lit(arg)


def evaluate_into_exprs(
    *exprs: IntoExpr | NonNestedLiteral | _1DArray,
    ns: CompliantNamespaceAny,
    str_as_lit: bool,
    allow_multi_output: bool,
) -> Iterator[CompliantExprAny]:
    for expr in exprs:
        ret = _parse_into_expr(
            expr, str_as_lit=str_as_lit, backend=ns._implementation
        )._to_compliant_expr(ns)
        if not allow_multi_output and ret._metadata.expansion_kind.is_multi_output():
            msg = "Multi-output expressions are not allowed in this context."
            raise MultiOutputExpressionError(msg)
        yield ret


def maybe_broadcast_ces(*compliant_exprs: CompliantExprAny) -> list[CompliantExprAny]:
    broadcast = any(not is_scalar_like(ce) for ce in compliant_exprs)
    results: list[CompliantExprAny] = []
    for compliant_expr in compliant_exprs:
        if broadcast and is_scalar_like(compliant_expr):
            _compliant_expr: CompliantExprAny = compliant_expr.broadcast()
            # Make sure to preserve metadata.
            _compliant_expr._opt_metadata = compliant_expr._metadata
            results.append(_compliant_expr)
        else:
            results.append(compliant_expr)
    return results


def evaluate_root_node(node: ExprNode, ns: CompliantNamespaceAny) -> CompliantExprAny:
    if node.name in {"col", "exclude"}:
        # There's too much potential for Sequence[str] vs str bugs, so we pass down
        # `names` positionally rather than as a sequence of strings.
        ce = getattr(ns, node.name)(*node.kwargs["names"])
        ces = []
    else:
        if "." in node.name:
            module, method = node.name.split(".")
            func = getattr(getattr(ns, module), method)
        else:
            func = getattr(ns, node.name)
        ces = maybe_broadcast_ces(
            *evaluate_into_exprs(
                *node.exprs,
                ns=ns,
                str_as_lit=node.str_as_lit,
                allow_multi_output=node.allow_multi_output,
            )
        )
        ce = cast("CompliantExprAny", func(*ces, **node.kwargs))
    md = ExprMetadata.from_node(node, *ces)
    ce._opt_metadata = md
    return ce


def evaluate_node(
    compliant_expr: CompliantExprAny, node: ExprNode, ns: CompliantNamespaceAny
) -> CompliantExprAny:
    md: ExprMetadata = compliant_expr._metadata
    compliant_expr, *compliant_expr_args = maybe_broadcast_ces(
        compliant_expr,
        *evaluate_into_exprs(
            *node.exprs,
            ns=ns,
            str_as_lit=node.str_as_lit,
            allow_multi_output=node.allow_multi_output,
        ),
    )
    md = md.with_node(node, compliant_expr, *compliant_expr_args)
    if "." in node.name:
        accessor, method = node.name.split(".")
        func = getattr(getattr(compliant_expr, accessor), method)
    else:
        func = getattr(compliant_expr, node.name)
    ret = cast("CompliantExprAny", func(*compliant_expr_args, **node.kwargs))
    ret._opt_metadata = md
    return ret


def evaluate_nodes(
    nodes: Sequence[ExprNode], ns: CompliantNamespaceAny
) -> CompliantExprAny:
    ce = evaluate_root_node(nodes[0], ns)
    for node in nodes[1:]:
        ce = evaluate_node(ce, node, ns)
    return ce