"""Optimizer module for simplifying and optimizing grouping specifications.

This module provides optimization capabilities for removing redundant
grouping fields and aggregation definitions.
"""

from __future__ import annotations

from dataclasses import dataclass

from therismos.grouping._grouping import GroupSpec


@dataclass(frozen=True)
class OptimizationRecord:
    """Record of an optimization transformation.

    :ivar before: The original grouping spec before optimization.
    :vartype before: GroupSpec
    :ivar after: The resulting grouping spec after optimization.
    :vartype after: GroupSpec
    :ivar reason: A description of the optimization rule that was applied.
    :vartype reason: str
    """

    before: GroupSpec
    after: GroupSpec
    reason: str


def optimize(
    spec: GroupSpec, records: list[OptimizationRecord] | None = None
) -> tuple[GroupSpec, list[OptimizationRecord]]:
    """Optimize a grouping specification.

    Applies the following optimization rules:
    1. Removes duplicate grouping fields (keeps last occurrence)
    2. Removes duplicate aggregation IDs (keeps last occurrence)

    :param spec: The grouping specification to optimize.
    :type spec: GroupSpec
    :param records: Optional list to collect optimization records. If None, a new list is created.
    :type records: list[OptimizationRecord] | None
    :returns: A tuple of (optimized grouping spec, list of optimization records).
    :rtype: tuple[GroupSpec, list[OptimizationRecord]]

    Example:
        >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
        >>> spec = GroupSpec(
        ...     group_by=["category", "region", "category"],  # duplicate
        ...     aggregations=[
        ...         Aggregation("total", AggregationFunction.COUNT),
        ...         Aggregation("min_price", AggregationFunction.MIN, "price"),
        ...         Aggregation("total", AggregationFunction.MAX, "quantity"),  # duplicate ID
        ...     ],
        ... )
        >>> optimized, _ = optimize(spec)
        >>> optimized.group_by
        ('region', 'category')
        >>> list(optimized.aggregations.keys())
        ['min_price', 'total']
        >>> optimized.aggregations['total'].function
        <AggregationFunction.MAX: 'max'>
    """
    if records is None:
        records = []

    # Start with the original spec
    current_spec = spec

    # Apply optimization rules
    current_spec = _remove_duplicate_grouping_fields(current_spec, records)
    current_spec = _remove_duplicate_aggregations(current_spec, records)

    return current_spec, records


def _remove_duplicate_grouping_fields(
    spec: GroupSpec, records: list[OptimizationRecord]
) -> GroupSpec:
    """Remove duplicate grouping fields, keeping only the last occurrence.

    When a field appears multiple times in group_by, only the last (rightmost)
    occurrence is kept.

    :param spec: The grouping specification to optimize.
    :type spec: GroupSpec
    :param records: List to collect optimization records.
    :type records: list[OptimizationRecord]
    :returns: The optimized grouping specification.
    :rtype: GroupSpec
    """
    # Track seen field names and their last occurrence index
    seen_fields: dict[str, int] = {}
    for i, field in enumerate(spec.group_by):
        seen_fields[field] = i

    # Keep only the last occurrence of each field
    result_fields = []
    for i, field in enumerate(spec.group_by):
        # Keep this field if it's the last occurrence
        if seen_fields[field] == i:
            result_fields.append(field)

    # If anything was removed, record it
    if len(result_fields) < len(spec.group_by):
        result = GroupSpec(
            group_by=tuple(result_fields),
            aggregations=spec.aggregations,
        )
        removed_count = len(spec.group_by) - len(result_fields)
        records.append(
            OptimizationRecord(
                spec,
                result,
                f"Removed {removed_count} duplicate grouping field(s) "
                f"(overridden by later occurrences)",
            )
        )
        return result

    return spec


def _remove_duplicate_aggregations(spec: GroupSpec, records: list[OptimizationRecord]) -> GroupSpec:
    """Remove duplicate aggregations, keeping only the last occurrence of each ID.

    When an aggregation ID appears multiple times, only the last (rightmost)
    definition is kept.

    :param spec: The grouping specification to optimize.
    :type spec: GroupSpec
    :param records: List to collect optimization records.
    :type records: list[OptimizationRecord]
    :returns: The optimized grouping specification.
    :rtype: GroupSpec
    """
    # Convert aggregations to list to preserve insertion order and find duplicates
    agg_list = list(spec.aggregations.values())

    # Track seen aggregation IDs and their last occurrence index
    seen_ids: dict[str, int] = {}
    for i, agg in enumerate(agg_list):
        seen_ids[agg.id] = i

    # Keep only the last occurrence of each aggregation ID
    result_aggs = []
    for i, agg in enumerate(agg_list):
        # Keep this aggregation if it's the last occurrence of its ID
        if seen_ids[agg.id] == i:
            result_aggs.append(agg)

    # If anything was removed, record it
    if len(result_aggs) < len(agg_list):
        result = GroupSpec(
            group_by=spec.group_by,
            aggregations=result_aggs,
        )
        removed_count = len(agg_list) - len(result_aggs)
        records.append(
            OptimizationRecord(
                spec,
                result,
                f"Removed {removed_count} duplicate aggregation(s) "
                f"(overridden by later occurrences)",
            )
        )
        return result

    return spec


__all__ = [
    "OptimizationRecord",
    "optimize",
]
