"""Tests for Pandas grouping specification visitor."""

import pytest

pd = pytest.importorskip("pandas")

from therismos.grouping import Aggregation, AggregationFunction, GroupSpec
from therismos.grouping.visitors.pandas import PandasGroupSpec, PandasGroupSpecVisitor


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


@pytest.fixture()
def df() -> "pd.DataFrame":
    """Return a sample DataFrame for integration tests."""
    return pd.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 TestPandasGroupSpecVisitorBasic:
    """Tests for basic PandasGroupSpecVisitor functionality."""

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

    def test_group_by_preserved(self, visitor: PandasGroupSpecVisitor) -> 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_dict_keys(self, visitor: PandasGroupSpecVisitor) -> None:
        """Test that agg dict keys match aggregation IDs."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
            ],
        )
        result = spec.accept(visitor)
        assert "total" in result.agg
        assert "avg_price" in result.agg

    def test_agg_values_are_named_agg(self, visitor: PandasGroupSpecVisitor) -> None:
        """Test that agg values are pd.NamedAgg instances."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert isinstance(result.agg["total"], pd.NamedAgg)


class TestPandasGroupSpecVisitorAggregations:
    """Tests for all aggregation function conversions."""

    @pytest.mark.parametrize(
        "func,expected_aggfunc",
        [
            (AggregationFunction.SUM, "sum"),
            (AggregationFunction.MIN, "min"),
            (AggregationFunction.MAX, "max"),
            (AggregationFunction.AVERAGE, "mean"),
            (AggregationFunction.STDDEV, "std"),
            (AggregationFunction.MEDIAN, "median"),
        ],
    )
    def test_simple_aggfuncs(
        self,
        visitor: PandasGroupSpecVisitor,
        func: AggregationFunction,
        expected_aggfunc: str,
    ) -> None:
        """Test that simple aggregations map to correct aggfunc strings."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("result", func, "price")],
        )
        result = spec.accept(visitor)
        assert result.agg["result"].aggfunc == expected_aggfunc
        assert result.agg["result"].column == "price"

    @pytest.mark.parametrize(
        "func",
        [
            AggregationFunction.Q1,
            AggregationFunction.Q3,
            AggregationFunction.P01,
            AggregationFunction.P05,
            AggregationFunction.P10,
            AggregationFunction.P90,
            AggregationFunction.P95,
            AggregationFunction.P99,
        ],
    )
    def test_percentile_aggfuncs_are_callable(
        self, visitor: PandasGroupSpecVisitor, func: AggregationFunction
    ) -> None:
        """Test that percentile aggregations use callable aggfunc."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("result", func, "price")],
        )
        result = spec.accept(visitor)
        assert callable(result.agg["result"].aggfunc)
        assert result.agg["result"].column == "price"

    def test_count_uses_first_group_field(self, visitor: PandasGroupSpecVisitor) -> None:
        """Test that COUNT uses the first group_by field as the column."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        result = spec.accept(visitor)
        assert result.agg["total"].column == "category"
        assert result.agg["total"].aggfunc == "count"

    def test_count_raises_without_group_by(self, visitor: PandasGroupSpecVisitor) -> None:
        """Test that COUNT raises ValueError when group_by is empty."""
        spec = GroupSpec(
            group_by=[],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )
        with pytest.raises(ValueError, match="COUNT aggregation requires at least one group_by"):
            spec.accept(visitor)


class TestPandasGroupSpecVisitorIntegration:
    """Integration tests for Pandas group visitor with actual DataFrames."""

    def test_count_aggregation(self, visitor: PandasGroupSpecVisitor, df: "pd.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.groupby(list(grp.group_by)).agg(**grp.agg).sort_index()
        assert result.loc["A", "total"] == 3
        assert result.loc["B", "total"] == 2

    def test_sum_aggregation(self, visitor: PandasGroupSpecVisitor, df: "pd.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.groupby(list(grp.group_by)).agg(**grp.agg).sort_index()
        assert abs(result.loc["A", "total_price"] - 37.0) < 0.001

    def test_average_aggregation(self, visitor: PandasGroupSpecVisitor, df: "pd.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.groupby(list(grp.group_by)).agg(**grp.agg).sort_index()
        assert abs(result.loc["B", "avg_price"] - 22.5) < 0.001

    def test_min_max_aggregations(
        self, visitor: PandasGroupSpecVisitor, df: "pd.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.groupby(list(grp.group_by)).agg(**grp.agg).sort_index()
        assert result.loc["A", "min_price"] == 10.0
        assert result.loc["A", "max_price"] == 15.0

    def test_percentile_aggregation(
        self, visitor: PandasGroupSpecVisitor, df: "pd.DataFrame"
    ) -> None:
        """Test Q1 and Q3 percentile aggregations produce numeric results."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("q1_price", AggregationFunction.Q1, "price"),
                Aggregation("q3_price", AggregationFunction.Q3, "price"),
            ],
        )
        grp = spec.accept(visitor)
        result = df.groupby(list(grp.group_by)).agg(**grp.agg)
        assert "q1_price" in result.columns
        assert "q3_price" in result.columns
        row_a = result.loc["A"]
        assert row_a["q1_price"] <= row_a["q3_price"]

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

    def test_multiple_aggregations_combined(
        self, visitor: PandasGroupSpecVisitor, df: "pd.DataFrame"
    ) -> None:
        """Test multiple aggregation functions 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"),
            ],
        )
        grp = spec.accept(visitor)
        result = df.groupby(list(grp.group_by)).agg(**grp.agg).sort_index()
        assert "total" in result.columns
        assert "avg_price" in result.columns
        assert "max_revenue" in result.columns
        assert result.loc["A", "total"] == 3
