"""Polars grouping specification visitor for therismos grouping.

This module provides a visitor that converts therismos grouping specifications into
Polars-compatible group_by and aggregation parameters via the PolarsGroupSpec dataclass.

Requires Polars to be installed. Install with: pip install therismos[polars]

Example
-------
>>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
>>> from therismos.grouping.visitors.polars import PolarsGroupSpecVisitor
>>>
>>> spec = GroupSpec(
...     group_by=["category", "region"],
...     aggregations=[
...         Aggregation("total", AggregationFunction.COUNT),
...         Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
...     ],
... )
>>>
>>> visitor = PolarsGroupSpecVisitor()
>>> grp = spec.accept(visitor)
>>>
>>> # Use with eager DataFrame
>>> # df.group_by(list(grp.group_by)).agg(list(grp.agg))
>>>
>>> # Use with lazy LazyFrame
>>> # lf.group_by(list(grp.group_by)).agg(list(grp.agg))

Notes
-----
- COUNT uses pl.len() to count all rows in the group
- Percentile functions use pl.col(field).quantile(value)
- PolarsGroupSpec is not frozen because pl.Expr is not hashable
"""

from __future__ import annotations

from dataclasses import dataclass, field

try:
    import polars as pl
except ImportError as exc:
    msg = "polars is required. Install with: pip install therismos[polars]"
    raise ImportError(msg) from exc

from therismos.grouping._grouping import Aggregation, AggregationFunction, GroupSpec


@dataclass
class PolarsGroupSpec:
    """Dataclass holding Polars group_by and aggregation parameters.

    Not frozen because pl.Expr is not hashable.

    :ivar group_by: Tuple of field names to group by.
    :vartype group_by: tuple[str, ...]
    :ivar agg: Tuple of Polars expressions for aggregation.
    :vartype agg: tuple[pl.Expr, ...]

    Example
    -------
    >>> grp = PolarsGroupSpec(
    ...     group_by=("category",),
    ...     agg=(pl.len().alias("total"), pl.col("price").mean().alias("avg_price")),
    ... )
    >>> df.group_by(list(grp.group_by)).agg(list(grp.agg))
    """

    group_by: tuple[str, ...] = field(default_factory=tuple)
    agg: tuple[pl.Expr, ...] = field(default_factory=tuple)


class PolarsGroupSpecVisitor:
    """Visitor that converts grouping specifications to Polars group_by parameters.

    This visitor traverses grouping specifications and generates a PolarsGroupSpec
    dataclass that can be passed to DataFrame.group_by().agg() or LazyFrame equivalents.

    Example
    -------
    >>> import polars as pl
    >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
    >>> from therismos.grouping.visitors.polars import PolarsGroupSpecVisitor
    >>>
    >>> spec = GroupSpec(
    ...     group_by=["status"],
    ...     aggregations=[
    ...         Aggregation("count", AggregationFunction.COUNT),
    ...         Aggregation("avg_age", AggregationFunction.AVERAGE, "age"),
    ...     ],
    ... )
    >>>
    >>> visitor = PolarsGroupSpecVisitor()
    >>> grp = spec.accept(visitor)
    >>> df = pl.DataFrame({"status": ["a", "b", "a"], "age": [20, 30, 25]})
    >>> df.group_by(list(grp.group_by)).agg(list(grp.agg))
    """

    def _convert_aggregation(self, agg: Aggregation) -> pl.Expr:
        """Convert a single aggregation to a Polars expression.

        :param agg: The aggregation to convert.
        :type agg: Aggregation
        :returns: Polars aggregation expression.
        :rtype: pl.Expr
        :raises ValueError: If the aggregation function is not supported.
        """
        func = agg.function
        agg_id = agg.id

        match func:
            case AggregationFunction.COUNT:
                return pl.len().alias(agg_id)
            case AggregationFunction.SUM:
                return pl.col(agg.field).sum().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.MIN:
                return pl.col(agg.field).min().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.MAX:
                return pl.col(agg.field).max().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.AVERAGE:
                return pl.col(agg.field).mean().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.STDDEV:
                return pl.col(agg.field).std().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.MEDIAN:
                return pl.col(agg.field).median().alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.Q1:
                return pl.col(agg.field).quantile(0.25).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.Q3:
                return pl.col(agg.field).quantile(0.75).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P01:
                return pl.col(agg.field).quantile(0.01).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P05:
                return pl.col(agg.field).quantile(0.05).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P10:
                return pl.col(agg.field).quantile(0.10).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P90:
                return pl.col(agg.field).quantile(0.90).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P95:
                return pl.col(agg.field).quantile(0.95).alias(agg_id)  # type: ignore[arg-type]
            case AggregationFunction.P99:
                return pl.col(agg.field).quantile(0.99).alias(agg_id)  # type: ignore[arg-type]
            case _:
                msg = f"Unsupported aggregation function: {func}"
                raise ValueError(msg)

    def visit_group_spec(self, spec: GroupSpec) -> PolarsGroupSpec:
        """Convert a grouping specification to Polars group_by parameters.

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        :returns: PolarsGroupSpec with group_by fields and aggregation expressions.
        :rtype: PolarsGroupSpec
        """
        agg_exprs = tuple(self._convert_aggregation(agg) for agg in spec.aggregations.values())
        return PolarsGroupSpec(group_by=spec.group_by, agg=agg_exprs)


__all__ = [
    "PolarsGroupSpec",
    "PolarsGroupSpecVisitor",
]
