"""Optimizer module for simplifying and optimizing expression trees.

This module provides optimization capabilities for detecting contradictions,
tautologies, and simplification opportunities in expression trees.

.. note::
    - The optimizer assumes single-valued semantics for comparison expressions.
    - Optimization rules are conservative and remain correct for multi-valued fields,
      but may miss some optimization opportunities in those cases.
    - Contradiction detection for :class:`Eq`, :class:`Ne`, :class:`Lt`, :class:`Le`,
      :class:`Gt`, :class:`Ge`
      assumes single values.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from therismos.expr import (
    FALSE,
    TRUE,
    AllExpr,
    AnyExpr,
    Eq,
    Expr,
    FalseExpr,
    Field,
    Ge,
    Gt,
    In,
    IsNull,
    Le,
    Lt,
    Ne,
    NotExpr,
    TrueExpr,
)


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

    :ivar before: The expression before optimization.
    :vartype before: Expr
    :ivar after: The resulting expression after optimization.
    :vartype after: Expr
    :ivar reason: A description of the applied optimization rule.
    :vartype reason: str
    """

    before: Expr
    after: Expr
    reason: str


def optimize(
    expr: Expr, records: list[OptimizationRecord] | None = None
) -> tuple[Expr, list[OptimizationRecord]]:
    """Optimize an expression tree by applying simplification rules.

    :param expr: The expression to be optimized.
    :type expr: Expr
    :param records: An optional list to collect optimization records.
                    If ``None``, a new list is created.
    :type records: list[OptimizationRecord] | None
    :return: A tuple containing the optimized expression and a list of optimization records.
    :rtype: tuple[Expr, list[OptimizationRecord]]
    """
    if records is None:
        records = []

    def _record(
        records_list: list[OptimizationRecord],
        before: Expr,
        after: Expr,
        reason: str,
    ) -> None:
        reason = reason.strip()
        records_list.append(OptimizationRecord(before, after, reason))

    def _optimize_expr(e: Expr, _records: list[OptimizationRecord]) -> Expr:
        """Recursively optimize an expression."""
        # First, optimize children
        if isinstance(e, NotExpr):
            optimized_child = _optimize_expr(e.expr, _records)
            if optimized_child is not e.expr:
                new_expr: Expr = NotExpr(optimized_child)
                # Record reconstruction if the child was optimized
                if new_expr != e:
                    _record(_records, e, new_expr, "Optimize child in NOT")
                    e = new_expr
        elif isinstance(e, AllExpr):
            optimized_children = tuple(_optimize_expr(child, _records) for child in e.exprs)
            if optimized_children != e.exprs:
                new_expr = AllExpr(*optimized_children)
                if new_expr != e:
                    _record(_records, e, new_expr, "Optimize children in AND")
                    e = new_expr
        elif isinstance(e, AnyExpr):
            optimized_children = tuple(_optimize_expr(child, _records) for child in e.exprs)
            if optimized_children != e.exprs:
                new_expr = AnyExpr(*optimized_children)
                if new_expr != e:
                    _record(_records, e, new_expr, "Optimize children in OR")
                    e = new_expr

        # Then, apply optimization rules
        while True:
            optimized = _apply_rules(e, _records)
            if optimized is e:
                break
            # Recording is now done within individual optimization functions
            e = optimized

        return e

    def _apply_rules(e: Expr, _records: list[OptimizationRecord]) -> Expr:
        """Apply a single optimization rule if applicable."""
        # IN simplifications
        if isinstance(e, In):
            return _optimize_in(e, _records)

        # NOT simplifications
        if isinstance(e, NotExpr):
            return _optimize_not(e, _records)

        # AllExpr (AND) simplifications
        if isinstance(e, AllExpr):
            return _optimize_all(e, _records)

        # AnyExpr (OR) simplifications
        if isinstance(e, AnyExpr):
            return _optimize_any(e, _records)

        return e

    def _optimize_in(e: In, _records: list[OptimizationRecord]) -> Expr:
        """Optimize IN expressions."""
        # IN([]) → FALSE
        if not e.values:
            _record(_records, e, FALSE, "IN with empty set → FALSE")
            return FALSE

        # IN([v]) → Eq(v)
        if len(e.values) == 1:
            result = Eq(e.field, e.values[0])
            _record(_records, e, result, "IN with single value → Eq")
            return result

        return e

    def _optimize_not(e: NotExpr, _records: list[OptimizationRecord]) -> Expr:
        """Optimize NOT expressions."""
        # NOT(TRUE) → FALSE
        if isinstance(e.expr, TrueExpr):
            _record(_records, e, FALSE, "NOT(TRUE) → FALSE")
            return FALSE

        # NOT(FALSE) → TRUE
        if isinstance(e.expr, FalseExpr):
            _record(_records, e, TRUE, "NOT(FALSE) → TRUE")
            return TRUE

        # NOT(NOT(x)) → x
        if isinstance(e.expr, NotExpr):
            result = e.expr.expr
            _record(_records, e, result, "NOT(NOT(x)) → x")
            return result

        # De Morgan's laws
        # NOT(AllExpr(a, b, ...)) → AnyExpr(NOT(a), NOT(b), ...)
        if isinstance(e.expr, AllExpr):
            result = AnyExpr(*[NotExpr(child) for child in e.expr.exprs])
            _record(_records, e, result, "De Morgan's law: NOT(AND(...)) → OR(NOT(...))")
            return result

        # NOT(AnyExpr(a, b, ...)) → AllExpr(NOT(a), NOT(b), ...)
        if isinstance(e.expr, AnyExpr):
            result = AllExpr(*[NotExpr(child) for child in e.expr.exprs])
            _record(_records, e, result, "De Morgan's law: NOT(OR(...)) → AND(NOT(...))")
            return result

        return e

    def _optimize_all(e: AllExpr, _records: list[OptimizationRecord]) -> Expr:
        """Optimize AllExpr (AND) expressions."""
        # AllExpr([]) → TRUE
        if not e.exprs:
            _record(_records, e, TRUE, "AllExpr([]) → TRUE")
            return TRUE

        # Single operand: AllExpr([x]) → x
        if len(e.exprs) == 1:
            result = e.exprs[0]
            _record(_records, e, result, "AllExpr([x]) → x")
            return result

        # AllExpr([..., FALSE, ...]) → FALSE
        if any(isinstance(child, FalseExpr) for child in e.exprs):
            _record(_records, e, FALSE, "AllExpr([..., FALSE, ...]) → FALSE")
            return FALSE

        # Remove TRUE from operands
        filtered = tuple(child for child in e.exprs if not isinstance(child, TrueExpr))
        if len(filtered) != len(e.exprs):
            if not filtered:
                _record(_records, e, TRUE, "AllExpr([TRUE, ...]) → TRUE")
                return TRUE
            if len(filtered) == 1:
                result = filtered[0]
                _record(_records, e, result, "AllExpr([TRUE, ..., x]) → x")
                return result
            result = AllExpr(*filtered)
            _record(_records, e, result, "Identity elimination: remove TRUE from AND")
            return result

        # Intersect AND equality/IN chains
        intersected = _intersect_and_equalities(e.exprs, _records)
        if intersected is not None:
            if intersected is FALSE:
                _record(_records, e, FALSE, "Contradiction detected in AND")
            else:
                _record(_records, e, intersected, "AND equality/IN chain intersection")
            return intersected

        # Check for other contradictions
        if _has_contradiction(e.exprs, _records):
            _record(_records, e, FALSE, "Contradiction detected in AND")
            return FALSE

        return e

    def _optimize_any(e: AnyExpr, _records: list[OptimizationRecord]) -> Expr:
        """Optimize AnyExpr (OR) expressions."""
        # AnyExpr([]) → FALSE
        if not e.exprs:
            _record(_records, e, FALSE, "AnyExpr([]) → FALSE")
            return FALSE

        # Single operand: AnyExpr([x]) → x
        if len(e.exprs) == 1:
            result = e.exprs[0]
            _record(_records, e, result, "AnyExpr([x]) → x")
            return result

        # AnyExpr([..., TRUE, ...]) → TRUE
        if any(isinstance(child, TrueExpr) for child in e.exprs):
            _record(_records, e, TRUE, "AnyExpr([..., TRUE, ...]) → TRUE")
            return TRUE

        # Remove FALSE from operands
        filtered = tuple(child for child in e.exprs if not isinstance(child, FalseExpr))
        if len(filtered) != len(e.exprs):
            if not filtered:
                _record(_records, e, FALSE, "AnyExpr([FALSE, ...]) → FALSE")
                return FALSE
            if len(filtered) == 1:
                result = filtered[0]
                _record(_records, e, result, "AnyExpr([FALSE, ..., x]) → x")
                return result
            result = AnyExpr(*filtered)
            _record(_records, e, result, "Identity elimination: remove FALSE from OR")
            return result

        # Check for tautologies
        if _has_tautology(e.exprs, _records):
            _record(_records, e, TRUE, "Tautology detected in OR")
            return TRUE

        # Aggregate OR equality chains to IN
        aggregated = _aggregate_or_equalities(e.exprs, _records)
        if aggregated is not None:
            _record(_records, e, aggregated, "OR equality chain aggregation to IN")
            return aggregated

        return e

    def _has_contradiction(exprs: tuple[Expr, ...], _records: list[OptimizationRecord]) -> bool:
        """Check if a list of AND-connected expressions contains contradictions."""
        by_field: dict[str, list[Expr]] = {}
        for expr in exprs:
            if hasattr(expr, "field") and isinstance(expr.field, Field):
                field_name = expr.field.name
                if field_name not in by_field:
                    by_field[field_name] = []
                by_field[field_name].append(expr)

        return any(
            _field_has_contradiction(field_exprs, _records) for field_exprs in by_field.values()
        )

    def _field_has_contradiction(exprs: list[Expr], _records: list[OptimizationRecord]) -> bool:
        """Check if expressions on a single field contain contradictions."""
        eq_values, ne_values = set(), set()
        is_null_count, is_not_null_count = 0, 0
        lt_values, le_values, gt_values, ge_values = [], [], [], []

        for expr in exprs:
            if isinstance(expr, Eq):
                eq_values.add(expr.value)
            elif isinstance(expr, Ne):
                ne_values.add(expr.value)
            elif isinstance(expr, IsNull):
                if expr.is_null:
                    is_null_count += 1
                else:
                    is_not_null_count += 1
            elif isinstance(expr, Lt):
                lt_values.append(expr.value)
            elif isinstance(expr, Le):
                le_values.append(expr.value)
            elif isinstance(expr, Gt):
                gt_values.append(expr.value)
            elif isinstance(expr, Ge):
                ge_values.append(expr.value)

        if eq_values & ne_values:
            return True

        if is_null_count > 0 and is_not_null_count > 0:
            return True

        # Check range contradictions
        for lt_val in lt_values:
            for gt_val in gt_values:
                try:
                    if gt_val >= lt_val:
                        return True
                except TypeError:
                    pass  # Cannot compare, skip

        for le_val in le_values:
            for gt_val in gt_values:
                try:
                    if gt_val >= le_val:
                        return True
                except TypeError:
                    pass

        for ge_val in ge_values:
            for lt_val in lt_values:
                try:
                    if ge_val >= lt_val:
                        return True
                except TypeError:
                    pass

        return False

    def _has_tautology(exprs: tuple[Expr, ...], _records: list[OptimizationRecord]) -> bool:
        """Check if a list of OR-connected expressions contains tautologies."""
        by_field: dict[str, list[Expr]] = {}
        for expr in exprs:
            if hasattr(expr, "field") and isinstance(expr.field, Field):
                field_name = expr.field.name
                if field_name not in by_field:
                    by_field[field_name] = []
                by_field[field_name].append(expr)

        return any(_field_has_tautology(field_exprs, _records) for field_exprs in by_field.values())

    def _field_has_tautology(exprs: list[Expr], _records: list[OptimizationRecord]) -> bool:
        """Check if expressions on a single field contain tautologies."""
        eq_values, ne_values = set(), set()
        is_null_count, is_not_null_count = 0, 0
        has_lt, has_ge, has_le, has_gt = False, False, False, False
        lt_ge_pairs, le_gt_pairs = [], []

        for expr in exprs:
            if isinstance(expr, Eq):
                eq_values.add(expr.value)
            elif isinstance(expr, Ne):
                ne_values.add(expr.value)
            elif isinstance(expr, IsNull):
                if expr.is_null:
                    is_null_count += 1
                else:
                    is_not_null_count += 1
            elif isinstance(expr, Lt):
                has_lt = True
                lt_ge_pairs.append(("lt", expr.value))
            elif isinstance(expr, Le):
                has_le = True
                le_gt_pairs.append(("le", expr.value))
            elif isinstance(expr, Gt):
                has_gt = True
                le_gt_pairs.append(("gt", expr.value))
            elif isinstance(expr, Ge):
                has_ge = True
                lt_ge_pairs.append(("ge", expr.value))

        if eq_values & ne_values:
            return True

        if is_null_count > 0 and is_not_null_count > 0:
            return True

        if has_lt and has_ge:
            lt_vals = {val for op, val in lt_ge_pairs if op == "lt"}
            ge_vals = {val for op, val in lt_ge_pairs if op == "ge"}
            if lt_vals & ge_vals:
                return True

        if has_le and has_gt:
            le_vals = {val for op, val in le_gt_pairs if op == "le"}
            gt_vals = {val for op, val in le_gt_pairs if op == "gt"}
            if le_vals & gt_vals:
                return True

        return False

    def _aggregate_or_equalities(
        exprs: tuple[Expr, ...], _records: list[OptimizationRecord]
    ) -> Expr | None:
        """Aggregate OR equality and IN chains to IN expressions.

        Patterns:
        - (f == v1) OR (f == v2) OR ... OR (f == vk) → f IN (v1, v2, ..., vk)
        - (f == v1) OR (f IN (v2, v3)) → f IN (v1, v2, v3)
        - (f IN (v1, v2)) OR (f IN (v3, v4)) → f IN (v1, v2, v3, v4)

        Only applies if ALL expressions on a field are Eq or In.
        """
        by_field: dict[str, list[Expr]] = {}
        non_field_exprs: list[Expr] = []

        for expr in exprs:
            if hasattr(expr, "field") and isinstance(expr.field, Field):
                field_name = expr.field.name
                if field_name not in by_field:
                    by_field[field_name] = []
                by_field[field_name].append(expr)
            else:
                non_field_exprs.append(expr)

        aggregated_exprs: list[Expr] = list(non_field_exprs)
        made_change = False

        for _, field_exprs in by_field.items():
            if all(isinstance(e, (Eq, In)) for e in field_exprs) and len(field_exprs) > 1:
                all_values: list[Any] = []
                field = None
                for e in field_exprs:
                    if isinstance(e, Eq):
                        if field is None:
                            field = e.field
                        all_values.append(e.value)
                    elif isinstance(e, In):
                        if field is None:
                            field = e.field
                        all_values.extend(e.values)

                if field is not None:
                    aggregated_exprs.append(In(field, tuple(all_values)))
                    made_change = True
            else:
                aggregated_exprs.extend(field_exprs)

        if not made_change:
            return None

        return aggregated_exprs[0] if len(aggregated_exprs) == 1 else AnyExpr(*aggregated_exprs)

    def _intersect_and_equalities(
        exprs: tuple[Expr, ...], _records: list[OptimizationRecord]
    ) -> Expr | None:
        """Intersect AND equality and IN chains.

        Patterns:
        - (f == v1) AND (f == v2) → FALSE if v1 != v2, or f == v1 if v1 == v2
        - (f == v) AND (f IN (v, ...)) → f == v
        - (f == v) AND (f IN (...)) → FALSE if v not in set
        - (f IN (v1, v2)) AND (f IN (v3, v4)) → FALSE if no intersection
        - (f IN (v1, v2)) AND (f IN (v2, v3)) → f == v2 if intersection is single
        - (f IN (v1, v2, v3)) AND (f IN (v2, v3, v4)) → f IN (v2, v3)

        Only applies if ALL expressions on a field are Eq or In.
        """
        by_field: dict[str, list[Expr]] = {}
        non_field_exprs: list[Expr] = []

        for expr in exprs:
            if hasattr(expr, "field") and isinstance(expr.field, Field):
                field_name = expr.field.name
                if field_name not in by_field:
                    by_field[field_name] = []
                by_field[field_name].append(expr)
            else:
                non_field_exprs.append(expr)

        intersected_exprs: list[Expr] = list(non_field_exprs)
        made_change = False

        for _, field_exprs in by_field.items():
            if all(isinstance(e, (Eq, In)) for e in field_exprs) and len(field_exprs) > 1:
                field = None
                intersection_set: set[Any] | None = None

                for e in field_exprs:
                    if isinstance(e, Eq):
                        if field is None:
                            field = e.field
                        if intersection_set is None:
                            intersection_set = {e.value}
                        else:
                            intersection_set &= {e.value}
                    elif isinstance(e, In):
                        if field is None:
                            field = e.field
                        if intersection_set is None:
                            intersection_set = set(e.values)
                        else:
                            intersection_set &= set(e.values)

                if field is not None and intersection_set is not None:
                    if not intersection_set:
                        return FALSE
                    if len(intersection_set) == 1:
                        intersected_exprs.append(Eq(field, next(iter(intersection_set))))
                    else:
                        intersected_exprs.append(In(field, tuple(intersection_set)))
                    made_change = True
            else:
                intersected_exprs.extend(field_exprs)

        if not made_change:
            return None

        return intersected_exprs[0] if len(intersected_exprs) == 1 else AllExpr(*intersected_exprs)

    result = _optimize_expr(expr, records)
    return result, records


__all__ = ["OptimizationRecord", "optimize"]
