"""Term Position Index for RBSP.

Stores the exact positions where each term appears in each document,
enabling phrase queries ("exact match") and proximity scoring.

Internal storage: ``term -> doc_id -> sorted positions``.

Thread-safe for concurrent reads and writes.
"""

from __future__ import annotations

import threading
from collections import defaultdict
from typing import Dict, List, Set


# ---------------------------------------------------------------------------
# Position Index
# ---------------------------------------------------------------------------

class PositionIndex:
    """Thread-safe term position index for phrase and proximity queries.

    Tracks where each term appears (by position) in each document, which
    enables two capabilities beyond boolean keyword matching:

    1. **Phrase queries** -- verify that terms appear adjacently (or within
       a configurable *slop* distance) in a document.
    2. **Proximity scoring** -- score documents by how close query terms
       appear to each other (closer = higher score).

    Usage::

        pidx = PositionIndex()
        pidx.add_document("doc1", ["the", "quick", "brown", "fox"])
        pidx.phrase_match("doc1", ["quick", "brown"])        # True
        pidx.phrase_match("doc1", ["brown", "quick"])        # False
        pidx.phrase_search(["quick", "brown"])               # {"doc1"}
        pidx.proximity_score("doc1", ["quick", "fox"])       # 0.333...
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        # term -> doc_id -> sorted list of positions
        self._index: Dict[str, Dict[str, List[int]]] = defaultdict(dict)
        # doc_id -> set of terms (for efficient removal)
        self._doc_terms: Dict[str, Set[str]] = {}

    # ------------------------------------------------------------------
    # Mutation
    # ------------------------------------------------------------------

    def add_document(self, doc_id: str, tokens: List[str]) -> None:
        """Index all token positions for a document.

        If the document was already indexed, old positions are replaced.

        Args:
            doc_id: Unique document identifier.
            tokens: Ordered list of tokens as they appear in the document.
        """
        # Build position lists outside the lock for minimal contention.
        positions: Dict[str, List[int]] = defaultdict(list)
        for pos, term in enumerate(tokens):
            positions[term].append(pos)

        with self._lock:
            # Remove old data if re-indexing
            old_terms = self._doc_terms.get(doc_id)
            if old_terms is not None:
                for t in old_terms:
                    posting = self._index.get(t)
                    if posting is not None:
                        posting.pop(doc_id, None)
                        if not posting:
                            del self._index[t]

            self._doc_terms[doc_id] = set(positions.keys())
            for term, pos_list in positions.items():
                self._index[term][doc_id] = pos_list

    def remove_document(self, doc_id: str) -> bool:
        """Remove all position data for a document.

        Args:
            doc_id: The document to remove.

        Returns:
            True if the document was found and removed.
        """
        with self._lock:
            old_terms = self._doc_terms.pop(doc_id, None)
            if old_terms is None:
                return False
            for t in old_terms:
                posting = self._index.get(t)
                if posting is not None:
                    posting.pop(doc_id, None)
                    if not posting:
                        del self._index[t]
            return True

    # ------------------------------------------------------------------
    # Lookup
    # ------------------------------------------------------------------

    def get_positions(self, doc_id: str, term: str) -> List[int]:
        """Get sorted positions of *term* in *doc_id*.

        Args:
            doc_id: The document identifier.
            term:   The term to look up.

        Returns:
            Sorted list of positions, or empty list if not found.
        """
        with self._lock:
            posting = self._index.get(term)
            if posting is None:
                return []
            positions = posting.get(doc_id)
            if positions is None:
                return []
            return list(positions)

    def phrase_match(self, doc_id: str, terms: List[str], slop: int = 0) -> bool:
        """Check if *terms* appear as a phrase in *doc_id*.

        A phrase match requires the terms to appear in the given order with
        consecutive positions.  The *slop* parameter allows gaps: a slop of
        1 means one extra word can appear between adjacent phrase terms.

        Args:
            doc_id: The document to check.
            terms:  Ordered list of terms forming the phrase.
            slop:   Maximum allowed gap between consecutive terms beyond
                    the expected distance of 1 (default 0 = exact phrase).

        Returns:
            True if the phrase matches in the document.
        """
        if not terms:
            return True
        if len(terms) == 1:
            with self._lock:
                posting = self._index.get(terms[0])
                return posting is not None and doc_id in posting

        # Gather positions for each term under the lock.
        with self._lock:
            all_positions: List[List[int]] = []
            for term in terms:
                posting = self._index.get(term)
                if posting is None:
                    return False
                positions = posting.get(doc_id)
                if positions is None:
                    return False
                all_positions.append(positions)

        # Check if any starting position in the first term leads to a
        # valid phrase chain through all subsequent terms.
        return _phrase_check(all_positions, slop)

    def phrase_search(self, terms: List[str], slop: int = 0) -> Set[str]:
        """Find all documents containing *terms* as a phrase.

        Args:
            terms: Ordered list of terms forming the phrase.
            slop:  Maximum allowed gap between consecutive terms (default 0).

        Returns:
            Set of matching document IDs.
        """
        if not terms:
            return set()

        with self._lock:
            # Start with docs that contain all terms (intersection).
            posting_first = self._index.get(terms[0])
            if posting_first is None:
                return set()
            candidates = set(posting_first.keys())

            for term in terms[1:]:
                posting = self._index.get(term)
                if posting is None:
                    return set()
                candidates &= set(posting.keys())
                if not candidates:
                    return set()

            # For each candidate, verify the phrase match.
            results: Set[str] = set()
            for doc_id in candidates:
                all_positions: List[List[int]] = []
                for term in terms:
                    all_positions.append(self._index[term][doc_id])
                if _phrase_check(all_positions, slop):
                    results.add(doc_id)

        return results

    def proximity_score(self, doc_id: str, terms: List[str]) -> float:
        """Score a document based on how close the terms appear together.

        Computes the minimum span (window) covering at least one occurrence
        of every term, then returns ``1.0 / span``.  A span of 1 (all terms
        at the same position) returns 1.0; larger spans yield lower scores.

        If any term is missing from the document, returns 0.0.

        Args:
            doc_id: The document to score.
            terms:  Terms to measure proximity for.

        Returns:
            Float in (0.0, 1.0] where higher is closer, or 0.0 if any
            term is absent.
        """
        if not terms:
            return 0.0

        with self._lock:
            all_positions: List[List[int]] = []
            for term in terms:
                posting = self._index.get(term)
                if posting is None:
                    return 0.0
                positions = posting.get(doc_id)
                if positions is None:
                    return 0.0
                all_positions.append(positions)

        if len(terms) == 1:
            return 1.0

        return _min_span_score(all_positions)

    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------

    def doc_count(self) -> int:
        """Number of indexed documents."""
        with self._lock:
            return len(self._doc_terms)

    def term_count(self) -> int:
        """Number of unique terms in the index."""
        with self._lock:
            return len(self._index)

    def __len__(self) -> int:
        return self.doc_count()

    def __contains__(self, doc_id: str) -> bool:
        with self._lock:
            return doc_id in self._doc_terms

    def clear(self) -> None:
        """Remove all entries."""
        with self._lock:
            self._index.clear()
            self._doc_terms.clear()


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _phrase_check(all_positions: List[List[int]], slop: int) -> bool:
    """Check if a phrase exists across sorted position lists.

    Uses a pointer-based scan: for each starting position in the first
    term's list, try to advance through subsequent terms finding positions
    that are within ``previous_pos + 1 + slop``.

    Args:
        all_positions: Position lists for each term (must be sorted).
        slop: Maximum allowed gap.

    Returns:
        True if a valid phrase chain exists.
    """
    if len(all_positions) == 1:
        return bool(all_positions[0])

    first_positions = all_positions[0]
    for start_pos in first_positions:
        current_pos = start_pos
        matched = True
        for i in range(1, len(all_positions)):
            # Find the smallest position in all_positions[i] that is
            # >= current_pos + 1 and <= current_pos + 1 + slop.
            next_positions = all_positions[i]
            found = False
            lo, hi = 0, len(next_positions)
            # Binary search for first position >= current_pos + 1
            while lo < hi:
                mid = (lo + hi) // 2
                if next_positions[mid] < current_pos + 1:
                    lo = mid + 1
                else:
                    hi = mid
            if lo < len(next_positions) and next_positions[lo] <= current_pos + 1 + slop:
                current_pos = next_positions[lo]
                found = True
            if not found:
                matched = False
                break
        if matched:
            return True
    return False


def _min_span_score(all_positions: List[List[int]]) -> float:
    """Compute proximity score from the minimum span of positions.

    Uses a pointer-advancing algorithm similar to the merge step of
    merge-sort to find the smallest window containing at least one
    position from each term.

    Args:
        all_positions: Position lists for each term (must be sorted,
                       length >= 2).

    Returns:
        ``1.0 / min_span`` where *min_span* is the size of the smallest
        window.  Returns 0.0 if a span cannot be formed.
    """
    n = len(all_positions)
    # Pointers into each position list
    ptrs = [0] * n
    min_span = float("inf")

    while True:
        # Find current min and max positions across all pointers.
        cur_min = float("inf")
        cur_max = float("-inf")
        min_idx = 0
        for i in range(n):
            val = all_positions[i][ptrs[i]]
            if val < cur_min:
                cur_min = val
                min_idx = i
            if val > cur_max:
                cur_max = val

        span = cur_max - cur_min + 1
        if span < min_span:
            min_span = span

        # Advance the pointer with the minimum position.
        ptrs[min_idx] += 1
        if ptrs[min_idx] >= len(all_positions[min_idx]):
            break

    if min_span == float("inf"):
        return 0.0
    return 1.0 / min_span
