"""
RBSP Phrase Query Evaluation.

Evaluates exact and near-exact phrase matches by checking that query terms
appear in the correct order within indexed documents.  Supports configurable
"slop" — the maximum total positional distance allowed between consecutive
terms beyond their ideal adjacent positions.

Phrase matching algorithm:
    1. Collect positional posting lists for each query term.
    2. Starting from positions of the first term, walk forward through
       subsequent terms checking that each appears within (prev_pos + 1 + slop).
    3. Track the minimum span of each valid phrase occurrence.
    4. Score by occurrence count with a tightness boost for lower actual slop.

Scoring is normalised to [0, 1]:
    raw_score = occurrences * tightness_boost
    final_score = raw_score / (raw_score + 1)   (saturating normalisation)
"""

from __future__ import annotations

import heapq
from dataclasses import dataclass
from typing import Dict, List, Set, Tuple


@dataclass(frozen=True, slots=True)
class PhraseQuery:
    """A phrase query requiring ordered term proximity.

    Attributes:
        terms: Ordered tuple of terms that must appear in sequence.
        slop:  Maximum total positional slack allowed across the phrase.
               0 means terms must be strictly adjacent.
    """

    terms: tuple[str, ...]
    slop: int = 0


class PhraseScorer:
    """Scores documents against phrase queries using positional data.

    Args:
        position_data: Mapping of term -> doc_id -> sorted list of positions.
                       Positions are 0-based token offsets within the document.
    """

    def __init__(self, position_data: Dict[str, Dict[str, List[int]]]) -> None:
        self._positions = position_data

    def score_document(self, query: PhraseQuery, doc_id: str) -> float:
        """Compute a phrase match score for a single document.

        Returns:
            Score in [0.0, 1.0].  0.0 if no phrase match is found.
        """
        if len(query.terms) == 0:
            return 0.0

        # Gather per-term position lists for this document
        positions_list: List[List[int]] = []
        for term in query.terms:
            term_positions = self._positions.get(term, {})
            doc_positions = term_positions.get(doc_id, [])
            if not doc_positions:
                return 0.0
            positions_list.append(doc_positions)

        # Single-term "phrase" — presence is a match
        if len(positions_list) == 1:
            count = len(positions_list[0])
            return count / (count + 1.0)

        phrase_hits = self._find_phrase_positions(positions_list, query.slop)
        if not phrase_hits:
            return 0.0

        # Each hit is the starting position; compute span-based tightness
        ideal_span = len(query.terms) - 1
        total_boost = 0.0
        for start_pos, actual_span in phrase_hits:
            if ideal_span == 0:
                tightness = 1.0
            else:
                tightness = ideal_span / max(actual_span, ideal_span)
            total_boost += tightness

        # Saturating normalisation: score = raw / (raw + 1)
        return total_boost / (total_boost + 1.0)

    def find_matches(
        self,
        query: PhraseQuery,
        candidate_docs: Set[str],
    ) -> List[Tuple[str, float]]:
        """Score all candidate documents against a phrase query.

        Args:
            query:          The phrase query to evaluate.
            candidate_docs: Set of document IDs to check.

        Returns:
            List of (doc_id, score) tuples for documents with non-zero
            scores, sorted by descending score.
        """
        scored: List[Tuple[float, str]] = []
        for doc_id in candidate_docs:
            score = self.score_document(query, doc_id)
            if score > 0.0:
                scored.append((score, doc_id))

        # Return sorted by descending score using heapq for efficiency
        top = heapq.nlargest(len(scored), scored, key=lambda x: x[0])
        return [(doc_id, score) for score, doc_id in top]

    @staticmethod
    def _find_phrase_positions(
        positions_list: List[List[int]],
        slop: int,
    ) -> List[Tuple[int, int]]:
        """Find positions where a phrase occurs within the allowed slop.

        Uses a pointer-advancement algorithm over sorted position lists.
        For each starting position of the first term, greedily finds the
        closest valid position of each subsequent term.

        Args:
            positions_list: List of sorted position lists, one per term.
            slop:           Maximum total slack beyond ideal adjacent placement.

        Returns:
            List of (start_position, actual_span) for each phrase occurrence.
            actual_span is the distance from first term position to last term
            position.
        """
        if len(positions_list) < 2:
            return [(p, 0) for p in positions_list[0]] if positions_list else []

        results: List[Tuple[int, int]] = []

        for start_pos in positions_list[0]:
            prev_pos = start_pos
            total_span = 0
            matched = True

            for term_idx in range(1, len(positions_list)):
                next_positions = positions_list[term_idx]
                # Find the smallest position in next_positions that is
                # >= prev_pos + 1 (must come after) and
                # <= prev_pos + 1 + slop (within slop window)
                best_pos = _find_nearest_in_window(
                    next_positions, prev_pos + 1, prev_pos + 1 + slop
                )
                if best_pos is None:
                    matched = False
                    break
                total_span += best_pos - prev_pos
                prev_pos = best_pos

            if matched:
                results.append((start_pos, total_span))

        return results


def _find_nearest_in_window(
    sorted_positions: List[int], low: int, high: int
) -> int | None:
    """Binary search for the smallest position in [low, high].

    Args:
        sorted_positions: Ascending list of positions.
        low:  Minimum acceptable position (inclusive).
        high: Maximum acceptable position (inclusive).

    Returns:
        The smallest position in range, or None if no position qualifies.
    """
    lo, hi = 0, len(sorted_positions) - 1
    result = None

    while lo <= hi:
        mid = (lo + hi) // 2
        val = sorted_positions[mid]
        if val < low:
            lo = mid + 1
        elif val > high:
            hi = mid - 1
        else:
            result = val
            hi = mid - 1  # look for an even smaller qualifying value

    return result
