"""Typing stubs for PySpark-shaped DataFrame wrappers."""

from __future__ import annotations

from typing import Any
from typing import Literal as TypingLiteral

from pydantable.dataframe import DataFrame as CoreDataFrame
from pydantable.dataframe_model import DataFrameModel as CoreDataFrameModel
from pydantable.expressions import ColumnRef

class PySparkGroupedDataFrame:
    def pivot(
        self, pivot_col: str | ColumnRef, *, values: list[Any] | None = None
    ) -> PySparkPivotedGroupedDataFrame: ...
    def agg(
        self, *exprs: Any, streaming: bool | None = None, **aggregations: Any
    ) -> DataFrame: ...
    def sum(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def mean(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def min(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def max(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def count(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def len(self, *, streaming: bool | None = None) -> DataFrame: ...

class PySparkGroupedDataFrameModel:
    def pivot(
        self, pivot_col: str | ColumnRef, *, values: list[Any] | None = None
    ) -> PySparkPivotedGroupedDataFrameModel: ...
    def agg(self, *exprs: Any, **aggregations: Any) -> DataFrameModel: ...
    def sum(self, *columns: str, streaming: bool | None = None) -> DataFrameModel: ...
    def mean(self, *columns: str, streaming: bool | None = None) -> DataFrameModel: ...
    def min(self, *columns: str, streaming: bool | None = None) -> DataFrameModel: ...
    def max(self, *columns: str, streaming: bool | None = None) -> DataFrameModel: ...
    def count(self, *columns: str, streaming: bool | None = None) -> DataFrameModel: ...
    def len(self, *, streaming: bool | None = None) -> DataFrameModel: ...

class PySparkPivotedGroupedDataFrame:
    def agg(
        self,
        aggDict: Any | None = None,
        *,
        streaming: bool | None = None,
        **aggregations: Any,
    ) -> DataFrame: ...
    def count(self, *, streaming: bool | None = None) -> DataFrame: ...
    def sum(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def avg(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def mean(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def min(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def max(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...

class PySparkPivotedGroupedDataFrameModel:
    def agg(self, **aggregations: Any) -> DataFrame: ...
    def count(self, *, streaming: bool | None = None) -> DataFrame: ...
    def sum(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def avg(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def mean(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def min(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...
    def max(self, *columns: str, streaming: bool | None = None) -> DataFrame: ...

class DataFrameNaFunctions:
    def drop(
        self,
        how: str = "any",
        thresh: int | None = None,
        subset: str | list[str] | None = None,
    ) -> DataFrame: ...
    def fill(self, value: Any, subset: str | list[str] | None = None) -> DataFrame: ...

class DataFrameModelNaFunctions:
    def drop(
        self,
        how: str = "any",
        thresh: int | None = None,
        subset: str | list[str] | None = None,
    ) -> DataFrameModel: ...
    def fill(
        self, value: Any, subset: str | list[str] | None = None
    ) -> DataFrameModel: ...

class DataFrame(CoreDataFrame[Any]):
    def group_by(  # type: ignore[override]
        self,
        *keys: str | ColumnRef,
        maintain_order: bool = False,
        drop_nulls: bool = True,
    ) -> PySparkGroupedDataFrame: ...
    def groupBy(
        self,
        *keys: str | ColumnRef,
        maintain_order: bool = False,
        drop_nulls: bool = True,
    ) -> PySparkGroupedDataFrame: ...
    def sort(  # type: ignore[override]
        self,
        *columns: str,
        ascending: bool | list[bool] | None = None,
    ) -> DataFrame: ...
    def crossJoin(self, other: DataFrame) -> DataFrame: ...
    def count(self) -> int: ...
    def unionByName(
        self,
        other: DataFrame,
        *,
        allowMissingColumns: bool = False,
    ) -> DataFrame: ...
    def join(  # type: ignore[override]
        self,
        other: DataFrame,
        *,
        on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        left_on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        right_on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        how: str = "inner",
        suffix: str = "_right",
        coalesce: bool | None = None,
        validate: str | None = None,
        join_nulls: bool | None = None,
        maintain_order: bool | str | None = None,
        streaming: bool | None = None,
        keepRightJoinKeys: bool = False,
        keepLeftJoinKeys: bool = False,
    ) -> DataFrame: ...
    def intersect(self, other: DataFrame) -> DataFrame: ...
    def subtract(self, other: DataFrame) -> DataFrame: ...
    def except_(self, other: DataFrame) -> DataFrame: ...
    def exceptAll(self, other: DataFrame) -> DataFrame: ...
    def intersectAll(self, other: DataFrame) -> DataFrame: ...
    def fillna(
        self,
        value: Any = None,
        *,
        subset: str | list[str] | None = None,
    ) -> DataFrame: ...
    def dropna(
        self,
        how: str = "any",
        *,
        thresh: int | None = None,
        subset: str | list[str] | None = None,
    ) -> DataFrame: ...
    @property
    def na(self) -> DataFrameNaFunctions: ...
    def printSchema(self, level: int | None = None) -> None: ...
    def explain(  # type: ignore[override]
        self,
        extended: bool | None = None,
        mode: str | None = None,
        *,
        format: TypingLiteral["text", "json"] = "text",
    ) -> None: ...
    def describe(self) -> str: ...
    def toPandas(self) -> Any: ...
    def sample(  # type: ignore[override]
        self,
        withReplacement: bool | None = None,
        fraction: float | None = None,
        seed: int | None = None,
    ) -> DataFrame: ...
    def explode(
        self,
        column: Any,
        *,
        outer: bool = False,
        streaming: bool | None = None,
    ) -> DataFrame: ...
    def explode_outer(
        self, column: Any, *, streaming: bool | None = None
    ) -> DataFrame: ...
    def explode_all(self, *, streaming: bool | None = None) -> DataFrame: ...
    def posexplode(
        self,
        column: str,
        *,
        pos: str = ...,
        value: str | None = ...,
        outer: bool = ...,
        streaming: bool | None = ...,
    ) -> DataFrame: ...
    def posexplode_outer(
        self,
        column: str,
        *,
        pos: str = ...,
        value: str | None = ...,
        streaming: bool | None = ...,
    ) -> DataFrame: ...
    def unnest(self, column: Any, *, streaming: bool | None = None) -> DataFrame: ...
    def unnest_all(self, *, streaming: bool | None = None) -> DataFrame: ...

class DataFrameModel(CoreDataFrameModel):
    def group_by(  # type: ignore[override]
        self,
        *keys: Any,
        maintain_order: bool = False,
        drop_nulls: bool = True,
    ) -> PySparkGroupedDataFrameModel: ...
    def groupBy(
        self,
        *keys: Any,
        maintain_order: bool = False,
        drop_nulls: bool = True,
    ) -> PySparkGroupedDataFrameModel: ...
    def sort(  # type: ignore[override]
        self,
        *columns: str,
        ascending: bool | list[bool] | None = None,
    ) -> DataFrameModel: ...
    def crossJoin(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def count(self) -> int: ...
    def unionByName(
        self,
        other: DataFrameModel | DataFrame,
        *,
        allowMissingColumns: bool = False,
    ) -> DataFrameModel: ...
    def join(  # type: ignore[override]
        self,
        other: DataFrameModel | DataFrame,
        *,
        on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        left_on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        right_on: str
        | ColumnRef
        | list[str | ColumnRef]
        | tuple[str | ColumnRef, ...]
        | None = None,
        how: str = "inner",
        suffix: str = "_right",
        coalesce: bool | None = None,
        validate: str | None = None,
        join_nulls: bool | None = None,
        maintain_order: bool | str | None = None,
        streaming: bool | None = None,
        keepRightJoinKeys: bool = False,
        keepLeftJoinKeys: bool = False,
    ) -> DataFrameModel: ...
    def intersect(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def subtract(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def except_(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def exceptAll(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def intersectAll(self, other: DataFrameModel | DataFrame) -> DataFrameModel: ...
    def fillna(
        self,
        value: Any = None,
        *,
        subset: str | list[str] | None = None,
    ) -> DataFrameModel: ...
    def dropna(
        self,
        how: str = "any",
        *,
        thresh: int | None = None,
        subset: str | list[str] | None = None,
    ) -> DataFrameModel: ...
    @property
    def na(self) -> DataFrameModelNaFunctions: ...
    def printSchema(self, level: int | None = None) -> None: ...
    def explain(  # type: ignore[override]
        self,
        extended: bool | None = None,
        mode: str | None = None,
        *,
        format: TypingLiteral["text", "json"] = "text",
    ) -> None: ...
    def describe(self) -> str: ...
    def toPandas(self) -> Any: ...
    def sample(
        self,
        withReplacement: bool | None = None,
        fraction: float | None = None,
        seed: int | None = None,
    ) -> DataFrameModel: ...
    def explode(
        self,
        columns: Any,
        *,
        outer: bool = False,
        streaming: bool | None = None,
    ) -> DataFrameModel: ...
    def explode_outer(
        self, columns: Any, *, streaming: bool | None = None
    ) -> DataFrameModel: ...
    def explode_all(self, *, streaming: bool | None = None) -> DataFrameModel: ...
    def posexplode(
        self,
        column: str,
        *,
        pos: str = ...,
        value: str | None = ...,
        outer: bool = ...,
        streaming: bool | None = ...,
    ) -> DataFrameModel: ...
    def posexplode_outer(
        self,
        column: str,
        *,
        pos: str = ...,
        value: str | None = ...,
        streaming: bool | None = ...,
    ) -> DataFrameModel: ...
    def unnest(
        self, columns: Any, *, streaming: bool | None = None
    ) -> DataFrameModel: ...
    def unnest_all(self, *, streaming: bool | None = None) -> DataFrameModel: ...
