"""
Aggregation service for DataFrame operations.

This service provides aggregation operations using composition instead of mixin inheritance.
"""

from typing import TYPE_CHECKING, Any, Union, cast

from ...core.exceptions.operation import SparkColumnNotFoundError
from ...functions import Column, ColumnOperation, AggregateFunction
from ..grouped import GroupedData

if TYPE_CHECKING:
    from ..dataframe import DataFrame
    from ..protocols import SupportsDataFrameOps


class AggregationService:
    """Service providing aggregation operations for DataFrame."""

    def __init__(self, df: "DataFrame"):
        """Initialize aggregation service with DataFrame instance."""
        self._df = df

    def groupBy(self, *columns: Union[str, Column]) -> GroupedData:
        """Group DataFrame by columns for aggregation operations.

        Args:
            *columns: Column names or Column objects to group by.

        Returns:
            GroupedData for aggregation operations.

        Example:
            >>> df.groupBy("category").count()
            >>> df.groupBy("dept", "year").avg("salary")
        """
        col_names = []
        for col in columns:
            if isinstance(col, Column):
                col_names.append(col.name)
            else:
                col_names.append(col)

        # Validate that all columns exist
        for col_name in col_names:
            if col_name not in [field.name for field in self._df.schema.fields]:
                available_columns = [field.name for field in self._df.schema.fields]
                raise SparkColumnNotFoundError(col_name, available_columns)

        # Cast to SupportsDataFrameOps to satisfy type checker
        # DataFrame implements the protocol at runtime, but mypy can't verify
        # due to minor signature differences (e.g., approxQuantile accepts list[str])
        return GroupedData(cast("SupportsDataFrameOps", self._df), col_names)

    def groupby(self, *cols: Union[str, Column], **kwargs: Any) -> GroupedData:
        """Lowercase alias for groupBy() (all PySpark versions).

        Args:
            *cols: Column names or Column objects to group by
            **kwargs: Additional grouping options

        Returns:
            GroupedData object
        """
        return self.groupBy(*cols, **kwargs)

    def rollup(self, *columns: Union[str, Column]) -> Any:  # Returns RollupGroupedData
        """Create rollup grouped data for hierarchical grouping.

        Args:
            *columns: Columns to rollup.

        Returns:
            RollupGroupedData for hierarchical grouping.

        Example:
            >>> df.rollup("country", "state").sum("sales")
        """
        col_names = []
        for col in columns:
            if isinstance(col, Column):
                col_names.append(col.name)
            else:
                col_names.append(col)

        # Validate that all columns exist
        for col_name in col_names:
            if col_name not in [field.name for field in self._df.schema.fields]:
                available_columns = [field.name for field in self._df.schema.fields]
                raise SparkColumnNotFoundError(col_name, available_columns)

        from ..grouped.rollup import RollupGroupedData

        # Cast to SupportsDataFrameOps to satisfy type checker
        return RollupGroupedData(cast("SupportsDataFrameOps", self._df), col_names)

    def cube(self, *columns: Union[str, Column]) -> Any:  # Returns CubeGroupedData
        """Create cube grouped data for multi-dimensional grouping.

        Args:
            *columns: Columns to cube.

        Returns:
            CubeGroupedData for multi-dimensional grouping.

        Example:
            >>> df.cube("year", "month").sum("revenue")
        """
        col_names = []
        for col in columns:
            if isinstance(col, Column):
                col_names.append(col.name)
            else:
                col_names.append(col)

        # Validate that all columns exist
        for col_name in col_names:
            if col_name not in [field.name for field in self._df.schema.fields]:
                available_columns = [field.name for field in self._df.schema.fields]
                raise SparkColumnNotFoundError(col_name, available_columns)

        from ..grouped.cube import CubeGroupedData

        # Cast to SupportsDataFrameOps to satisfy type checker
        return CubeGroupedData(cast("SupportsDataFrameOps", self._df), col_names)

    def agg(
        self, *exprs: Union[str, Column, ColumnOperation, dict[str, str]]
    ) -> "SupportsDataFrameOps":
        """Aggregate DataFrame without grouping (global aggregation).

        Args:
            *exprs: Aggregation expressions, column names, or dictionary mapping
                   column names to aggregation functions.

        Returns:
            DataFrame with aggregated results.

        Example:
            >>> df.agg(F.max("age"), F.min("age"))
            >>> df.agg({"age": "max", "salary": "avg"})
        """
        # Handle dictionary syntax: {"col": "agg_func"}
        if len(exprs) == 1 and isinstance(exprs[0], dict):
            from ...functions import F

            agg_dict = exprs[0]
            converted_exprs: list[
                Union[str, Column, ColumnOperation, AggregateFunction]
            ] = []
            for col_name, agg_func in agg_dict.items():
                if agg_func == "sum":
                    converted_exprs.append(F.sum(col_name))
                elif agg_func == "avg" or agg_func == "mean":
                    converted_exprs.append(F.avg(col_name))
                elif agg_func == "max":
                    converted_exprs.append(F.max(col_name))
                elif agg_func == "min":
                    converted_exprs.append(F.min(col_name))
                elif agg_func == "count":
                    converted_exprs.append(F.count(col_name))
                elif agg_func == "stddev":
                    converted_exprs.append(F.stddev(col_name))
                elif agg_func == "variance":
                    converted_exprs.append(F.variance(col_name))
                else:
                    # Fallback to string expression
                    converted_exprs.append(f"{agg_func}({col_name})")
            # Create a grouped data object with empty group columns for global aggregation
            grouped = GroupedData(cast("SupportsDataFrameOps", self._df), [])
            return cast("SupportsDataFrameOps", grouped.agg(*converted_exprs))

        # Create a grouped data object with empty group columns for global aggregation
        grouped = GroupedData(cast("SupportsDataFrameOps", self._df), [])
        return cast("SupportsDataFrameOps", grouped.agg(*exprs))
