"""Core grouping and aggregation module for therismos library.

This module provides classes for modeling SQL-like GROUP BY operations
with aggregation functions.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum
from typing import Generic, Protocol, TypeVar

T = TypeVar("T", covariant=True)


class AggregationFunction(StrEnum):
    """Enumeration of supported aggregation functions.

    :ivar COUNT: Count of items in the group.
    :ivar SUM: Sum of values in the group.
    :ivar MIN: Minimum value in the group.
    :ivar MAX: Maximum value in the group.
    :ivar AVERAGE: Average (mean) value in the group.
    :ivar STDDEV: Standard deviation of values in the group.
    :ivar MEDIAN: Median value in the group.
    :ivar Q1: First quartile (25th percentile) value in the group.
    :ivar Q3: Third quartile (75th percentile) value in the group.
    :ivar P01: 1st percentile value in the group.
    :ivar P05: 5th percentile value in the group.
    :ivar P10: 10th percentile value in the group.
    :ivar P90: 90th percentile value in the group.
    :ivar P95: 95th percentile value in the group.
    :ivar P99: 99th percentile value in the group.
    """

    COUNT = "count"
    SUM = "sum"
    MIN = "min"
    MAX = "max"
    AVERAGE = "average"
    STDDEV = "stddev"
    MEDIAN = "median"
    Q1 = "q1"
    Q3 = "q3"
    P01 = "p01"
    P05 = "p05"
    P10 = "p10"
    P90 = "p90"
    P95 = "p95"
    P99 = "p99"

    @property
    def requires_field(self) -> bool:
        """Check if this aggregation function requires a field.

        :returns: True if the function requires a field, False otherwise.
        :rtype: bool
        """
        return self != AggregationFunction.COUNT


class GroupSpecVisitor(Protocol, Generic[T]):
    """Visitor interface for traversing group specifications.

    Implementations should provide a visit method for group specifications.
    The generic type parameter T specifies the return type of visit methods.
    """

    def visit_group_spec(self, spec: GroupSpec) -> T:
        """Visit a group specification.

        :param spec: The group specification to visit.
        :type spec: GroupSpec
        :returns: The result of visiting the specification.
        :rtype: T
        """
        ...


@dataclass(frozen=True)
class Aggregation:
    """Represents a single aggregation operation.

    An aggregation consists of an identifier (name), a function to apply,
    and optionally a field to aggregate (required for all functions except COUNT).

    :ivar id: The identifier for this aggregation (used to reference the result).
    :vartype id: str
    :ivar function: The aggregation function to apply.
    :vartype function: AggregationFunction
    :ivar field: The field name to aggregate (None for COUNT, required for others).
    :vartype field: str | None
    """

    id: str
    function: AggregationFunction
    field: str | None = None

    def __post_init__(self) -> None:
        """Validate field requirement after initialization.

        :raises ValueError: If COUNT has a field or if non-COUNT function lacks a field.
        """
        # Validate field requirement
        if self.function == AggregationFunction.COUNT and self.field is not None:
            msg = "COUNT aggregation must not have a field"
            raise ValueError(msg)
        if self.function != AggregationFunction.COUNT and self.field is None:
            msg = f"{self.function.value} aggregation requires a field"
            raise ValueError(msg)

    def __str__(self) -> str:
        """Return a string representation of the aggregation.

        :returns: String in format "id:function" or "id:function:field".
        :rtype: str
        """
        if self.field is None:
            return f"{self.id}:{self.function.value}"
        return f"{self.id}:{self.function.value}:{self.field}"

    @property
    def has_field(self) -> bool:
        """Check if this aggregation has a field.

        :returns: True if the aggregation has a field, False otherwise.
        :rtype: bool
        """
        return self.field is not None


@dataclass(frozen=True)
class GroupSpec:
    """Represents a grouping specification with aggregations.

    A GroupSpec defines how to group data by one or more fields and what
    aggregations to compute for each group. This is equivalent to SQL's
    GROUP BY clause with aggregate functions.

    :ivar group_by: Tuple of field names to group by.
    :vartype group_by: tuple[str, ...]
    :ivar aggregations: Dictionary mapping aggregation IDs to Aggregation objects.
    :vartype aggregations: dict[str, Aggregation]
    """

    group_by: tuple[str, ...]
    aggregations: dict[str, Aggregation]

    def __init__(
        self,
        group_by: tuple[str, ...] | list[str],
        aggregations: dict[str, Aggregation] | list[Aggregation],
    ) -> None:
        """Initialize a GroupSpec with normalization and validation.

        :param group_by: Field names to group by (tuple or list).
        :type group_by: tuple[str, ...] | list[str]
        :param aggregations: Aggregations to compute (dict or list).
        :type aggregations: dict[str, Aggregation] | list[Aggregation]
        :raises ValueError: If aggregation dict keys don't match aggregation IDs.
        """
        # Convert group_by to tuple if needed
        if isinstance(group_by, list):
            group_by = tuple(group_by)

        # Convert aggregations list to dict if needed
        if isinstance(aggregations, list):
            aggregations = {agg.id: agg for agg in aggregations}

        # Validate that dict keys match aggregation IDs (before setting)
        for key, agg in aggregations.items():
            if key != agg.id:
                msg = f"Aggregation dict key '{key}' does not match aggregation ID '{agg.id}'"
                raise ValueError(msg)

        # Use object.__setattr__ since the class is frozen
        object.__setattr__(self, "group_by", group_by)
        object.__setattr__(self, "aggregations", aggregations)

    def __str__(self) -> str:
        """Return a string representation of the group specification.

        Returns a tuple-like format: ("field1,field2", "agg1:count,agg2:min:field")

        :returns: String representation of the grouping specification.
        :rtype: str
        """
        group_by_str = ",".join(self.group_by)
        agg_strs = [str(agg) for agg in self.aggregations.values()]
        agg_str = ",".join(agg_strs)
        return f'("{group_by_str}", "{agg_str}")'

    def accept(self, visitor: GroupSpecVisitor[T]) -> T:
        """Accept a visitor and dispatch to the appropriate visit method.

        :param visitor: The visitor to accept.
        :type visitor: GroupSpecVisitor[T]
        :returns: The result of the visitor's visit method.
        :rtype: T
        """
        return visitor.visit_group_spec(self)


__all__ = [
    "Aggregation",
    "AggregationFunction",
    "GroupSpec",
    "GroupSpecVisitor",
]
