"""Tests for the field pruning and projection utilities."""

import pytest

from therismos import (
    FALSE,
    TRUE,
    F,
    FieldSelection,
    Le,
    PruneMode,
    prune_fields,
)


class TestSimplePruning:
    """Pruning a single field with RESTRICT (default) and RELAX modes."""

    @pytest.mark.parametrize(
        "expr,fields,mode,expected",
        [
            (F("age") > 18, frozenset({"age"}), PruneMode.RESTRICT, FALSE),
            (F("age") > 18, frozenset({"age"}), PruneMode.RELAX, TRUE),
            (F("age") > 18, frozenset({"status"}), PruneMode.RESTRICT, F("age") > 18),
            (F("age") > 18, frozenset({"status"}), PruneMode.RELAX, F("age") > 18),
            (F("age") > 18, frozenset(), PruneMode.RESTRICT, F("age") > 18),
        ],
    )
    def test_atomic_prune(
        self,
        expr: object,
        fields: frozenset[str],
        mode: PruneMode,
        expected: object,
    ) -> None:
        """Pruning a targeted atomic expression yields the correct constant or unchanged node."""
        assert prune_fields(expr, fields, mode=mode) == expected  # type: ignore[arg-type]

    @pytest.mark.parametrize(
        "expr_type,field_name",
        [
            (lambda f: f == 5, "age"),
            (lambda f: f != 5, "age"),
            (lambda f: f < 5, "age"),
            (lambda f: f <= 5, "age"),
            (lambda f: f > 5, "age"),
            (lambda f: f >= 5, "age"),
            (lambda f: f.between(1, 10), "age"),
            (lambda f: f.matches(r"\d+"), "age"),
            (lambda f: f.is_in(1, 2, 3), "age"),
            (lambda f: f.is_null(), "age"),
            (lambda f: f.is_not_null(), "age"),
        ],
    )
    def test_all_field_based_subtypes_are_replaced(
        self, expr_type: object, field_name: str
    ) -> None:
        """All FieldBasedExpr subtypes are replaced by FALSE in RESTRICT mode."""
        expr = expr_type(F(field_name))  # type: ignore[operator]
        result = prune_fields(expr, frozenset({field_name}))
        assert result is FALSE


class TestConstantsPassThrough:
    """TrueExpr and FalseExpr are unaffected by pruning."""

    @pytest.mark.parametrize("constant", [TRUE, FALSE])
    def test_constant_unchanged(self, constant: object) -> None:
        """Constant expressions are returned unchanged regardless of fields."""
        assert prune_fields(constant, frozenset({"any_field"})) is constant  # type: ignore[arg-type]


class TestPolarityUnderNot:
    """Polarity flips correctly when a pruned leaf is inside NOT."""

    @pytest.mark.parametrize(
        "mode,expected",
        [
            (PruneMode.RESTRICT, FALSE),
            (PruneMode.RELAX, TRUE),
        ],
    )
    def test_not_wrapping_single_pruned_field(self, mode: PruneMode, expected: object) -> None:
        """NOT(pruned) collapses to the same constant as a bare pruned leaf."""
        expr = ~(F("age") > 18)
        assert prune_fields(expr, frozenset({"age"}), mode=mode) == expected

    @pytest.mark.parametrize(
        "mode,expected",
        [
            (PruneMode.RESTRICT, FALSE),
            (PruneMode.RELAX, TRUE),
        ],
    )
    def test_double_not_restores_polarity(self, mode: PruneMode, expected: object) -> None:
        """NOT(NOT(pruned)) restores polarity — same result as a bare pruned leaf."""
        expr = ~~(F("age") > 18)
        assert prune_fields(expr, frozenset({"age"}), mode=mode) == expected

    def test_not_wrapping_kept_field(self) -> None:
        """NOT around a non-pruned field is simplified to its complement."""
        age = F("age")
        result = prune_fields(~(age > 18), frozenset({"status"}))
        assert result == Le(age, 18)


class TestAndPropagation:
    """Optimizer collapses AND trees correctly after substitution."""

    def test_restrict_any_pruned_field_in_and_yields_false(self) -> None:
        """AND(pruned_field, other_field) → FALSE in RESTRICT mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset({"age"}))
        assert result is FALSE

    def test_relax_pruned_field_in_and_drops_constraint(self) -> None:
        """AND(pruned_field, other_field) → other_field in RELAX mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset({"age"}), mode=PruneMode.RELAX)
        assert result == (F("status") == "active")

    def test_all_fields_pruned_and_restrict(self) -> None:
        """AND with all fields pruned → FALSE in RESTRICT mode."""
        expr = (F("age") > 18) & (F("score") < 100)
        assert prune_fields(expr, frozenset({"age", "score"})) is FALSE

    def test_all_fields_pruned_and_relax(self) -> None:
        """AND with all fields pruned → TRUE in RELAX mode."""
        expr = (F("age") > 18) & (F("score") < 100)
        assert prune_fields(expr, frozenset({"age", "score"}), mode=PruneMode.RELAX) is TRUE


class TestOrPropagation:
    """Optimizer collapses OR trees correctly after substitution."""

    def test_restrict_pruned_field_in_or_drops_branch(self) -> None:
        """OR(pruned_field, other_field) → other_field in RESTRICT mode."""
        expr = (F("age") > 18) | (F("status") == "active")
        result = prune_fields(expr, frozenset({"age"}))
        assert result == (F("status") == "active")

    def test_relax_pruned_field_in_or_yields_true(self) -> None:
        """OR(pruned_field, other_field) → TRUE in RELAX mode."""
        expr = (F("age") > 18) | (F("status") == "active")
        result = prune_fields(expr, frozenset({"age"}), mode=PruneMode.RELAX)
        assert result is TRUE

    def test_all_fields_pruned_or_restrict(self) -> None:
        """OR with all fields pruned → FALSE in RESTRICT mode."""
        expr = (F("age") > 18) | (F("score") < 100)
        assert prune_fields(expr, frozenset({"age", "score"})) is FALSE

    def test_all_fields_pruned_or_relax(self) -> None:
        """OR with all fields pruned → TRUE in RELAX mode."""
        expr = (F("age") > 18) | (F("score") < 100)
        assert prune_fields(expr, frozenset({"age", "score"}), mode=PruneMode.RELAX) is TRUE


class TestMixedTrees:
    """Complex nested trees with multiple fields, some pruned and some retained."""

    def test_and_not_pruned_restrict(self) -> None:
        """AND(other, NOT(pruned)) → FALSE in RESTRICT mode."""
        expr = (F("status") == "active") & ~(F("age") > 18)
        result = prune_fields(expr, frozenset({"age"}))
        assert result is FALSE

    def test_and_not_pruned_relax(self) -> None:
        """AND(other, NOT(pruned)) → other in RELAX mode."""
        expr = (F("status") == "active") & ~(F("age") > 18)
        result = prune_fields(expr, frozenset({"age"}), mode=PruneMode.RELAX)
        assert result == (F("status") == "active")

    def test_or_not_pruned_restrict(self) -> None:
        """OR(other, NOT(pruned)) → other in RESTRICT mode (FALSE is identity for OR)."""
        expr = (F("status") == "active") | ~(F("age") > 18)
        result = prune_fields(expr, frozenset({"age"}))
        assert result == (F("status") == "active")

    def test_or_not_pruned_relax(self) -> None:
        """OR(other, NOT(pruned)) → TRUE in RELAX mode."""
        expr = (F("status") == "active") | ~(F("age") > 18)
        result = prune_fields(expr, frozenset({"age"}), mode=PruneMode.RELAX)
        assert result is TRUE

    def test_deeply_nested(self) -> None:
        """Deeply nested tree: pruning propagates correctly through all levels.

        ``(a==1 AND b==2) OR (c==3 AND d==4)`` relaxing ``{a, c}`` replaces each
        with TRUE, collapsing both inner ANDs to single terms, yielding ``b==2 OR d==4``.
        """
        expr = ((F("a") == 1) & (F("b") == 2)) | ((F("c") == 3) & (F("d") == 4))
        result = prune_fields(expr, frozenset({"a", "c"}), mode=PruneMode.RELAX)
        assert result == ((F("b") == 2) | (F("d") == 4))

    def test_unpruned_fields_unmodified(self) -> None:
        """Fields not in the pruned set are returned structurally identical."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset({"score"}), mode=PruneMode.RESTRICT)
        assert result == expr


class TestFieldSelectionKeep:
    """FieldSelection.KEEP retains only the specified fields."""

    def test_keep_single_field_restrict(self) -> None:
        """AND(kept, other) → kept field only in RESTRICT/KEEP mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset({"age"}), selection=FieldSelection.KEEP)
        assert result is FALSE

    def test_keep_single_field_relax(self) -> None:
        """AND(kept, other) → kept field only in RELAX/KEEP mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(
            expr, frozenset({"age"}), selection=FieldSelection.KEEP, mode=PruneMode.RELAX
        )
        assert result == (F("age") > 18)

    def test_keep_all_fields_unchanged(self) -> None:
        """When all referenced fields are kept, the expression is unchanged."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset({"age", "status"}), selection=FieldSelection.KEEP)
        assert result == expr

    def test_keep_no_fields_restrict(self) -> None:
        """Keeping an empty set removes everything → FALSE in RESTRICT mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(expr, frozenset(), selection=FieldSelection.KEEP)
        assert result is FALSE

    def test_keep_no_fields_relax(self) -> None:
        """Keeping an empty set removes everything → TRUE in RELAX mode."""
        expr = (F("age") > 18) & (F("status") == "active")
        result = prune_fields(
            expr, frozenset(), selection=FieldSelection.KEEP, mode=PruneMode.RELAX
        )
        assert result is TRUE

    def test_keep_is_dual_of_prune(self) -> None:
        """KEEP on {A} gives the same result as PRUNE on {B} when fields are {A, B}."""
        expr = (F("age") > 18) & (F("status") == "active")
        via_keep = prune_fields(
            expr, frozenset({"age"}), selection=FieldSelection.KEEP, mode=PruneMode.RELAX
        )
        via_prune = prune_fields(expr, frozenset({"status"}), mode=PruneMode.RELAX)
        assert via_keep == via_prune
