"""Tests for Polars grouping specification visitor."""

import pytest

pl = pytest.importorskip("polars")

from therismos.grouping import Aggregation, AggregationFunction, GroupSpec
from therismos.grouping.visitors.polars import PolarsGroupSpec, PolarsGroupSpecVisitor


@pytest.fixture()
def visitor() -> PolarsGroupSpecVisitor:
    """Return a PolarsGroupSpecVisitor instance."""
    return PolarsGroupSpecVisitor()


@pytest.fixture()
def df() -> "pl.DataFrame":
    """Return a sample DataFrame for integration tests."""
    return pl.DataFrame(
        {
            "category": ["A", "B", "A", "B", "A"],
            "region": ["east", "east", "west", "west", "east"],
            "price": [10.0, 20.0, 15.0, 25.0, 12.0],
            "revenue": [100.0, 200.0, 150.0, 250.0, 120.0],
        }
    )


class TestPolarsGroupSpecVisitorBasic:
    """Tests for basic PolarsGroupSpecVisitor functionality."""

    def test_result_is_polars_group_spec(self, visitor: PolarsGroupSpecVisitor) -> None:
        """Test that result is a PolarsGroupSpec."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert isinstance(result, PolarsGroupSpec)

    def test_group_by_preserved(self, visitor: PolarsGroupSpecVisitor) -> None:
        """Test that group_by fields are preserved correctly."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert result.group_by == ("category", "region")

    def test_agg_count(self, visitor: PolarsGroupSpecVisitor) -> None:
        """Test that COUNT produces a Polars expression."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert len(result.agg) == 1
        assert isinstance(result.agg[0], pl.Expr)

    def test_empty_group_by(self, visitor: PolarsGroupSpecVisitor) -> None:
        """Test with empty group_by."""
        spec = GroupSpec(
            group_by=[],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert result.group_by == ()


class TestPolarsGroupSpecVisitorAggregations:
    """Tests for all aggregation functions."""

    @pytest.mark.parametrize(
        "func",
        [
            AggregationFunction.SUM,
            AggregationFunction.MIN,
            AggregationFunction.MAX,
            AggregationFunction.AVERAGE,
            AggregationFunction.STDDEV,
            AggregationFunction.MEDIAN,
        ],
    )
    def test_field_aggregations(
        self, visitor: PolarsGroupSpecVisitor, func: AggregationFunction
    ) -> None:
        """Test that field-based aggregations produce Polars expressions."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("result", func, "price")],
        )
        result = spec.accept(visitor)
        assert len(result.agg) == 1
        assert isinstance(result.agg[0], pl.Expr)

    @pytest.mark.parametrize(
        "func",
        [
            AggregationFunction.Q1,
            AggregationFunction.Q3,
            AggregationFunction.P01,
            AggregationFunction.P05,
            AggregationFunction.P10,
            AggregationFunction.P90,
            AggregationFunction.P95,
            AggregationFunction.P99,
        ],
    )
    def test_percentile_aggregations(
        self, visitor: PolarsGroupSpecVisitor, func: AggregationFunction
    ) -> None:
        """Test that percentile aggregations produce Polars expressions."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("result", func, "price")],
        )
        result = spec.accept(visitor)
        assert len(result.agg) == 1
        assert isinstance(result.agg[0], pl.Expr)

    def test_multiple_aggregations(self, visitor: PolarsGroupSpecVisitor) -> None:
        """Test multiple aggregations in one spec."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
                Aggregation("max_revenue", AggregationFunction.MAX, "revenue"),
            ],
        )
        result = spec.accept(visitor)
        assert len(result.agg) == 3


class TestPolarsGroupSpecVisitorIntegration:
    """Integration tests for Polars group visitor with actual DataFrames."""

    def test_count_aggregation(self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame") -> None:
        """Test COUNT aggregation produces correct row counts."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        grp = spec.accept(visitor)
        result = df.group_by(list(grp.group_by)).agg(list(grp.agg)).sort("category")
        category_a = result.filter(pl.col("category") == "A")["total"][0]
        category_b = result.filter(pl.col("category") == "B")["total"][0]
        assert category_a == 3
        assert category_b == 2

    def test_sum_aggregation(self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame") -> None:
        """Test SUM aggregation produces correct sums."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total_price", AggregationFunction.SUM, "price")],
        )
        grp = spec.accept(visitor)
        result = df.group_by(list(grp.group_by)).agg(list(grp.agg)).sort("category")
        total_a = result.filter(pl.col("category") == "A")["total_price"][0]
        assert abs(total_a - 37.0) < 0.001

    def test_average_aggregation(self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame") -> None:
        """Test AVERAGE aggregation produces correct means."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("avg_price", AggregationFunction.AVERAGE, "price")],
        )
        grp = spec.accept(visitor)
        result = df.group_by(list(grp.group_by)).agg(list(grp.agg)).sort("category")
        avg_b = result.filter(pl.col("category") == "B")["avg_price"][0]
        assert abs(avg_b - 22.5) < 0.001

    def test_min_max_aggregations(
        self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame"
    ) -> None:
        """Test MIN and MAX aggregations."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("max_price", AggregationFunction.MAX, "price"),
            ],
        )
        grp = spec.accept(visitor)
        result = df.group_by(list(grp.group_by)).agg(list(grp.agg)).sort("category")
        row_a = result.filter(pl.col("category") == "A")
        assert row_a["min_price"][0] == 10.0
        assert row_a["max_price"][0] == 15.0

    def test_lazy_frame_group_by(self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame") -> None:
        """Test group_by on LazyFrame."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        grp = spec.accept(visitor)
        result = (
            df.lazy().group_by(list(grp.group_by)).agg(list(grp.agg)).collect().sort("category")
        )
        assert len(result) == 2

    def test_multi_group_by(self, visitor: PolarsGroupSpecVisitor, df: "pl.DataFrame") -> None:
        """Test multiple group_by fields."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        grp = spec.accept(visitor)
        result = df.group_by(list(grp.group_by)).agg(list(grp.agg))
        assert len(result) == 4
