"""Tests for grouping specification visitors."""

import pytest

from therismos.grouping import Aggregation, AggregationFunction, GroupSpec
from therismos.grouping.visitors import DictVisitor, FieldGathererVisitor, StringVisitor
from therismos.grouping.visitors.mongo import MongoVisitor


class TestStringVisitor:
    """Tests for StringVisitor."""

    def test_simple_group_spec(self):
        """Test converting a simple GroupSpec to string."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = StringVisitor()
        result = spec.accept(visitor)

        assert result == '("category", "total:count")'

    def test_multiple_fields_and_aggregations(self):
        """Test converting GroupSpec with multiple fields and aggregations."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
            ],
        )

        visitor = StringVisitor()
        result = spec.accept(visitor)

        expected = '("category,region", "total:count,min_price:min:price,avg_price:average:price")'
        assert result == expected

    def test_empty_group_by(self):
        """Test converting GroupSpec with no grouping fields."""
        spec = GroupSpec(
            group_by=[],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = StringVisitor()
        result = spec.accept(visitor)

        assert result == '("", "total:count")'

    def test_empty_aggregations(self):
        """Test converting GroupSpec with no aggregations."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[],
        )

        visitor = StringVisitor()
        result = spec.accept(visitor)

        assert result == '("category", "")'


class TestDictVisitor:
    """Tests for DictVisitor."""

    def test_simple_group_spec(self):
        """Test converting a simple GroupSpec to dictionary."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = DictVisitor()
        result = spec.accept(visitor)

        assert result == {
            "group_by": ["category"],
            "aggregations": [
                {"id": "total", "function": "count", "field": None},
            ],
        }

    def test_multiple_fields_and_aggregations(self):
        """Test converting GroupSpec with multiple fields and aggregations."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("max_price", AggregationFunction.MAX, "price"),
            ],
        )

        visitor = DictVisitor()
        result = spec.accept(visitor)

        assert result == {
            "group_by": ["category", "region"],
            "aggregations": [
                {"id": "total", "function": "count", "field": None},
                {"id": "min_price", "function": "min", "field": "price"},
                {"id": "max_price", "function": "max", "field": "price"},
            ],
        }

    def test_percentile_aggregations(self):
        """Test converting GroupSpec with percentile aggregations."""
        spec = GroupSpec(
            group_by=["service"],
            aggregations=[
                Aggregation("p95", AggregationFunction.P95, "latency"),
                Aggregation("p99", AggregationFunction.P99, "latency"),
            ],
        )

        visitor = DictVisitor()
        result = spec.accept(visitor)

        assert result["aggregations"] == [
            {"id": "p95", "function": "p95", "field": "latency"},
            {"id": "p99", "function": "p99", "field": "latency"},
        ]


class TestFieldGathererVisitor:
    """Tests for FieldGathererVisitor."""

    def test_simple_group_spec(self):
        """Test gathering fields from a simple GroupSpec."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = FieldGathererVisitor()
        spec.accept(visitor)

        assert visitor.field_names == {"category"}

    def test_multiple_fields_and_aggregations(self):
        """Test gathering fields from GroupSpec with multiple fields."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("max_revenue", AggregationFunction.MAX, "revenue"),
            ],
        )

        visitor = FieldGathererVisitor()
        spec.accept(visitor)

        assert visitor.field_names == {"category", "region", "price", "revenue"}

    def test_duplicate_fields(self):
        """Test that duplicate field names are collected only once."""
        spec = GroupSpec(
            group_by=["category", "region", "category"],
            aggregations=[
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("max_price", AggregationFunction.MAX, "price"),
            ],
        )

        visitor = FieldGathererVisitor()
        spec.accept(visitor)

        # Set should contain unique values
        assert visitor.field_names == {"category", "region", "price"}

    def test_empty_group_by(self):
        """Test gathering fields with no grouping fields."""
        spec = GroupSpec(
            group_by=[],
            aggregations=[
                Aggregation("total", AggregationFunction.COUNT),
                Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
            ],
        )

        visitor = FieldGathererVisitor()
        spec.accept(visitor)

        # Only aggregation field
        assert visitor.field_names == {"price"}


class TestMongoVisitor:
    """Tests for MongoVisitor."""

    def test_simple_count_aggregation(self):
        """Test converting simple COUNT aggregation to MongoDB format."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result == {
            "$group": {
                "_id": "$category",
                "total": {"$sum": 1},
            }
        }

    def test_multiple_grouping_fields(self):
        """Test MongoDB format with multiple grouping fields."""
        spec = GroupSpec(
            group_by=["category", "region"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result == {
            "$group": {
                "_id": {
                    "category": "$category",
                    "region": "$region",
                },
                "total": {"$sum": 1},
            }
        }

    def test_various_aggregation_functions(self):
        """Test various aggregation functions in MongoDB format."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("count", AggregationFunction.COUNT),
                Aggregation("total_revenue", AggregationFunction.SUM, "revenue"),
                Aggregation("min_price", AggregationFunction.MIN, "price"),
                Aggregation("max_price", AggregationFunction.MAX, "price"),
                Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
                Aggregation("stddev_price", AggregationFunction.STDDEV, "price"),
            ],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result["$group"]["count"] == {"$sum": 1}
        assert result["$group"]["total_revenue"] == {"$sum": "$revenue"}
        assert result["$group"]["min_price"] == {"$min": "$price"}
        assert result["$group"]["max_price"] == {"$max": "$price"}
        assert result["$group"]["avg_price"] == {"$avg": "$price"}
        assert result["$group"]["stddev_price"] == {"$stdDevPop": "$price"}

    def test_median_aggregation(self):
        """Test MEDIAN aggregation in MongoDB format."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("median_price", AggregationFunction.MEDIAN, "price")],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result["$group"]["median_price"] == {
            "$median": {
                "input": "$price",
                "method": "approximate",
            }
        }

    def test_quartile_aggregations(self):
        """Test Q1 and Q3 aggregations in MongoDB format."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[
                Aggregation("q1", AggregationFunction.Q1, "salary"),
                Aggregation("q3", AggregationFunction.Q3, "salary"),
            ],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result["$group"]["q1"] == {
            "$percentile": {
                "input": "$salary",
                "p": [0.25],
                "method": "approximate",
            }
        }
        assert result["$group"]["q3"] == {
            "$percentile": {
                "input": "$salary",
                "p": [0.75],
                "method": "approximate",
            }
        }

    @pytest.mark.parametrize(
        "func,percentile",
        [
            (AggregationFunction.P01, 0.01),
            (AggregationFunction.P05, 0.05),
            (AggregationFunction.P10, 0.10),
            (AggregationFunction.P90, 0.90),
            (AggregationFunction.P95, 0.95),
            (AggregationFunction.P99, 0.99),
        ],
    )
    def test_percentile_aggregations(self, func, percentile):
        """Test various percentile aggregations in MongoDB format."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("pct", func, "latency")],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result["$group"]["pct"] == {
            "$percentile": {
                "input": "$latency",
                "p": [percentile],
                "method": "approximate",
            }
        }

    def test_no_grouping_fields(self):
        """Test MongoDB format with no grouping fields (global aggregation)."""
        spec = GroupSpec(
            group_by=[],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = MongoVisitor()
        result = spec.accept(visitor)

        assert result == {
            "$group": {
                "_id": None,
                "total": {"$sum": 1},
            }
        }

    def test_simplify_single_group_disabled(self):
        """Test with simplify_single_group=False."""
        spec = GroupSpec(
            group_by=["category"],
            aggregations=[Aggregation("total", AggregationFunction.COUNT)],
        )

        visitor = MongoVisitor(simplify_single_group=False)
        result = spec.accept(visitor)

        # Should use composite format even for single field
        assert result == {
            "$group": {
                "_id": {"category": "$category"},
                "total": {"$sum": 1},
            }
        }
