"""
Protocol interfaces for DataFrame handlers.

This module defines type-safe protocols for the specialized handler classes
that implement the Single Responsibility Principle for DataFrame operations.
"""

from collections.abc import Iterator
from typing import TYPE_CHECKING, Any, Optional, Protocol, Union

try:  # Python <3.11 compatibility
    from typing import Self  # type: ignore[attr-defined]
except ImportError:  # pragma: no cover - fallback for older versions
    from typing_extensions import Self

if TYPE_CHECKING:
    from ..spark_types import Row, StructType


class HasSchema(Protocol):
    """Objects that expose a schema and can project pending operations."""

    _schema: "StructType"

    @property
    def schema(self) -> "StructType": ...

    @schema.setter
    def schema(self, value: "StructType") -> None: ...

    def _project_schema_with_operations(self) -> "StructType": ...


class HasData(Protocol):
    """Objects that carry in-memory row data."""

    data: list[dict[str, Any]]


class HasStorage(Protocol):
    """Objects that reference a storage backend."""

    storage: Any


class OperationQueueAware(Protocol):
    """Objects that queue operations prior to materialisation."""

    _operations_queue: list[tuple[str, Any]]

    def _queue_op(self, operation: str, payload: Any) -> Self: ...


class LazyMaterializable(Protocol):
    """Objects that can defer execution until explicitly materialised."""

    _cached_count: Optional[int]

    def _materialize_if_lazy(self) -> Self: ...


class ColumnAware(Protocol):
    """Objects that expose column metadata."""

    @property
    def columns(self) -> list[str]: ...


class HandlerProvider(Protocol):
    """Objects that vend auxiliary handlers used by mixins."""

    def _get_collection_handler(self) -> "CollectionHandler": ...

    def _get_validation_handler(self) -> "ValidationHandler": ...

    def _get_condition_handler(self) -> "ConditionHandler": ...

    def _get_lazy_engine(self) -> Any: ...


class CollectionSupport(Protocol):
    """Objects exposing collection-style helpers."""

    def collect(self) -> list[Any]: ...

    def orderBy(self, *columns: Any, **kwargs: Any) -> Self: ...

    def count(self) -> int: ...


class LogicalOpSupport(Protocol):
    """Objects that provide logical dataframe operations used by mixins."""

    def select(self, *columns: Any, **kwargs: Any) -> Self: ...

    def filter(self, condition: Any) -> Self: ...

    def distinct(self) -> Self: ...

    def dropDuplicates(self, subset: Optional[list[str]] = None) -> Self: ...

    def drop_duplicates(self, subset: Optional[list[str]] = None) -> Self: ...

    def union(self, other: Any) -> Self: ...

    def groupBy(self, *columns: Any, **kwargs: Any) -> Any: ...


class MutationSupport(Protocol):
    """Objects that provide mutation and advanced DataFrame operations."""

    def withColumn(self, col_name: str, col: Any) -> Self: ...

    def withColumns(self, cols_map: dict[str, Any]) -> Self: ...

    def withColumnRenamed(self, existing: str, new: str) -> Self: ...

    def withColumnsRenamed(self, cols_map: dict[str, str]) -> Self: ...

    def drop(self, *cols: str) -> Self: ...

    def replace(
        self, to_replace: Any, value: Any = ..., subset: Optional[list[str]] = None
    ) -> Self: ...

    def sample(
        self,
        fraction: float,
        seed: Optional[int] = None,
        withReplacement: bool = False,
    ) -> Self: ...

    def sampleBy(
        self, col: str, fractions: dict[Any, float], seed: Optional[int] = None
    ) -> Self: ...

    def randomSplit(
        self, weights: list[float], seed: Optional[int] = None
    ) -> list[Self]: ...

    def describe(self, *cols: str) -> Self: ...

    def summary(self, *stats: str) -> Self: ...

    def crosstab(self, col1: str, col2: str) -> Self: ...

    def freqItems(self, cols: list[str], support: Optional[float] = None) -> Self: ...

    def approxQuantile(
        self, col: Any, probabilities: list[float], relativeError: float
    ) -> Any: ...

    def cov(self, col1: str, col2: str) -> float: ...

    def transform(self, func: Any) -> Self: ...

    def mapPartitions(self, func: Any, preservesPartitioning: bool = False) -> Self: ...

    def mapInPandas(self, func: Any, schema: Any) -> Self: ...

    def unpivot(
        self,
        ids: Union[str, list[str]],
        values: Union[str, list[str]],
        variableColumnName: str = "variable",
        valueColumnName: str = "value",
    ) -> Self: ...

    def melt(
        self,
        ids: Optional[list[str]] = None,
        values: Optional[list[str]] = None,
        variableColumnName: str = "variable",
        valueColumnName: str = "value",
    ) -> Self: ...

    def toDF(self, *cols: str) -> Self: ...

    def alias(self, alias: str) -> Self: ...

    def hint(self, name: str, *parameters: Any) -> Self: ...

    def withWatermark(self, eventTime: str, delayThreshold: str) -> Self: ...

    def toJSON(self) -> Self: ...

    def checkpoint(self, eager: bool = False) -> Self: ...

    def localCheckpoint(self, eager: bool = True) -> Self: ...

    def limit(self, n: int) -> Self: ...

    def offset(self, n: int) -> Self: ...

    def repartition(self, numPartitions: int, *cols: Any) -> Self: ...

    def coalesce(self, numPartitions: int) -> Self: ...

    def repartitionByRange(self, numPartitions: Any, *cols: Any) -> Self: ...

    def sortWithinPartitions(self, *cols: Any, **kwargs: Any) -> Self: ...

    def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator["Row"]: ...

    def foreach(self, f: Any) -> None: ...

    def foreachPartition(self, f: Any) -> None: ...

    def writeTo(self, table: str) -> Any: ...

    def sameSemantics(self, other: Any) -> bool: ...

    def semanticHash(self) -> int: ...

    def inputFiles(self) -> list[str]: ...

    def isLocal(self) -> bool: ...

    @property
    def isStreaming(self) -> bool: ...


class WatermarkAware(Protocol):
    """Objects that keep watermark metadata for streaming semantics."""

    _watermark_col: Optional[str]
    _watermark_delay: Optional[str]


class SupportsDataFrameOps(
    HasSchema,
    HasData,
    HasStorage,
    OperationQueueAware,
    LazyMaterializable,
    ColumnAware,
    HandlerProvider,
    CollectionSupport,
    LogicalOpSupport,
    MutationSupport,
    WatermarkAware,
    Protocol,
):
    """Composite protocol satisfied by `mock_spark.dataframe.DataFrame`."""

    def _validate_expression_columns(
        self,
        expression: Any,
        operation: str,
        in_lazy_materialization: bool = False,
    ) -> None: ...

    def _validate_filter_expression(
        self,
        condition: Any,
        operation: str,
        has_pending_joins: bool = False,
    ) -> None: ...

    def _validate_column_exists(
        self, column_name: str, operation: str, allow_ambiguous: bool = False
    ) -> None: ...

    def _evaluate_column_expression(
        self,
        row: dict[str, Any],
        column_expression: Any,
    ) -> Any: ...


class WindowFunctionHandler(Protocol):
    """Protocol for window function evaluation."""

    def evaluate_window_functions(
        self, data: list[dict[str, Any]], window_functions: list[tuple[Any, ...]]
    ) -> list[dict[str, Any]]:
        """Evaluate window functions across all rows."""
        ...

    def _evaluate_lag_lead(
        self,
        result_data: list[dict[str, Any]],
        window_func: Any,
        col_name: str,
        is_lead: bool,
    ) -> None:
        """Evaluate LAG/LEAD functions."""
        ...

    def _apply_ordering_to_indices(
        self, data: list[dict[str, Any]], indices: list[int], order_by_cols: list[Any]
    ) -> list[int]:
        """Apply ordering to row indices."""
        ...

    def _apply_lag_lead_to_partition(
        self,
        result_data: list[dict[str, Any]],
        partition_indices: list[int],
        window_func: Any,
        col_name: str,
        is_lead: bool,
    ) -> None:
        """Apply LAG/LEAD to a partition."""
        ...

    def _evaluate_rank_functions(
        self, result_data: list[dict[str, Any]], window_func: Any, col_name: str
    ) -> None:
        """Evaluate RANK/DENSE_RANK functions."""
        ...

    def _apply_rank_to_partition(
        self,
        result_data: list[dict[str, Any]],
        partition_indices: list[int],
        window_func: Any,
        col_name: str,
    ) -> None:
        """Apply ranking to a partition."""
        ...

    def _evaluate_aggregate_window_functions(
        self, result_data: list[dict[str, Any]], window_func: Any, col_name: str
    ) -> None:
        """Evaluate aggregate window functions (SUM, AVG, etc.)."""
        ...

    def _apply_aggregate_to_partition(
        self,
        result_data: list[dict[str, Any]],
        partition_indices: list[int],
        window_func: Any,
        col_name: str,
    ) -> None:
        """Apply aggregate functions to a partition."""
        ...


class CollectionHandler(Protocol):
    """Protocol for collection operations."""

    def collect(self, data: list[dict[str, Any]], schema: "StructType") -> list["Row"]:
        """Convert data to Row objects."""
        ...

    def take(
        self, data: list[dict[str, Any]], schema: "StructType", n: int
    ) -> list["Row"]:
        """Take first n rows."""
        ...

    def head(self, data: list[dict[str, Any]], schema: "StructType", n: int = 1) -> Any:
        """Get first row(s)."""
        ...

    def tail(self, data: list[dict[str, Any]], schema: "StructType", n: int = 1) -> Any:
        """Get last n rows."""
        ...

    def to_local_iterator(
        self,
        data: list[dict[str, Any]],
        schema: "StructType",
        prefetch: bool = False,
    ) -> Iterator["Row"]:
        """Return iterator over rows."""
        ...


class ValidationHandler(Protocol):
    """Protocol for data validation."""

    def validate_column_exists(
        self, schema: "StructType", column_name: str, operation: str
    ) -> None:
        """Validate single column exists."""
        ...

    def validate_columns_exist(
        self, schema: "StructType", column_names: list[str], operation: str
    ) -> None:
        """Validate multiple columns exist."""
        ...

    def validate_filter_expression(
        self,
        schema: "StructType",
        condition: Any,
        operation: str,
        has_pending_joins: bool = False,
    ) -> None:
        """Validate filter expression."""
        ...

    def validate_expression_columns(
        self,
        schema: "StructType",
        expression: Any,
        operation: str,
        in_lazy_materialization: bool = False,
    ) -> None:
        """Validate columns in expression exist."""
        ...


class ConditionHandler(Protocol):
    """Protocol for condition evaluation."""

    def apply_condition(
        self, data: list[dict[str, Any]], condition: Any
    ) -> list[dict[str, Any]]:
        """Filter data based on condition."""
        ...

    def evaluate_condition(self, row: dict[str, Any], condition: Any) -> bool:
        """Evaluate condition for a single row."""
        ...

    def evaluate_column_expression(
        self, row: dict[str, Any], column_expression: Any
    ) -> Any:
        """Evaluate column expression."""
        ...

    def evaluate_case_when(self, row: dict[str, Any], case_when_obj: Any) -> Any:
        """Evaluate CASE WHEN expression."""
        ...

    def _evaluate_case_when_condition(
        self, row: dict[str, Any], condition: Any
    ) -> bool:
        """Helper for case when condition evaluation."""
        ...
