"""Tests for rbsp.query.boolean_query — Boolean query tree with cost-based optimization."""

from __future__ import annotations

import pytest

from rbsp.query.boolean_query import (
    MUST,
    MUST_NOT,
    SHOULD,
    BooleanQuery,
    BoostQuery,
    Clause,
    ConstantScoreQuery,
    DisjunctionMaxQuery,
    QueryNode,
    TermQuery,
    optimize_query,
)


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

DOC_SETS: dict[str, set[str]] = {
    "python": {"d1", "d2", "d3"},
    "java": {"d2", "d4"},
    "rust": {"d3", "d5"},
    "go": {"d4", "d5", "d6"},
    "test": {"d1", "d2", "d3", "d4", "d5", "d6"},
    "alpha": {"d1", "d2"},
    "beta": {"d2", "d3"},
    "gamma": {"d3", "d4"},
    "empty": set(),
}


# ===================================================================
# TermQuery
# ===================================================================


class TestTermQuery:
    def test_evaluate_existing_term(self) -> None:
        q = TermQuery("python")
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_evaluate_missing_term(self) -> None:
        q = TermQuery("haskell")
        assert q.evaluate(DOC_SETS) == set()

    def test_evaluate_empty_term(self) -> None:
        q = TermQuery("empty")
        assert q.evaluate(DOC_SETS) == set()

    def test_score_matching_doc(self) -> None:
        q = TermQuery("python")
        assert q.score("d1", DOC_SETS) == 1.0

    def test_score_non_matching_doc(self) -> None:
        q = TermQuery("python")
        assert q.score("d4", DOC_SETS) == 0.0

    def test_estimate_cost(self) -> None:
        q = TermQuery("python")
        assert q.estimate_cost(DOC_SETS) == 3

    def test_estimate_cost_missing(self) -> None:
        q = TermQuery("haskell")
        assert q.estimate_cost(DOC_SETS) == 0

    def test_equality(self) -> None:
        assert TermQuery("python") == TermQuery("python")
        assert TermQuery("python") != TermQuery("java")
        assert TermQuery("python") != "python"

    def test_hash(self) -> None:
        assert hash(TermQuery("python")) == hash(TermQuery("python"))

    def test_repr(self) -> None:
        assert repr(TermQuery("python")) == "TermQuery('python')"


# ===================================================================
# Clause
# ===================================================================


class TestClause:
    def test_valid_occur_types(self) -> None:
        for occur in (MUST, SHOULD, MUST_NOT):
            c = Clause(query=TermQuery("x"), occur=occur)
            assert c.occur == occur

    def test_invalid_occur_type(self) -> None:
        with pytest.raises(ValueError, match="Invalid occur type"):
            Clause(query=TermQuery("x"), occur="MAYBE")

    def test_frozen(self) -> None:
        c = Clause(query=TermQuery("x"), occur=MUST)
        with pytest.raises(AttributeError):
            c.occur = SHOULD  # type: ignore[misc]


# ===================================================================
# BooleanQuery — MUST (AND)
# ===================================================================


class TestBooleanQueryMust:
    def test_single_must(self) -> None:
        q = BooleanQuery([Clause(TermQuery("python"), MUST)])
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_two_must_intersection(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), MUST),
        ])
        assert q.evaluate(DOC_SETS) == {"d2"}

    def test_must_empty_intersection(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("go"), MUST),
        ])
        assert q.evaluate(DOC_SETS) == set()

    def test_must_with_missing_term(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("haskell"), MUST),
        ])
        assert q.evaluate(DOC_SETS) == set()


# ===================================================================
# BooleanQuery — SHOULD (OR)
# ===================================================================


class TestBooleanQueryShould:
    def test_single_should(self) -> None:
        q = BooleanQuery([Clause(TermQuery("python"), SHOULD)])
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_two_should_union(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), SHOULD),
            Clause(TermQuery("java"), SHOULD),
        ])
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3", "d4"}

    def test_minimum_should_match_2(self) -> None:
        q = BooleanQuery(
            [
                Clause(TermQuery("alpha"), SHOULD),
                Clause(TermQuery("beta"), SHOULD),
                Clause(TermQuery("gamma"), SHOULD),
            ],
            minimum_should_match=2,
        )
        # alpha={d1,d2}, beta={d2,d3}, gamma={d3,d4}
        # d2 matches alpha+beta (2), d3 matches beta+gamma (2)
        assert q.evaluate(DOC_SETS) == {"d2", "d3"}

    def test_minimum_should_match_3(self) -> None:
        q = BooleanQuery(
            [
                Clause(TermQuery("alpha"), SHOULD),
                Clause(TermQuery("beta"), SHOULD),
                Clause(TermQuery("gamma"), SHOULD),
            ],
            minimum_should_match=3,
        )
        # No doc matches all three
        assert q.evaluate(DOC_SETS) == set()

    def test_minimum_should_match_exceeds_clauses(self) -> None:
        q = BooleanQuery(
            [Clause(TermQuery("python"), SHOULD)],
            minimum_should_match=5,
        )
        assert q.evaluate(DOC_SETS) == set()


# ===================================================================
# BooleanQuery — MUST_NOT (NOT)
# ===================================================================


class TestBooleanQueryMustNot:
    def test_must_not_excludes(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), MUST_NOT),
        ])
        # python={d1,d2,d3} minus java={d2,d4} = {d1,d3}
        assert q.evaluate(DOC_SETS) == {"d1", "d3"}

    def test_must_not_alone_returns_empty(self) -> None:
        """MUST_NOT alone produces nothing (no positive clause)."""
        q = BooleanQuery([Clause(TermQuery("python"), MUST_NOT)])
        assert q.evaluate(DOC_SETS) == set()

    def test_should_with_must_not(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), SHOULD),
            Clause(TermQuery("java"), SHOULD),
            Clause(TermQuery("rust"), MUST_NOT),
        ])
        # union(python,java)={d1,d2,d3,d4} minus rust={d3,d5} = {d1,d2,d4}
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d4"}


# ===================================================================
# BooleanQuery — Mixed MUST + SHOULD
# ===================================================================


class TestBooleanQueryMixed:
    def test_must_and_should_intersection(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("test"), MUST),
            Clause(TermQuery("python"), SHOULD),
        ])
        # MUST: test={d1..d6}, SHOULD: python={d1,d2,d3}
        # With default msm=1, intersect MUST with SHOULD
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_must_and_should_msm_zero(self) -> None:
        """When msm=0, SHOULD is optional (boost only)."""
        q = BooleanQuery(
            [
                Clause(TermQuery("test"), MUST),
                Clause(TermQuery("python"), SHOULD),
            ],
            minimum_should_match=0,
        )
        # MUST drives the result, SHOULD is boost-only
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3", "d4", "d5", "d6"}

    def test_must_should_must_not_combined(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("test"), MUST),
            Clause(TermQuery("python"), SHOULD),
            Clause(TermQuery("java"), MUST_NOT),
        ])
        # MUST & SHOULD = {d1,d2,d3}, minus java={d2,d4} = {d1,d3}
        assert q.evaluate(DOC_SETS) == {"d1", "d3"}


# ===================================================================
# BooleanQuery — Empty / edge cases
# ===================================================================


class TestBooleanQueryEdgeCases:
    def test_empty_clauses(self) -> None:
        q = BooleanQuery([])
        assert q.evaluate(DOC_SETS) == set()

    def test_score_matching_doc(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), MUST),
        ])
        # d2 matches both MUST clauses
        score = q.score("d2", DOC_SETS)
        assert score > 0.0

    def test_score_non_matching_doc(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), MUST),
        ])
        assert q.score("d5", DOC_SETS) == 0.0

    def test_estimate_cost_must(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),  # cost=3
            Clause(TermQuery("java"), MUST),    # cost=2
        ])
        # min of MUST costs
        assert q.estimate_cost(DOC_SETS) == 2

    def test_estimate_cost_should(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("python"), SHOULD),  # cost=3
            Clause(TermQuery("java"), SHOULD),     # cost=2
        ])
        # sum of SHOULD costs
        assert q.estimate_cost(DOC_SETS) == 5


# ===================================================================
# BooleanQuery — Nested
# ===================================================================


class TestBooleanQueryNested:
    def test_nested_must(self) -> None:
        """AND(python, AND(java, rust)) should work correctly."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(TermQuery("test"), MUST),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(inner, MUST),
        ])
        # python={d1,d2,d3} & java={d2,d4} & test={d1..d6} = {d2}
        assert outer.evaluate(DOC_SETS) == {"d2"}

    def test_nested_should(self) -> None:
        """OR(python, OR(java, go)) should union all."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), SHOULD),
            Clause(TermQuery("go"), SHOULD),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), SHOULD),
            Clause(inner, SHOULD),
        ])
        # python={d1,d2,d3} | java={d2,d4} | go={d4,d5,d6}
        assert outer.evaluate(DOC_SETS) == {"d1", "d2", "d3", "d4", "d5", "d6"}


# ===================================================================
# DisjunctionMaxQuery
# ===================================================================


class TestDisjunctionMaxQuery:
    def test_basic_union(self) -> None:
        q = DisjunctionMaxQuery([TermQuery("python"), TermQuery("java")])
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3", "d4"}

    def test_empty_queries(self) -> None:
        q = DisjunctionMaxQuery([])
        assert q.evaluate(DOC_SETS) == set()

    def test_score_max(self) -> None:
        q = DisjunctionMaxQuery([TermQuery("python"), TermQuery("java")])
        # d2 matches both; max(1.0, 1.0) = 1.0
        assert q.score("d2", DOC_SETS) == 1.0

    def test_score_with_tie_breaker(self) -> None:
        q = DisjunctionMaxQuery(
            [TermQuery("python"), TermQuery("java")],
            tie_breaker=0.1,
        )
        # d2 matches both: max=1.0, tie=0.1*1.0 = 1.1
        assert q.score("d2", DOC_SETS) == pytest.approx(1.1)

    def test_score_single_match(self) -> None:
        q = DisjunctionMaxQuery(
            [TermQuery("python"), TermQuery("java")],
            tie_breaker=0.1,
        )
        # d1 matches only python: 1.0 + 0.0 = 1.0
        assert q.score("d1", DOC_SETS) == 1.0

    def test_score_no_match(self) -> None:
        q = DisjunctionMaxQuery([TermQuery("python")])
        assert q.score("d6", DOC_SETS) == 0.0

    def test_estimate_cost(self) -> None:
        q = DisjunctionMaxQuery([TermQuery("python"), TermQuery("java")])
        # sum of costs: 3 + 2 = 5
        assert q.estimate_cost(DOC_SETS) == 5

    def test_repr(self) -> None:
        q = DisjunctionMaxQuery([TermQuery("a")], tie_breaker=0.5)
        assert "DisjunctionMaxQuery" in repr(q)

    def test_equality(self) -> None:
        q1 = DisjunctionMaxQuery([TermQuery("a")], tie_breaker=0.1)
        q2 = DisjunctionMaxQuery([TermQuery("a")], tie_breaker=0.1)
        assert q1 == q2


# ===================================================================
# BoostQuery
# ===================================================================


class TestBoostQuery:
    def test_evaluate_delegates(self) -> None:
        q = BoostQuery(TermQuery("python"), boost=2.5)
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_score_boosted(self) -> None:
        q = BoostQuery(TermQuery("python"), boost=2.5)
        assert q.score("d1", DOC_SETS) == pytest.approx(2.5)

    def test_score_zero_for_non_match(self) -> None:
        q = BoostQuery(TermQuery("python"), boost=2.5)
        assert q.score("d6", DOC_SETS) == 0.0

    def test_estimate_cost(self) -> None:
        q = BoostQuery(TermQuery("python"), boost=2.5)
        assert q.estimate_cost(DOC_SETS) == 3

    def test_equality(self) -> None:
        q1 = BoostQuery(TermQuery("python"), boost=2.0)
        q2 = BoostQuery(TermQuery("python"), boost=2.0)
        assert q1 == q2
        assert q1 != BoostQuery(TermQuery("python"), boost=3.0)

    def test_repr(self) -> None:
        q = BoostQuery(TermQuery("python"), boost=1.5)
        assert "BoostQuery" in repr(q)
        assert "1.5" in repr(q)


# ===================================================================
# ConstantScoreQuery
# ===================================================================


class TestConstantScoreQuery:
    def test_evaluate_delegates(self) -> None:
        q = ConstantScoreQuery(TermQuery("python"), constant_score=5.0)
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3"}

    def test_score_constant(self) -> None:
        q = ConstantScoreQuery(TermQuery("python"), constant_score=5.0)
        assert q.score("d1", DOC_SETS) == 5.0
        assert q.score("d2", DOC_SETS) == 5.0

    def test_score_zero_for_non_match(self) -> None:
        q = ConstantScoreQuery(TermQuery("python"), constant_score=5.0)
        assert q.score("d6", DOC_SETS) == 0.0

    def test_default_score(self) -> None:
        q = ConstantScoreQuery(TermQuery("python"))
        assert q.score("d1", DOC_SETS) == 1.0

    def test_estimate_cost(self) -> None:
        q = ConstantScoreQuery(TermQuery("python"))
        assert q.estimate_cost(DOC_SETS) == 3

    def test_equality(self) -> None:
        q1 = ConstantScoreQuery(TermQuery("a"), constant_score=2.0)
        q2 = ConstantScoreQuery(TermQuery("a"), constant_score=2.0)
        assert q1 == q2
        assert q1 != ConstantScoreQuery(TermQuery("a"), constant_score=3.0)

    def test_repr(self) -> None:
        q = ConstantScoreQuery(TermQuery("x"), constant_score=7.0)
        assert "ConstantScoreQuery" in repr(q)


# ===================================================================
# QueryNode base class
# ===================================================================


class TestQueryNodeBase:
    def test_evaluate_raises(self) -> None:
        with pytest.raises(NotImplementedError):
            QueryNode().evaluate({})

    def test_score_delegates_to_evaluate(self) -> None:
        """Default score calls evaluate; since evaluate raises, score also raises."""
        with pytest.raises(NotImplementedError):
            QueryNode().score("d1", {})


# ===================================================================
# optimize_query
# ===================================================================


class TestOptimizeQuery:
    def test_term_query_unchanged(self) -> None:
        q = TermQuery("python")
        assert optimize_query(q) is q

    def test_single_must_unwrapped(self) -> None:
        """BooleanQuery with a single MUST clause should simplify."""
        q = BooleanQuery([Clause(TermQuery("python"), MUST)])
        result = optimize_query(q)
        assert isinstance(result, TermQuery)
        assert result.term == "python"

    def test_flatten_nested_must(self) -> None:
        """AND(a, AND(b, c)) -> AND(a, b, c)."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(TermQuery("rust"), MUST),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(inner, MUST),
        ])
        result = optimize_query(outer)
        assert isinstance(result, BooleanQuery)
        must_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == MUST and isinstance(c.query, TermQuery)
        ]
        assert set(must_terms) == {"python", "java", "rust"}

    def test_flatten_nested_should(self) -> None:
        """OR(a, OR(b, c)) -> OR(a, b, c)."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), SHOULD),
            Clause(TermQuery("rust"), SHOULD),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), SHOULD),
            Clause(inner, SHOULD),
        ])
        result = optimize_query(outer)
        assert isinstance(result, BooleanQuery)
        should_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == SHOULD and isinstance(c.query, TermQuery)
        ]
        assert set(should_terms) == {"python", "java", "rust"}

    def test_no_flatten_mixed_child(self) -> None:
        """Don't flatten a child that has mixed occur types."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(TermQuery("rust"), SHOULD),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(inner, MUST),
        ])
        result = optimize_query(outer)
        assert isinstance(result, BooleanQuery)
        # inner should not be flattened (it has mixed MUST + SHOULD)
        inner_clauses = [
            c for c in result.clauses if isinstance(c.query, BooleanQuery)
        ]
        assert len(inner_clauses) == 1

    def test_deduplicate_must(self) -> None:
        """Duplicate MUST clauses should be removed."""
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), MUST),
        ])
        result = optimize_query(q)
        assert isinstance(result, BooleanQuery)
        must_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == MUST
        ]
        assert must_terms == ["python", "java"]

    def test_deduplicate_should(self) -> None:
        q = BooleanQuery([
            Clause(TermQuery("alpha"), SHOULD),
            Clause(TermQuery("alpha"), SHOULD),
            Clause(TermQuery("beta"), SHOULD),
        ])
        result = optimize_query(q)
        assert isinstance(result, BooleanQuery)
        should_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == SHOULD
        ]
        assert should_terms == ["alpha", "beta"]

    def test_must_clauses_sorted_cheapest_first(self) -> None:
        """TermQuery MUST clauses should appear before complex MUST clauses."""
        inner = BooleanQuery([
            Clause(TermQuery("x"), SHOULD),
            Clause(TermQuery("y"), SHOULD),
        ])
        q = BooleanQuery([
            Clause(inner, MUST),
            Clause(TermQuery("python"), MUST),
        ])
        result = optimize_query(q)
        assert isinstance(result, BooleanQuery)
        must_clauses = [c for c in result.clauses if c.occur == MUST]
        # TermQuery should come first
        assert isinstance(must_clauses[0].query, TermQuery)

    def test_optimize_boost_query(self) -> None:
        inner = BooleanQuery([Clause(TermQuery("python"), MUST)])
        q = BoostQuery(inner, boost=2.0)
        result = optimize_query(q)
        assert isinstance(result, BoostQuery)
        # Inner single-MUST should be unwrapped
        assert isinstance(result.query, TermQuery)

    def test_optimize_constant_score_query(self) -> None:
        inner = BooleanQuery([Clause(TermQuery("python"), MUST)])
        q = ConstantScoreQuery(inner, constant_score=3.0)
        result = optimize_query(q)
        assert isinstance(result, ConstantScoreQuery)
        assert isinstance(result.query, TermQuery)

    def test_optimize_disjunction_max(self) -> None:
        inner = BooleanQuery([Clause(TermQuery("java"), MUST)])
        q = DisjunctionMaxQuery([TermQuery("python"), inner])
        result = optimize_query(q)
        assert isinstance(result, DisjunctionMaxQuery)
        # inner single-MUST should be unwrapped to TermQuery
        assert isinstance(result.queries[1], TermQuery)

    def test_optimize_preserves_semantics(self) -> None:
        """Optimized query must produce the same result as the original."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(TermQuery("test"), MUST),
        ])
        original = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(inner, MUST),
            Clause(TermQuery("go"), MUST_NOT),
        ])
        optimized = optimize_query(original)
        assert original.evaluate(DOC_SETS) == optimized.evaluate(DOC_SETS)

    def test_flatten_must_with_must_not_child(self) -> None:
        """AND(a, AND(b, NOT c)) should flatten: AND(a, b, NOT c)."""
        inner = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(TermQuery("go"), MUST_NOT),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(inner, MUST),
        ])
        original_result = outer.evaluate(DOC_SETS)
        result = optimize_query(outer)
        assert isinstance(result, BooleanQuery)
        assert result.evaluate(DOC_SETS) == original_result

        must_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == MUST and isinstance(c.query, TermQuery)
        ]
        assert "python" in must_terms
        assert "java" in must_terms
        must_not_terms = [
            c.query.term  # type: ignore[union-attr]
            for c in result.clauses
            if c.occur == MUST_NOT and isinstance(c.query, TermQuery)
        ]
        assert "go" in must_not_terms


# ===================================================================
# Integration: complex queries
# ===================================================================


class TestComplexQueries:
    def test_elasticsearch_style_query(self) -> None:
        """Simulate a typical ES boolean query with boost and filter."""
        q = BooleanQuery([
            Clause(BoostQuery(TermQuery("python"), boost=2.0), MUST),
            Clause(ConstantScoreQuery(TermQuery("test"), constant_score=0.5), MUST),
            Clause(TermQuery("java"), MUST_NOT),
        ])
        result = q.evaluate(DOC_SETS)
        # python={d1,d2,d3} & test={d1..d6} - java={d2,d4} = {d1,d3}
        assert result == {"d1", "d3"}

    def test_dis_max_in_boolean(self) -> None:
        """dis_max as a clause inside a boolean query."""
        dis_max = DisjunctionMaxQuery(
            [TermQuery("python"), TermQuery("java")],
            tie_breaker=0.1,
        )
        q = BooleanQuery([
            Clause(dis_max, MUST),
            Clause(TermQuery("test"), MUST),
        ])
        # dis_max matches {d1,d2,d3,d4}, test matches {d1..d6}
        # intersection = {d1,d2,d3,d4}
        assert q.evaluate(DOC_SETS) == {"d1", "d2", "d3", "d4"}

    def test_deeply_nested_optimization(self) -> None:
        """Triple nesting should flatten completely."""
        innermost = BooleanQuery([Clause(TermQuery("rust"), MUST)])
        middle = BooleanQuery([
            Clause(TermQuery("java"), MUST),
            Clause(innermost, MUST),
        ])
        outer = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(middle, MUST),
        ])
        original_result = outer.evaluate(DOC_SETS)
        optimized = optimize_query(outer)
        assert optimized.evaluate(DOC_SETS) == original_result

        # Should be fully flattened: AND(python, java, rust)
        assert isinstance(optimized, BooleanQuery)
        must_terms = sorted(
            c.query.term  # type: ignore[union-attr]
            for c in optimized.clauses
            if c.occur == MUST and isinstance(c.query, TermQuery)
        )
        assert must_terms == ["java", "python", "rust"]

    def test_scoring_boosted_boolean(self) -> None:
        """Verify scoring flows through boost and constant score wrappers."""
        q = BooleanQuery([
            Clause(BoostQuery(TermQuery("python"), boost=3.0), SHOULD),
            Clause(ConstantScoreQuery(TermQuery("java"), constant_score=0.5), SHOULD),
        ])
        # d2 matches both python and java
        score_d2 = q.score("d2", DOC_SETS)
        # d1 matches only python
        score_d1 = q.score("d1", DOC_SETS)
        assert score_d2 > 0.0
        assert score_d1 > 0.0
        # Both docs should have positive scores from matching SHOULD clauses
        # d1 matches boosted python (high score), d2 matches both
        assert score_d2 > 0.0 and score_d1 > 0.0

    def test_all_doc_sets_empty(self) -> None:
        empty: dict[str, set[str]] = {}
        q = BooleanQuery([
            Clause(TermQuery("python"), MUST),
            Clause(TermQuery("java"), SHOULD),
        ])
        assert q.evaluate(empty) == set()
