"""
RBSP Query Plan Optimizer.

Analyzes a query and chooses the optimal execution strategy based on term
statistics, analogous to a database query planner.  Given document-frequency
information the planner decides:

    - **Term ordering**: for AND queries, rarest terms first (smallest posting
      list drives the intersection); for OR queries, most common terms first
      (biggest lists first for early accumulation).
    - **Cost estimation**: posting-list sizes with a selectivity shrinkage
      factor for intersections.
    - **Boolean planning**: complex multi-clause queries decomposed into
      ordered steps.

Cost model:
    - Single term lookup:  cost = df(term)
    - AND intersection:    cost = min(df(t1), df(t2)) * selectivity_factor
    - OR union:            cost = df(t1) + df(t2)
    - Selectivity factor:  0.1 (independent terms assumption)

Zero external dependencies.
"""

from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_SELECTIVITY_FACTOR: float = 0.1


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class PlanStep:
    """A single step in a query execution plan.

    Attributes:
        operation:      The operation type (e.g. "lookup", "intersect",
                        "union", "negate").
        term:           The term this step operates on (empty for
                        intermediate operations).
        estimated_docs: Estimated number of documents touched.
        cost:           Estimated cost for this step.
    """

    operation: str
    term: str = ""
    estimated_docs: int = 0
    cost: float = 0.0


@dataclass(frozen=True, slots=True)
class QueryPlan:
    """A complete query execution plan.

    Attributes:
        steps:          Ordered list of plan steps to execute.
        estimated_cost: Total estimated cost of the plan.
        strategy:       Human-readable strategy label (e.g.
                        "and_rarest_first", "or_largest_first",
                        "single_term", "boolean_mixed").
    """

    steps: list[PlanStep]
    estimated_cost: float
    strategy: str


# ---------------------------------------------------------------------------
# QueryPlanner
# ---------------------------------------------------------------------------

class QueryPlanner:
    """Plans optimal query execution order based on term statistics.

    Args:
        doc_freqs:  Mapping of term -> document frequency.  When *None*,
                    all unknown terms are assumed to have frequency 0.
        total_docs: Total number of documents in the index.  Used for
                    relative frequency calculations.
    """

    def __init__(
        self,
        doc_freqs: Optional[Dict[str, int]] = None,
        total_docs: int = 0,
    ) -> None:
        self._doc_freqs: Dict[str, int] = doc_freqs if doc_freqs is not None else {}
        self._total_docs: int = total_docs

    # -- public helpers -----------------------------------------------------

    @property
    def doc_freqs(self) -> Dict[str, int]:
        """Current document-frequency map (read-only copy)."""
        return dict(self._doc_freqs)

    @property
    def total_docs(self) -> int:
        return self._total_docs

    def _df(self, term: str) -> int:
        """Return the document frequency of *term*, defaulting to 0."""
        return self._doc_freqs.get(term, 0)

    # -- plan() -------------------------------------------------------------

    def plan(
        self,
        query_terms: List[str],
        query_type: str = "AND",
    ) -> QueryPlan:
        """Build an execution plan for a flat term query.

        Args:
            query_terms: List of query terms.
            query_type:  ``"AND"`` or ``"OR"``.

        Returns:
            A ``QueryPlan`` with ordered steps and cost estimate.
        """
        if not query_terms:
            return QueryPlan(steps=[], estimated_cost=0.0, strategy="empty")

        if len(query_terms) == 1:
            term = query_terms[0]
            df = self._df(term)
            step = PlanStep(
                operation="lookup",
                term=term,
                estimated_docs=df,
                cost=float(df),
            )
            return QueryPlan(
                steps=[step],
                estimated_cost=float(df),
                strategy="single_term",
            )

        upper = query_type.upper()
        if upper == "AND":
            return self._plan_and(query_terms)
        if upper == "OR":
            return self._plan_or(query_terms)
        # Default to AND for unknown types
        return self._plan_and(query_terms)

    # -- AND planning -------------------------------------------------------

    def _plan_and(self, terms: List[str]) -> QueryPlan:
        """Plan an AND query: rarest term first."""
        # Sort by ascending document frequency (rarest first)
        sorted_terms = sorted(terms, key=lambda t: self._df(t))

        steps: List[PlanStep] = []
        total_cost: float = 0.0
        running_docs: float = 0.0

        for i, term in enumerate(sorted_terms):
            df = self._df(term)
            if i == 0:
                # First term: full posting list scan
                step_cost = float(df)
                running_docs = float(df)
                steps.append(PlanStep(
                    operation="lookup",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            else:
                # Intersection: cost bounded by current candidate set
                intersect_docs = running_docs * _SELECTIVITY_FACTOR
                step_cost = min(running_docs, float(df))
                running_docs = intersect_docs
                steps.append(PlanStep(
                    operation="intersect",
                    term=term,
                    estimated_docs=max(int(intersect_docs), 0),
                    cost=step_cost,
                ))
            total_cost += step_cost

        return QueryPlan(
            steps=steps,
            estimated_cost=total_cost,
            strategy="and_rarest_first",
        )

    # -- OR planning --------------------------------------------------------

    def _plan_or(self, terms: List[str]) -> QueryPlan:
        """Plan an OR query: most common term first."""
        # Sort by descending document frequency (most common first)
        sorted_terms = sorted(terms, key=lambda t: self._df(t), reverse=True)

        steps: List[PlanStep] = []
        total_cost: float = 0.0

        for i, term in enumerate(sorted_terms):
            df = self._df(term)
            step_cost = float(df)
            if i == 0:
                steps.append(PlanStep(
                    operation="lookup",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            else:
                steps.append(PlanStep(
                    operation="union",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            total_cost += step_cost

        return QueryPlan(
            steps=steps,
            estimated_cost=total_cost,
            strategy="or_largest_first",
        )

    # -- plan_boolean() -----------------------------------------------------

    def plan_boolean(self, clauses: List[Dict]) -> QueryPlan:
        """Plan a complex boolean query with mixed AND / OR / NOT clauses.

        Each clause is a dict with:
            - ``"op"``:   ``"AND"`` | ``"OR"`` | ``"NOT"``
            - ``"term"``: the query term

        The planner groups clauses by operator and orders them:
            1. NOT clauses (cheapest negation filters applied last).
            2. AND clauses ordered rarest-first.
            3. OR clauses ordered most-common-first.

        Args:
            clauses: List of clause dicts.

        Returns:
            A ``QueryPlan`` for the boolean expression.
        """
        if not clauses:
            return QueryPlan(steps=[], estimated_cost=0.0, strategy="empty")

        and_terms: List[str] = []
        or_terms: List[str] = []
        not_terms: List[str] = []

        for clause in clauses:
            op = clause.get("op", "AND").upper()
            term = clause.get("term", "")
            if not term:
                continue
            if op == "NOT":
                not_terms.append(term)
            elif op == "OR":
                or_terms.append(term)
            else:
                and_terms.append(term)

        steps: List[PlanStep] = []
        total_cost: float = 0.0
        running_docs: float = 0.0

        # Phase 1: AND terms (rarest first) — drive the candidate set down
        and_sorted = sorted(and_terms, key=lambda t: self._df(t))
        for i, term in enumerate(and_sorted):
            df = self._df(term)
            if i == 0 and not steps:
                step_cost = float(df)
                running_docs = float(df)
                steps.append(PlanStep(
                    operation="lookup",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            else:
                intersect_docs = running_docs * _SELECTIVITY_FACTOR
                step_cost = min(running_docs, float(df))
                running_docs = intersect_docs
                steps.append(PlanStep(
                    operation="intersect",
                    term=term,
                    estimated_docs=max(int(intersect_docs), 0),
                    cost=step_cost,
                ))
            total_cost += step_cost

        # Phase 2: OR terms (most common first) — expand the result set
        or_sorted = sorted(or_terms, key=lambda t: self._df(t), reverse=True)
        for term in or_sorted:
            df = self._df(term)
            step_cost = float(df)
            if not steps:
                running_docs = float(df)
                steps.append(PlanStep(
                    operation="lookup",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            else:
                running_docs += float(df)
                steps.append(PlanStep(
                    operation="union",
                    term=term,
                    estimated_docs=df,
                    cost=step_cost,
                ))
            total_cost += step_cost

        # Phase 3: NOT terms — filter at the end (cheapest when candidate
        # set is already small)
        not_sorted = sorted(not_terms, key=lambda t: self._df(t))
        for term in not_sorted:
            df = self._df(term)
            step_cost = running_docs  # scan remaining candidates
            running_docs = max(running_docs - float(df), 0.0)
            steps.append(PlanStep(
                operation="negate",
                term=term,
                estimated_docs=max(int(running_docs), 0),
                cost=step_cost,
            ))
            total_cost += step_cost

        strategy = "boolean_mixed"
        if not or_terms and not not_terms:
            strategy = "and_rarest_first"
        elif not and_terms and not not_terms:
            strategy = "or_largest_first"

        return QueryPlan(
            steps=steps,
            estimated_cost=total_cost,
            strategy=strategy,
        )

    # -- explain() ----------------------------------------------------------

    def explain(self, plan: QueryPlan) -> str:
        """Return a human-readable explanation of a query plan.

        Args:
            plan: The query plan to explain.

        Returns:
            Multi-line string describing the strategy, each step, and the
            total estimated cost.
        """
        lines: List[str] = []
        lines.append(f"Strategy: {plan.strategy}")
        lines.append(f"Estimated total cost: {plan.estimated_cost:.1f}")
        lines.append(f"Steps ({len(plan.steps)}):")

        for i, step in enumerate(plan.steps, start=1):
            term_part = f" '{step.term}'" if step.term else ""
            lines.append(
                f"  {i}. {step.operation}{term_part}"
                f" | est_docs={step.estimated_docs}"
                f" | cost={step.cost:.1f}"
            )

        return "\n".join(lines)
