"""Inverted Index for RBSP.

Maps tokens (words, code identifiers) to sets of signature IDs, enabling
O(1) single-term lookup and fast multi-term boolean queries via set operations.

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

from __future__ import annotations

import re
import threading
from collections import defaultdict
from typing import Dict, FrozenSet, Iterable, List, Optional, Set


# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------

# Splits on whitespace and common non-alphanumeric chars, also breaks
# camelCase / PascalCase.
_CAMEL_RE = re.compile(r"(?<=[a-z])(?=[A-Z])")
_SPLIT_RE = re.compile(r"[^a-zA-Z0-9]+")

# Minimal stop words to reduce noise without a heavyweight NLP dependency.
_STOP_WORDS: FrozenSet[str] = frozenset({
    "the", "a", "an", "and", "or", "not", "is", "it", "in", "of", "to",
    "for", "on", "with", "as", "at", "by", "be", "if", "no", "do",
})


def tokenize(text: str) -> List[str]:
    """Extract indexable tokens from *text*.

    Splits camelCase, lowercases, removes tokens shorter than 2 characters,
    and strips common stop words.

    Args:
        text: Arbitrary text (natural language, source code, etc.).

    Returns:
        Ordered list of extracted tokens (may contain duplicates).
    """
    # Break camelCase / PascalCase
    spaced = _CAMEL_RE.sub(" ", text)
    parts = _SPLIT_RE.split(spaced)
    tokens: List[str] = []
    for p in parts:
        t = p.lower().strip()
        if len(t) >= 2 and t not in _STOP_WORDS:
            tokens.append(t)
    return tokens


def tokenize_unique(text: str) -> Set[str]:
    """Like :func:`tokenize` but returns a deduplicated set."""
    return set(tokenize(text))


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

class _QueryNode:
    """Base class for query AST nodes."""

    def evaluate(self, index: "InvertedIndex") -> Set[str]:
        raise NotImplementedError


class TermNode(_QueryNode):
    """A single-term query."""

    __slots__ = ("term",)

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

    def evaluate(self, index: "InvertedIndex") -> Set[str]:
        return index.lookup(self.term)

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


class AndNode(_QueryNode):
    """Intersection of children."""

    __slots__ = ("children",)

    def __init__(self, children: List[_QueryNode]) -> None:
        self.children = children

    def evaluate(self, index: "InvertedIndex") -> Set[str]:
        if not self.children:
            return set()
        result = self.children[0].evaluate(index)
        for child in self.children[1:]:
            result = result & child.evaluate(index)
            if not result:
                break
        return result

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


class OrNode(_QueryNode):
    """Union of children."""

    __slots__ = ("children",)

    def __init__(self, children: List[_QueryNode]) -> None:
        self.children = children

    def evaluate(self, index: "InvertedIndex") -> Set[str]:
        result: Set[str] = set()
        for child in self.children:
            result = result | child.evaluate(index)
        return result

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


class NotNode(_QueryNode):
    """Complement relative to the full document set."""

    __slots__ = ("child",)

    def __init__(self, child: _QueryNode) -> None:
        self.child = child

    def evaluate(self, index: "InvertedIndex") -> Set[str]:
        excluded = self.child.evaluate(index)
        return index.all_doc_ids() - excluded

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


def parse_query(query_str: str) -> _QueryNode:
    """Parse a simple boolean query string.

    Supported syntax::

        word                    -> single term
        word1 AND word2         -> intersection
        word1 OR word2          -> union
        NOT word                -> complement
        word1 AND word2 OR word3  -> left-to-right (AND binds tighter)

    Parentheses are **not** supported in this minimal parser.  For complex
    queries, construct AST nodes directly.

    Args:
        query_str: The query string.

    Returns:
        A query AST node.
    """
    raw_tokens = query_str.strip().split()
    if not raw_tokens:
        return TermNode("")

    nodes: List[_QueryNode] = []
    pending_op: Optional[str] = None
    i = 0
    while i < len(raw_tokens):
        tok = raw_tokens[i]
        upper = tok.upper()

        if upper == "AND":
            pending_op = "AND"
            i += 1
            continue
        elif upper == "OR":
            pending_op = "OR"
            i += 1
            continue
        elif upper == "NOT":
            i += 1
            if i < len(raw_tokens):
                inner = TermNode(raw_tokens[i])
                node: _QueryNode = NotNode(inner)
            else:
                node = TermNode("")
        else:
            node = TermNode(tok)

        if pending_op == "AND" and nodes:
            prev = nodes.pop()
            node = AndNode([prev, node])
        elif pending_op == "OR" and nodes:
            prev = nodes.pop()
            node = OrNode([prev, node])
        pending_op = None
        nodes.append(node)
        i += 1

    if len(nodes) == 1:
        return nodes[0]
    # Default: implicit AND
    return AndNode(nodes)


# ---------------------------------------------------------------------------
# Inverted Index
# ---------------------------------------------------------------------------

class InvertedIndex:
    """Thread-safe inverted index mapping tokens to document/signature IDs.

    Provides O(1) single-term lookup and boolean query evaluation.

    Usage::

        idx = InvertedIndex()
        idx.add_document("sig_abc", ["python", "hash", "index"])
        hits = idx.lookup("python")          # {"sig_abc"}
        hits = idx.search("python AND hash") # {"sig_abc"}
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._index: Dict[str, Set[str]] = defaultdict(set)
        self._doc_tokens: Dict[str, Set[str]] = {}

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

    def add_document(self, doc_id: str, tokens: Iterable[str]) -> None:
        """Index a document.

        Args:
            doc_id: Unique identifier (typically a signature ID).
            tokens: Iterable of tokens to associate with *doc_id*.
        """
        token_set = set(tokens)
        with self._lock:
            # Remove old tokens if re-indexing the same doc
            old = self._doc_tokens.get(doc_id)
            if old is not None:
                for t in old:
                    self._index[t].discard(doc_id)
                    if not self._index[t]:
                        del self._index[t]
            self._doc_tokens[doc_id] = token_set
            for t in token_set:
                self._index[t].add(doc_id)

    def add_text(self, doc_id: str, text: str) -> None:
        """Tokenize *text* and index under *doc_id*."""
        tokens = tokenize_unique(text)
        self.add_document(doc_id, tokens)

    def remove_document(self, doc_id: str) -> bool:
        """Remove a document from the index.

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

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

    def lookup(self, term: str) -> Set[str]:
        """O(1) lookup: return all doc_ids containing *term*.

        Returns a **copy** of the internal set so callers can mutate freely.
        """
        with self._lock:
            return set(self._index.get(term.lower(), set()))

    def all_doc_ids(self) -> Set[str]:
        """Return the set of all indexed document IDs."""
        with self._lock:
            return set(self._doc_tokens.keys())

    def search(self, query_str: str) -> Set[str]:
        """Evaluate a boolean query string and return matching doc IDs.

        See :func:`parse_query` for supported syntax.
        """
        ast = parse_query(query_str)
        return ast.evaluate(self)

    def query(self, node: _QueryNode) -> Set[str]:
        """Evaluate a pre-built query AST node."""
        return node.evaluate(self)

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

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

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

    def term_frequency(self, term: str) -> int:
        """Number of documents containing *term*."""
        with self._lock:
            return len(self._index.get(term.lower(), set()))

    def doc_freqs(self) -> Dict[str, int]:
        """Return {term: num_docs_containing_term} for all terms.

        Returns a snapshot copy suitable for passing to BM25Scorer.update_stats().
        """
        with self._lock:
            return {t: len(docs) for t, docs in self._index.items()}

    def doc_freqs_for_terms(self, terms: Iterable[str]) -> Dict[str, int]:
        """Return document frequencies only for the given terms.

        More efficient than doc_freqs() when you only need a subset.

        Args:
            terms: Iterable of terms to look up.

        Returns:
            Dict mapping each term to its document frequency.
        """
        with self._lock:
            result: Dict[str, int] = {}
            for t in terms:
                t_lower = t.lower()
                posting = self._index.get(t_lower)
                if posting:
                    result[t_lower] = len(posting)
            return result

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

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

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