"""Field pruning and projection utilities for expression trees.

This module provides :func:`prune_fields`, which removes or projects
field-based constraints from an expression tree in a polarity-aware manner,
then delegates simplification to the existing optimizer.
"""

import enum

from therismos.expr._expr import (
    FALSE,
    TRUE,
    AllExpr,
    AnyExpr,
    Expr,
    FieldBasedExpr,
    NotExpr,
)
from therismos.expr.optimizer import optimize


class FieldSelection(enum.Enum):
    """Controls which fields the ``fields`` argument of :func:`prune_fields` targets.

    :cvar PRUNE: The supplied field names are **removed**; all others are kept.
    :cvar KEEP: The supplied field names are **kept**; all others are removed.
    """

    PRUNE = "prune"
    KEEP = "keep"


class PruneMode(enum.Enum):
    """Controls the replacement semantics when a targeted field expression is removed.

    The substitution is polarity-aware: a pruned leaf inside an odd number of
    ``NOT`` wrappers has its replacement flipped so that conservative or
    permissive semantics are preserved consistently throughout the tree.

    :cvar RESTRICT: Conservative mode.  A record whose matching constraints
        cannot be resolved is **excluded**.  Targeted leaves are replaced with
        ``FALSE`` at positive polarity and ``TRUE`` at negative polarity
        (so ``NOT(pruned)`` collapses to ``FALSE`` after optimization).
    :cvar RELAX: Permissive mode.  Targeted constraints are simply dropped and
        records **pass through**.  Targeted leaves are replaced with ``TRUE``
        at positive polarity and ``FALSE`` at negative polarity
        (so ``NOT(pruned)`` collapses to ``TRUE`` after optimization).
    """

    RESTRICT = "restrict"
    RELAX = "relax"


def prune_fields(
    expr: Expr,
    fields: frozenset[str],
    *,
    selection: FieldSelection = FieldSelection.PRUNE,
    mode: PruneMode = PruneMode.RESTRICT,
) -> Expr:
    """Remove or project field-based constraints from an expression tree.

    Walks *expr* and replaces every :class:`~therismos.expr.FieldBasedExpr`
    whose field matches the selection criteria with a boolean constant,
    then runs :func:`~therismos.expr.optimizer.optimize` to simplify the
    resulting tree.

    The replacement is **polarity-aware**: descending into a
    :class:`~therismos.expr.NotExpr` flips the current polarity so that
    conservative or permissive semantics are maintained correctly under
    negation (see :class:`PruneMode` for the substitution table).

    :param expr: The expression tree to process.
    :type expr: Expr
    :param fields: Field names that drive the selection.
    :type fields: frozenset[str]
    :param selection: Whether *fields* identifies expressions to remove
        (:attr:`FieldSelection.PRUNE`) or to keep (:attr:`FieldSelection.KEEP`).
    :type selection: FieldSelection
    :param mode: Replacement semantics — conservative (:attr:`PruneMode.RESTRICT`)
        or permissive (:attr:`PruneMode.RELAX`).
    :type mode: PruneMode
    :return: The simplified expression after pruning and optimization.
    :rtype: Expr

    :Example:
        >>> from therismos import F, prune_fields, FieldSelection, PruneMode
        >>> expr = (F("age") > 18) & (F("status") == "active")
        >>> prune_fields(expr, frozenset({"age"}))
        FalseExpr()
        >>> prune_fields(expr, frozenset({"age"}), mode=PruneMode.RELAX)
        Eq(field=Field(name='status', type_=None), value='active')
    """
    if selection is FieldSelection.PRUNE:
        should_replace = fields.__contains__
    else:
        should_replace = lambda name: name not in fields  # noqa: E731

    def _prune(e: Expr, positive_polarity: bool) -> Expr:
        if isinstance(e, FieldBasedExpr) and should_replace(e.field.name):
            if mode is PruneMode.RESTRICT:
                return FALSE if positive_polarity else TRUE
            return TRUE if positive_polarity else FALSE

        if isinstance(e, NotExpr):
            return NotExpr(_prune(e.expr, not positive_polarity))

        if isinstance(e, AllExpr):
            return AllExpr(*(_prune(c, positive_polarity) for c in e.exprs))

        if isinstance(e, AnyExpr):
            return AnyExpr(*(_prune(c, positive_polarity) for c in e.exprs))

        return e

    substituted = _prune(expr, positive_polarity=True)
    result, _ = optimize(substituted)
    return result


__all__ = ["FieldSelection", "PruneMode", "prune_fields"]
