"""
RBSP Boolean Query Tree.

Full boolean query tree with cost-based optimization, supporting AND, OR, NOT
with arbitrary nesting, minimum-should-match semantics, and cost-based
execution ordering.

Query node types:
    - TermQuery:            Matches a single term against a doc set.
    - BooleanQuery:         Combines clauses with MUST/SHOULD/MUST_NOT logic.
    - DisjunctionMaxQuery:  Takes the max score across sub-queries (ES dis_max).
    - BoostQuery:           Wraps a query with a score multiplier.
    - ConstantScoreQuery:   Wraps a query, returns a fixed score for all matches.

Clause occur types:
    - MUST:     Document must match (intersection).
    - SHOULD:   Document may match (union with minimum_should_match).
    - MUST_NOT: Document must not match (exclusion).

Optimizer:
    - Flattens nested boolean queries of the same type.
    - Orders MUST clauses cheapest-first (smallest result sets for early pruning).
    - Removes duplicate clauses.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple


# ---------------------------------------------------------------------------
# Occur types
# ---------------------------------------------------------------------------

MUST = "MUST"
SHOULD = "SHOULD"
MUST_NOT = "MUST_NOT"

_VALID_OCCURS = frozenset({MUST, SHOULD, MUST_NOT})


# ---------------------------------------------------------------------------
# Query node base class
# ---------------------------------------------------------------------------

class QueryNode:
    """Base class for all query tree nodes.

    Every node can:
      - ``evaluate(doc_sets)`` -> set of matching document IDs.
      - ``score(doc_id, doc_sets)`` -> relevance score for a single document.
      - ``estimate_cost(doc_sets)`` -> estimated result set size (for optimizer).
    """

    __slots__ = ()

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        """Return the set of document IDs matching this query node.

        Args:
            doc_sets: Mapping of term -> set of document IDs containing that term.

        Returns:
            Set of matching document IDs.
        """
        raise NotImplementedError

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        """Return a relevance score for a single document.

        Default implementation returns 1.0 if the document matches, 0.0 otherwise.

        Args:
            doc_id:   The document to score.
            doc_sets: Mapping of term -> set of document IDs.

        Returns:
            Relevance score >= 0.0.
        """
        return 1.0 if doc_id in self.evaluate(doc_sets) else 0.0

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        """Estimate the result set size for cost-based optimization.

        Args:
            doc_sets: Mapping of term -> set of document IDs.

        Returns:
            Estimated number of matching documents.
        """
        return len(self.evaluate(doc_sets))

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        return self.__dict__ == other.__dict__

    def __hash__(self) -> int:
        return id(self)


# ---------------------------------------------------------------------------
# Clause
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class Clause:
    """A single clause in a BooleanQuery.

    Attributes:
        query: The query node for this clause.
        occur: One of MUST, SHOULD, MUST_NOT.
    """

    query: QueryNode
    occur: str

    def __post_init__(self) -> None:
        if self.occur not in _VALID_OCCURS:
            raise ValueError(
                f"Invalid occur type: {self.occur!r}. "
                f"Must be one of {sorted(_VALID_OCCURS)}"
            )


# ---------------------------------------------------------------------------
# Term query
# ---------------------------------------------------------------------------

class TermQuery(QueryNode):
    """Matches documents containing a single term.

    Attributes:
        term: The term to match against doc_sets keys.
    """

    __slots__ = ("term",)

    def __init__(self, term: str) -> None:
        self.term = term

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        return doc_sets.get(self.term, set())

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        return 1.0 if doc_id in doc_sets.get(self.term, set()) else 0.0

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        return len(doc_sets.get(self.term, set()))

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        return self.term == other.term  # type: ignore[union-attr]

    def __hash__(self) -> int:
        return hash(("TermQuery", self.term))

    def __repr__(self) -> str:
        return f"TermQuery({self.term!r})"


# ---------------------------------------------------------------------------
# Boolean query
# ---------------------------------------------------------------------------

class BooleanQuery(QueryNode):
    """Combines clauses with boolean logic: MUST, SHOULD, MUST_NOT.

    Evaluation semantics (matching Elasticsearch/Lucene):
        1. All MUST clauses are intersected — a document must match every one.
        2. SHOULD clauses are unioned — a document must match at least
           ``minimum_should_match`` of them.  If there are no MUST clauses,
           at least one SHOULD must match by default.
        3. MUST_NOT clauses are excluded — matching documents are removed.

    When both MUST and SHOULD clauses are present, the SHOULD clauses act as
    boosters only (they affect scoring but not filtering) unless
    ``minimum_should_match`` is explicitly > 0.

    Attributes:
        clauses:              List of Clause objects.
        minimum_should_match: Minimum number of SHOULD clauses that must match.
    """

    __slots__ = ("clauses", "minimum_should_match")

    def __init__(
        self,
        clauses: List[Clause],
        minimum_should_match: int = 1,
    ) -> None:
        self.clauses = list(clauses)
        self.minimum_should_match = minimum_should_match

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        must_clauses: List[Clause] = []
        should_clauses: List[Clause] = []
        must_not_clauses: List[Clause] = []

        for clause in self.clauses:
            if clause.occur == MUST:
                must_clauses.append(clause)
            elif clause.occur == SHOULD:
                should_clauses.append(clause)
            else:
                must_not_clauses.append(clause)

        # --- MUST: intersect all ---
        result: Optional[Set[str]] = None
        for clause in must_clauses:
            matches = clause.query.evaluate(doc_sets)
            if result is None:
                result = set(matches)
            else:
                result &= matches
                if not result:
                    return set()  # short-circuit: empty intersection

        # --- SHOULD: union with minimum_should_match ---
        if should_clauses:
            msm = self.minimum_should_match

            if len(should_clauses) == 1 and msm <= 1:
                # Fast path: single SHOULD clause
                should_matches = should_clauses[0].query.evaluate(doc_sets)
            elif msm <= 1:
                # Fast path: just union all
                should_matches: Set[str] = set()
                for clause in should_clauses:
                    should_matches |= clause.query.evaluate(doc_sets)
            else:
                # Count per-doc how many SHOULD clauses match
                should_counts: Dict[str, int] = {}
                for clause in should_clauses:
                    for doc_id in clause.query.evaluate(doc_sets):
                        should_counts[doc_id] = should_counts.get(doc_id, 0) + 1
                should_matches = {
                    doc_id
                    for doc_id, count in should_counts.items()
                    if count >= msm
                }

            if result is None:
                # No MUST clauses — SHOULD drives the result set
                result = should_matches
            elif must_clauses and self.minimum_should_match <= 0:
                # MUST present, msm=0 -> SHOULD is optional (boost only)
                pass
            else:
                # Intersect MUST results with SHOULD results
                result &= should_matches

        if result is None:
            result = set()

        # --- MUST_NOT: exclude ---
        for clause in must_not_clauses:
            if not result:
                break
            excluded = clause.query.evaluate(doc_sets)
            result -= excluded

        return result

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        if doc_id not in self.evaluate(doc_sets):
            return 0.0

        total = 0.0
        count = 0
        for clause in self.clauses:
            if clause.occur == MUST_NOT:
                continue
            s = clause.query.score(doc_id, doc_sets)
            if s > 0.0:
                total += s
                count += 1

        return total / max(count, 1)

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        must_costs = [
            c.query.estimate_cost(doc_sets)
            for c in self.clauses
            if c.occur == MUST
        ]
        should_costs = [
            c.query.estimate_cost(doc_sets)
            for c in self.clauses
            if c.occur == SHOULD
        ]

        if must_costs:
            return min(must_costs)
        if should_costs:
            return sum(should_costs)
        return 0

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        o = other  # type: ignore[assignment]
        return (
            self.minimum_should_match == o.minimum_should_match
            and self.clauses == o.clauses
        )

    def __hash__(self) -> int:
        return hash(("BooleanQuery", self.minimum_should_match, id(self)))

    def __repr__(self) -> str:
        parts = ", ".join(
            f"{c.occur}:{c.query!r}" for c in self.clauses
        )
        return f"BooleanQuery([{parts}], msm={self.minimum_should_match})"


# ---------------------------------------------------------------------------
# Disjunction max query
# ---------------------------------------------------------------------------

class DisjunctionMaxQuery(QueryNode):
    """Returns docs matching any sub-query, scored by max sub-query score.

    This mirrors Elasticsearch's ``dis_max`` query.  The final score is:

        max(sub_scores) + tie_breaker * sum(other_sub_scores)

    Attributes:
        queries:     List of sub-query nodes.
        tie_breaker: Factor for non-max scores (default 0.0).
    """

    __slots__ = ("queries", "tie_breaker")

    def __init__(
        self,
        queries: List[QueryNode],
        tie_breaker: float = 0.0,
    ) -> None:
        self.queries = list(queries)
        self.tie_breaker = tie_breaker

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        result: Set[str] = set()
        for q in self.queries:
            result |= q.evaluate(doc_sets)
        return result

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        scores: List[float] = []
        for q in self.queries:
            s = q.score(doc_id, doc_sets)
            if s > 0.0:
                scores.append(s)

        if not scores:
            return 0.0

        scores.sort(reverse=True)
        best = scores[0]
        tie_sum = sum(scores[1:])
        return best + self.tie_breaker * tie_sum

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        return sum(q.estimate_cost(doc_sets) for q in self.queries)

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        o = other  # type: ignore[assignment]
        return self.tie_breaker == o.tie_breaker and self.queries == o.queries

    def __hash__(self) -> int:
        return hash(("DisjunctionMaxQuery", self.tie_breaker, id(self)))

    def __repr__(self) -> str:
        qs = ", ".join(repr(q) for q in self.queries)
        return f"DisjunctionMaxQuery([{qs}], tie_breaker={self.tie_breaker})"


# ---------------------------------------------------------------------------
# Boost query
# ---------------------------------------------------------------------------

class BoostQuery(QueryNode):
    """Wraps a query and multiplies its score by a boost factor.

    Attributes:
        query: The inner query node.
        boost: The multiplicative score factor.
    """

    __slots__ = ("query", "boost")

    def __init__(self, query: QueryNode, boost: float) -> None:
        self.query = query
        self.boost = boost

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        return self.query.evaluate(doc_sets)

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        return self.query.score(doc_id, doc_sets) * self.boost

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        return self.query.estimate_cost(doc_sets)

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        o = other  # type: ignore[assignment]
        return self.boost == o.boost and self.query == o.query

    def __hash__(self) -> int:
        return hash(("BoostQuery", self.boost, hash(self.query)))

    def __repr__(self) -> str:
        return f"BoostQuery({self.query!r}, boost={self.boost})"


# ---------------------------------------------------------------------------
# Constant score query
# ---------------------------------------------------------------------------

class ConstantScoreQuery(QueryNode):
    """Wraps a query and returns a constant score for all matching documents.

    Useful for filter clauses where relevance scoring is not needed.

    Attributes:
        query:          The inner query node (used only for matching).
        constant_score: The fixed score returned for every match.
    """

    __slots__ = ("query", "constant_score")

    def __init__(self, query: QueryNode, constant_score: float = 1.0) -> None:
        self.query = query
        self.constant_score = constant_score

    def evaluate(self, doc_sets: Dict[str, Set[str]]) -> Set[str]:
        return self.query.evaluate(doc_sets)

    def score(self, doc_id: str, doc_sets: Dict[str, Set[str]]) -> float:
        if doc_id in self.query.evaluate(doc_sets):
            return self.constant_score
        return 0.0

    def estimate_cost(self, doc_sets: Dict[str, Set[str]]) -> int:
        return self.query.estimate_cost(doc_sets)

    def __eq__(self, other: object) -> bool:
        if type(self) is not type(other):
            return NotImplemented
        o = other  # type: ignore[assignment]
        return self.constant_score == o.constant_score and self.query == o.query

    def __hash__(self) -> int:
        return hash(("ConstantScoreQuery", self.constant_score, hash(self.query)))

    def __repr__(self) -> str:
        return f"ConstantScoreQuery({self.query!r}, score={self.constant_score})"


# ---------------------------------------------------------------------------
# Query optimizer
# ---------------------------------------------------------------------------

def optimize_query(query: QueryNode) -> QueryNode:
    """Optimize a query tree for faster execution.

    Transformations applied:
        1. Flatten nested BooleanQuery nodes that share the same structure
           (nested MUST-only into parent MUST, nested SHOULD-only into parent
           SHOULD).
        2. Order MUST clauses cheapest-first so intersections prune early.
        3. Remove duplicate clauses within each BooleanQuery.
        4. Recursively optimize sub-queries.

    Args:
        query: The root query node to optimize.

    Returns:
        An optimized (possibly new) query node.
    """
    if isinstance(query, TermQuery):
        return query

    if isinstance(query, BoostQuery):
        return BoostQuery(optimize_query(query.query), query.boost)

    if isinstance(query, ConstantScoreQuery):
        return ConstantScoreQuery(
            optimize_query(query.query), query.constant_score
        )

    if isinstance(query, DisjunctionMaxQuery):
        optimized_qs = [optimize_query(q) for q in query.queries]
        return DisjunctionMaxQuery(optimized_qs, query.tie_breaker)

    if isinstance(query, BooleanQuery):
        return _optimize_boolean(query)

    return query


def _optimize_boolean(query: BooleanQuery) -> QueryNode:
    """Optimize a BooleanQuery: flatten, deduplicate, reorder."""

    # Step 1: Recursively optimize children and flatten
    must: List[QueryNode] = []
    should: List[QueryNode] = []
    must_not: List[QueryNode] = []

    for clause in query.clauses:
        optimized_child = optimize_query(clause.query)

        # Flatten: if child is a BooleanQuery that can be absorbed
        if isinstance(optimized_child, BooleanQuery):
            flattened = _try_flatten(optimized_child, clause.occur)
            if flattened is not None:
                f_must, f_should, f_must_not = flattened
                must.extend(f_must)
                should.extend(f_should)
                must_not.extend(f_must_not)
                continue

        if clause.occur == MUST:
            must.append(optimized_child)
        elif clause.occur == SHOULD:
            should.append(optimized_child)
        else:
            must_not.append(optimized_child)

    # Step 2: Deduplicate within each group
    must = _deduplicate(must)
    should = _deduplicate(should)
    must_not = _deduplicate(must_not)

    # Step 3: Reorder MUST clauses cheapest-first
    # We use a sentinel cost estimate (None means defer to evaluate-time).
    # For TermQuery nodes we can cheaply estimate via term string length
    # as a proxy; the real cost-ordering happens at evaluate time via
    # estimate_cost, but we can statically sort TermQueries to the front
    # since they are the cheapest node type.
    must = _sort_cheapest_first(must)

    # Rebuild clauses
    new_clauses: List[Clause] = []
    for q in must:
        new_clauses.append(Clause(query=q, occur=MUST))
    for q in should:
        new_clauses.append(Clause(query=q, occur=SHOULD))
    for q in must_not:
        new_clauses.append(Clause(query=q, occur=MUST_NOT))

    if not new_clauses:
        return BooleanQuery([], minimum_should_match=query.minimum_should_match)

    # Single-clause simplification: unwrap when safe
    if (
        len(new_clauses) == 1
        and new_clauses[0].occur == MUST
        and query.minimum_should_match <= 1
    ):
        return new_clauses[0].query

    return BooleanQuery(new_clauses, minimum_should_match=query.minimum_should_match)


def _try_flatten(
    child: BooleanQuery,
    parent_occur: str,
) -> Optional[Tuple[List[QueryNode], List[QueryNode], List[QueryNode]]]:
    """Try to flatten a child BooleanQuery into its parent.

    Flattening is safe when:
        - Parent is MUST and child has only MUST clauses (AND(AND(a,b),c) -> AND(a,b,c))
        - Parent is SHOULD and child has only SHOULD clauses with msm=1
          (OR(OR(a,b),c) -> OR(a,b,c))

    Returns:
        Tuple of (must, should, must_not) query lists if flattened, else None.
    """
    child_must = [c.query for c in child.clauses if c.occur == MUST]
    child_should = [c.query for c in child.clauses if c.occur == SHOULD]
    child_must_not = [c.query for c in child.clauses if c.occur == MUST_NOT]

    # MUST parent + child with only MUST clauses (possibly with MUST_NOT)
    if parent_occur == MUST and not child_should and child_must:
        return (child_must, [], child_must_not)

    # SHOULD parent + child with only SHOULD clauses, msm <= 1, no MUST_NOT
    if (
        parent_occur == SHOULD
        and not child_must
        and not child_must_not
        and child_should
        and child.minimum_should_match <= 1
    ):
        return ([], child_should, [])

    return None


def _deduplicate(queries: List[QueryNode]) -> List[QueryNode]:
    """Remove duplicate query nodes, preserving order of first occurrence."""
    seen: List[QueryNode] = []
    for q in queries:
        if not any(q == existing for existing in seen):
            seen.append(q)
    return seen


def _sort_cheapest_first(queries: List[QueryNode]) -> List[QueryNode]:
    """Sort query nodes so cheaper ones (TermQuery) come first.

    This enables early pruning during MUST intersection — starting with
    the smallest posting list reduces the working set fastest.

    Heuristic ordering:
        1. TermQuery (cheapest — single posting list lookup)
        2. ConstantScoreQuery / BoostQuery wrapping a TermQuery
        3. Everything else (preserved in original order)
    """

    def _cost_key(q: QueryNode) -> Tuple[int, int]:
        # Primary: node type priority, Secondary: tie-break by position
        if isinstance(q, TermQuery):
            return (0, 0)
        inner = _unwrap(q)
        if isinstance(inner, TermQuery):
            return (1, 0)
        return (2, 0)

    # Stable sort preserves original order within each priority bucket
    return sorted(queries, key=_cost_key)


def _unwrap(q: QueryNode) -> QueryNode:
    """Unwrap BoostQuery / ConstantScoreQuery to find the inner node."""
    while isinstance(q, (BoostQuery, ConstantScoreQuery)):
        q = q.query
    return q
