"""Visitor implementations for grouping specifications.

This module provides example visitors for converting grouping specifications to
various formats and for analyzing them.
"""

from __future__ import annotations

from typing import Any

from therismos.grouping._grouping import GroupSpec


class StringVisitor:
    """Visitor that converts grouping specifications to string representation.

    Produces the serialization format: ("field1,field2", "agg1:count,agg2:min:field")

    Example:
        >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
        >>> 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()
        >>> spec.accept(visitor)
        '("category,region", "total:count,min_price:min:price,avg_price:average:price")'
    """

    def visit_group_spec(self, spec: GroupSpec) -> str:
        """Convert a grouping specification to string.

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        :returns: String representation in the format ("fields", "aggregations").
        :rtype: str
        """
        return str(spec)


class DictVisitor:
    """Visitor that converts grouping specifications to dictionary representation.

    Example:
        >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
        >>> spec = GroupSpec(
        ...     group_by=["category", "region"],
        ...     aggregations=[
        ...         Aggregation("total", AggregationFunction.COUNT),
        ...         Aggregation("min_price", AggregationFunction.MIN, "price"),
        ...     ],
        ... )
        >>> visitor = DictVisitor()
        >>> spec.accept(visitor)
        {
            'group_by': ['category', 'region'],
            'aggregations': [
                {'id': 'total', 'function': 'count', 'field': None},
                {'id': 'min_price', 'function': 'min', 'field': 'price'}
            ]
        }
    """

    def visit_group_spec(self, spec: GroupSpec) -> dict[str, Any]:
        """Convert a grouping specification to dictionary.

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        :returns: Dictionary representation of the specification.
        :rtype: dict[str, Any]
        """
        return {
            "group_by": list(spec.group_by),
            "aggregations": [
                {
                    "id": agg.id,
                    "function": agg.function.value,
                    "field": agg.field,
                }
                for agg in spec.aggregations.values()
            ],
        }


class FieldGathererVisitor:
    """Visitor that collects the names of all fields used in a grouping specification.

    This includes both grouping fields and fields referenced in aggregations.

    Example:
        >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
        >>> spec = GroupSpec(
        ...     group_by=["category", "region"],
        ...     aggregations=[
        ...         Aggregation("total", AggregationFunction.COUNT),
        ...         Aggregation("min_price", AggregationFunction.MIN, "price"),
        ...         Aggregation("max_price", AggregationFunction.MAX, "price"),
        ...         Aggregation("avg_revenue", AggregationFunction.AVERAGE, "revenue"),
        ...     ],
        ... )
        >>> visitor = FieldGathererVisitor()
        >>> spec.accept(visitor)
        >>> sorted(visitor.field_names)
        ['category', 'price', 'region', 'revenue']
    """

    def __init__(self) -> None:
        """Initialize the visitor with an empty set of field names."""
        self.field_names: set[str] = set()

    def visit_group_spec(self, spec: GroupSpec) -> None:
        """Collect field names from a grouping specification.

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        """
        # Add grouping fields
        self.field_names.update(spec.group_by)

        # Add aggregation fields
        for agg in spec.aggregations.values():
            if agg.field is not None:
                self.field_names.add(agg.field)


__all__ = [
    "StringVisitor",
    "DictVisitor",
    "FieldGathererVisitor",
]
