"""MongoDB aggregation pipeline visitor for therismos grouping.

This module provides a visitor that converts therismos grouping specifications into
MongoDB aggregation pipeline $group stages compatible with PyMongo and Motor.

The MongoDB visitor does not require PyMongo or Motor to be installed for the
conversion itself - it simply returns dictionaries in MongoDB $group format.
However, you'll need one of these libraries to actually execute the queries:

- PyMongo: Synchronous MongoDB driver
- Motor: Asynchronous MongoDB driver for asyncio

Example
-------
>>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
>>> from therismos.grouping.visitors.mongo import MongoVisitor
>>>
>>> spec = GroupSpec(
...     group_by=["category", "region"],
...     aggregations=[
...         Aggregation("total", AggregationFunction.COUNT),
...         Aggregation("min_price", AggregationFunction.MIN, "price"),
...         Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
...     ],
... )
>>>
>>> visitor = MongoVisitor()
>>> group_stage = spec.accept(visitor)
>>> # Result: {
>>> #     "$group": {
>>> #         "_id": {"category": "$category", "region": "$region"},
>>> #         "total": {"$sum": 1},
>>> #         "min_price": {"$min": "$price"},
>>> #         "avg_price": {"$avg": "$price"}
>>> #     }
>>> # }
>>>
>>> # Use with PyMongo
>>> # pipeline = [group_stage]
>>> # results = collection.aggregate(pipeline)
>>>
>>> # Use with Motor
>>> # pipeline = [group_stage]
>>> # results = await collection.aggregate(pipeline).to_list(length=None)

Notes
-----
- COUNT aggregation converts to {"$sum": 1}
- SUM aggregation converts to {"$sum": "$field"}
- MIN/MAX convert to {"$min": "$field"} / {"$max": "$field"}
- AVERAGE converts to {"$avg": "$field"}
- STDDEV converts to {"$stdDevPop": "$field"}
- MEDIAN and percentiles use $percentile operator (MongoDB 7.0+)
- The _id field contains the grouping fields as a composite key
- For single grouping field, _id can be simplified to "$field"
"""

from __future__ import annotations

from typing import Any

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


class MongoVisitor:
    """Visitor that converts grouping specifications to MongoDB $group stages.

    This visitor traverses grouping specifications and generates MongoDB
    aggregation pipeline $group stages that can be used with PyMongo's
    aggregate() or Motor's aggregate() methods.

    :ivar simplify_single_group: If True, simplify _id to "$field" for single grouping field.
    :vartype simplify_single_group: bool

    Example
    -------
    >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
    >>> from therismos.grouping.visitors.mongo import MongoVisitor
    >>>
    >>> spec = GroupSpec(
    ...     group_by=["status"],
    ...     aggregations=[
    ...         Aggregation("count", AggregationFunction.COUNT),
    ...         Aggregation("avg_age", AggregationFunction.AVERAGE, "age"),
    ...     ],
    ... )
    >>>
    >>> visitor = MongoVisitor(simplify_single_group=True)
    >>> group_stage = spec.accept(visitor)
    >>> print(group_stage)
    {
        "$group": {
            "_id": "$status",
            "count": {"$sum": 1},
            "avg_age": {"$avg": "$age"}
        }
    }
    """

    def __init__(self, simplify_single_group: bool = True) -> None:
        """Initialize the MongoDB visitor.

        :param simplify_single_group: If True, use "$field" for single grouping field
            instead of {"field": "$field"}.
        :type simplify_single_group: bool
        """
        self.simplify_single_group = simplify_single_group

    def _convert_aggregation(self, agg: Aggregation) -> dict[str, Any]:
        """Convert an aggregation to MongoDB accumulator format.

        :param agg: The aggregation to convert.
        :type agg: Aggregation
        :returns: MongoDB accumulator expression.
        :rtype: dict[str, Any]
        :raises ValueError: If the aggregation function is not supported.
        """
        func = agg.function
        field_ref = f"${agg.field}" if agg.field else None

        if func == AggregationFunction.COUNT:
            return {"$sum": 1}
        elif func == AggregationFunction.SUM:
            return {"$sum": field_ref}
        elif func == AggregationFunction.MIN:
            return {"$min": field_ref}
        elif func == AggregationFunction.MAX:
            return {"$max": field_ref}
        elif func == AggregationFunction.AVERAGE:
            return {"$avg": field_ref}
        elif func == AggregationFunction.STDDEV:
            # Use population standard deviation (sample stddev would be $stdDevSamp)
            return {"$stdDevPop": field_ref}
        elif func == AggregationFunction.MEDIAN:
            # MongoDB 7.0+ percentile operator
            return {
                "$median": {
                    "input": field_ref,
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.Q1:
            # First quartile = 25th percentile
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.25],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.Q3:
            # Third quartile = 75th percentile
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.75],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P01:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.01],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P05:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.05],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P10:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.10],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P90:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.90],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P95:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.95],
                    "method": "approximate",
                }
            }
        elif func == AggregationFunction.P99:
            return {
                "$percentile": {
                    "input": field_ref,
                    "p": [0.99],
                    "method": "approximate",
                }
            }
        else:
            msg = f"Unsupported aggregation function: {func}"
            raise ValueError(msg)

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

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        :returns: MongoDB aggregation pipeline $group stage.
        :rtype: dict[str, Any]
        """
        # Build the _id field for grouping
        if len(spec.group_by) == 0:
            # No grouping fields means group all documents together
            group_id: Any = None
        elif len(spec.group_by) == 1 and self.simplify_single_group:
            # Single field: use simple "$field" format
            group_id = f"${spec.group_by[0]}"
        else:
            # Multiple fields: use composite key format
            group_id = {field: f"${field}" for field in spec.group_by}

        # Build the accumulator fields
        accumulators = {
            agg.id: self._convert_aggregation(agg) for agg in spec.aggregations.values()
        }

        # Construct the $group stage
        return {
            "$group": {
                "_id": group_id,
                **accumulators,
            }
        }


__all__ = [
    "MongoVisitor",
]
