from __future__ import annotations

from typing import Any

from .window_spec import WindowSpec

__all__ = [
    "BinaryOp",
    "ColumnRef",
    "CompareOp",
    "Expr",
    "Literal",
    "WhenChain",
    "coalesce",
    "concat",
    "dense_rank",
    "global_count",
    "global_max",
    "global_mean",
    "global_min",
    "global_row_count",
    "global_sum",
    "lag",
    "lead",
    "rank",
    "row_number",
    "when",
    "window_max",
    "window_mean",
    "window_min",
    "window_sum",
]

_MAX_EXPR_REPR_AST = 200

def _rust_expr_ast_snippet(rust_expr: Any) -> str: ...

class Expr:
    _rust_expr: Any

    def __init__(self, *, rust_expr: Any): ...
    @property
    def dtype(self) -> Any: ...
    def referenced_columns(self) -> set[str]: ...
    def __repr__(self) -> str: ...
    def _coerce_other(self, other: Any) -> Expr: ...
    def _binary(self, op_symbol: str, other: Any) -> Expr: ...
    def _binary_reflected(self, op_symbol: str, other: Any) -> Expr: ...
    def _compare(self, op_symbol: str, other: Any) -> Expr: ...
    def cast(self, dtype: Any) -> Expr: ...
    def is_null(self) -> Expr: ...
    def is_not_null(self) -> Expr: ...
    def over(
        self,
        partition_by: str | list[str] | tuple[str, ...] | None = None,
        order_by: str | list[str] | tuple[str, ...] | None = None,
    ) -> Expr: ...
    def __add__(self, other: Any) -> Expr: ...
    def __sub__(self, other: Any) -> Expr: ...
    def __mul__(self, other: Any) -> Expr: ...
    def __truediv__(self, other: Any) -> Expr: ...
    def __radd__(self, other: Any) -> Expr: ...
    def __rsub__(self, other: Any) -> Expr: ...
    def __rmul__(self, other: Any) -> Expr: ...
    def __rtruediv__(self, other: Any) -> Expr: ...
    def __eq__(self, other: Any) -> Any: ...
    def __ne__(self, other: Any) -> Any: ...
    def __lt__(self, other: Any) -> Expr: ...
    def __le__(self, other: Any) -> Expr: ...
    def __gt__(self, other: Any) -> Expr: ...
    def __ge__(self, other: Any) -> Expr: ...
    def isin(self, *values: Any) -> Expr: ...
    def between(self, low: Any, high: Any) -> Expr: ...
    def substr(self, start: Any, length: Any | None = None) -> Expr: ...
    def char_length(self) -> Expr: ...
    def struct_field(self, name: str) -> Expr: ...
    def abs(self) -> Expr: ...
    def round(self, decimals: int = 0) -> Expr: ...
    def floor(self) -> Expr: ...
    def ceil(self) -> Expr: ...
    def strip(self) -> Expr: ...
    def upper(self) -> Expr: ...
    def lower(self) -> Expr: ...
    def str_replace(
        self, pattern: str, replacement: str, *, literal: bool = True
    ) -> Expr: ...
    def starts_with(self, prefix: str) -> Expr: ...
    def ends_with(self, suffix: str) -> Expr: ...
    def str_contains(self, substring: str) -> Expr: ...
    def str_contains_pat(self, pattern: str, *, literal: bool = False) -> Expr: ...
    def str_split(self, delimiter: str) -> Expr: ...
    def str_reverse(self) -> Expr: ...
    def str_pad_start(self, length: int, fill_char: str = " ") -> Expr: ...
    def str_pad_end(self, length: int, fill_char: str = " ") -> Expr: ...
    def str_zfill(self, length: int) -> Expr: ...
    def str_extract_regex(self, pattern: str, group_index: int = 1) -> Expr: ...
    def str_json_path_match(self, path: str) -> Expr: ...
    def strip_prefix(self, prefix: str) -> Expr: ...
    def strip_suffix(self, suffix: str) -> Expr: ...
    def strip_chars(self, chars: str) -> Expr: ...
    def __and__(self, other: Any) -> Expr: ...
    def __rand__(self, other: Any) -> Expr: ...
    def __or__(self, other: Any) -> Expr: ...
    def __ror__(self, other: Any) -> Expr: ...
    def __invert__(self) -> Expr: ...
    def dt_year(self) -> Expr: ...
    def dt_month(self) -> Expr: ...
    def dt_day(self) -> Expr: ...
    def dt_hour(self) -> Expr: ...
    def dt_minute(self) -> Expr: ...
    def dt_second(self) -> Expr: ...
    def dt_nanosecond(self) -> Expr: ...
    def dt_weekday(self) -> Expr: ...
    def dt_quarter(self) -> Expr: ...
    def dt_week(self) -> Expr: ...
    def dt_date(self) -> Expr: ...
    def strptime(self, format: str, *, to_datetime: bool = False) -> Expr: ...
    def unix_timestamp(self, unit: str = "seconds") -> Expr: ...
    def binary_len(self) -> Expr: ...
    def map_len(self) -> Expr: ...
    def map_get(self, key: str) -> Expr: ...
    def map_contains_key(self, key: str) -> Expr: ...
    def map_keys(self) -> Expr: ...
    def map_values(self) -> Expr: ...
    def map_entries(self) -> Expr: ...
    def map_from_entries(self) -> Expr: ...
    def element_at(self, key: str) -> Expr: ...
    def list_len(self) -> Expr: ...
    def list_get(self, index: Any) -> Expr: ...
    def list_contains(self, value: Any) -> Expr: ...
    def list_min(self) -> Expr: ...
    def list_max(self) -> Expr: ...
    def list_sum(self) -> Expr: ...
    def list_mean(self) -> Expr: ...
    def list_join(self, separator: str, *, ignore_nulls: bool = False) -> Expr: ...
    def list_sort(
        self,
        *,
        descending: bool = False,
        nulls_last: bool = False,
        maintain_order: bool = False,
    ) -> Expr: ...
    def list_unique(self, *, stable: bool = False) -> Expr: ...

class WhenChain:
    def __repr__(self) -> str: ...
    def __init__(self, condition: Expr, value: Expr): ...
    def when(self, condition: Expr, value: Expr) -> WhenChain: ...
    def otherwise(self, value: Expr) -> Expr: ...

def when(condition: Expr, value: Expr) -> WhenChain: ...

class ColumnRef(Expr):
    def __init__(self, *, name: str, dtype: Any): ...
    def __repr__(self) -> str: ...

class Literal(Expr):
    def __init__(self, *, value: Any, dtype: Any = None): ...

class BinaryOp(Expr):
    def __init__(self, *, rust_expr: Any): ...

class CompareOp(Expr):
    def __init__(self, *, rust_expr: Any): ...

def coalesce(*exprs: Expr) -> Expr: ...
def concat(*exprs: Expr) -> Expr: ...

class _WindowFnPending:
    def __init__(self, kind: str): ...
    def __repr__(self) -> str: ...
    def over(self, window: WindowSpec) -> Expr: ...

class _WindowAggPending:
    def __init__(self, inner: Expr, kind: str): ...
    def __repr__(self) -> str: ...
    def over(self, window: WindowSpec) -> Expr: ...

class _WindowShiftPending:
    def __init__(self, inner: Expr, kind: str, n: int): ...
    def __repr__(self) -> str: ...
    def over(self, window: WindowSpec) -> Expr: ...

def row_number() -> _WindowFnPending: ...
def rank() -> _WindowFnPending: ...
def dense_rank() -> _WindowFnPending: ...
def window_sum(column: Expr) -> _WindowAggPending: ...
def window_mean(column: Expr) -> _WindowAggPending: ...
def window_min(column: Expr) -> _WindowAggPending: ...
def window_max(column: Expr) -> _WindowAggPending: ...
def lag(column: Expr, n: int = 1) -> _WindowShiftPending: ...
def lead(column: Expr, n: int = 1) -> _WindowShiftPending: ...
def global_sum(column: Expr) -> Expr: ...
def global_mean(column: Expr) -> Expr: ...
def global_count(column: Expr) -> Expr: ...
def global_min(column: Expr) -> Expr: ...
def global_max(column: Expr) -> Expr: ...
def global_row_count() -> Expr: ...

__all__ = [
    "BinaryOp",
    "ColumnRef",
    "CompareOp",
    "Expr",
    "Literal",
    "WhenChain",
    "coalesce",
    "concat",
    "dense_rank",
    "global_count",
    "global_max",
    "global_mean",
    "global_min",
    "global_row_count",
    "global_sum",
    "lag",
    "lead",
    "rank",
    "row_number",
    "when",
    "window_max",
    "window_mean",
    "window_min",
    "window_sum",
]
