dotfiles/python/ebook_search/protected_phrases/matching.py

"""Runtime protected-phrase matching and chunk mention indexing."""

from __future__ import annotations

import logging
from collections import defaultdict
from typing import TYPE_CHECKING

from sqlalchemy import and_, delete, or_, select, union

from python.ebook_search.protected_phrases.config import get_ignored_phrases
from python.ebook_search.protected_phrases.models import (
    PhraseLookup,
    PhraseMatch,
)
from python.ebook_search.protected_phrases.text_normalization import tokenize_with_offsets
from python.orm.richie import (
    EbookChunk,
    EbookChunkPhraseMention,
    EbookPhraseAlias,
    EbookProtectedPhrase,
)

if TYPE_CHECKING:
    from collections.abc import Iterator, Sequence

    from sqlalchemy.ext.asyncio import AsyncSession

    from python.ebook_search.config import EbookSearchConfig

logger = logging.getLogger(__name__)


async def detect_protected_phrases_for_query(
    session: AsyncSession,
    query_text: str,
    config: EbookSearchConfig,
) -> list[PhraseMatch]:
    """Find query phrases with indexed exact matches on canonical and alias norms.

    Args:
        session (AsyncSession): Active database session.
        query_text (str): User query text to detect phrases in.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        list[PhraseMatch]: Metadata-backed, overlap-resolved phrase matches for the query.
    """
    tokens_ = tokenize_with_offsets(query_text)
    windows_by_norm: defaultdict[str, list[tuple[int, int]]] = defaultdict(list)
    token_texts = [token.text for token in tokens_]
    max_tokens = max(config.phrase_max_tokens, config.phrase_max_entity_tokens)
    for phrase_norm, start, end in generate_query_ngrams(
        token_texts,
        min_n=config.phrase_min_tokens,
        max_n=max_tokens,
    ):
        windows_by_norm[phrase_norm].append((start, end))
    if not windows_by_norm:
        return []

    query_norms = tuple(windows_by_norm)
    matched_norms = union(
        select(
            EbookProtectedPhrase.id.label("phrase_id"),
            EbookProtectedPhrase.phrase_norm.label("matched_norm"),
        ).where(EbookProtectedPhrase.phrase_norm.in_(query_norms)),
        select(
            EbookPhraseAlias.phrase_id.label("phrase_id"),
            EbookPhraseAlias.alias_norm.label("matched_norm"),
        ).where(EbookPhraseAlias.alias_norm.in_(query_norms)),
    ).subquery()
    statement = select(EbookProtectedPhrase, matched_norms.c.matched_norm).join(
        matched_norms,
        matched_norms.c.phrase_id == EbookProtectedPhrase.id,
    )

    matches: list[PhraseMatch] = []
    for phrase, matched_norm in await session.execute(statement):
        for start, end in windows_by_norm[matched_norm]:
            matches.append(
                PhraseMatch(
                    phrase_id=phrase.id,
                    matched_norm=matched_norm,
                    phrase_text=phrase.phrase_text,
                    phrase_norm=phrase.phrase_norm,
                    canonical_id=phrase.canonical_id,
                    phrase_type=phrase.phrase_type,
                    token_count=end - start,
                    confidence=phrase.confidence,
                    importance=phrase.importance,
                    allow_nested=phrase.allow_nested,
                    suppress_children=phrase.suppress_children,
                    start_token=start,
                    end_token=end,
                    start_char=tokens_[start].start_char,
                    end_char=tokens_[end - 1].end_char,
                    book_id=phrase.book_id,
                    series_id=phrase.series_id,
                )
            )
    return resolve_overlaps(matches)


async def index_chunk_phrase_mentions_for_book(
    session: AsyncSession,
    book_id: int,
    config: EbookSearchConfig,
) -> int:
    """Rebuild chunk phrase mentions for all chunks in one book.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose chunk mentions are rebuilt.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        int: Total number of chunk phrase mentions indexed for the book.
    """
    lookup = await load_phrase_lookup(session, config, book_id=book_id)
    await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
    chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
    count = 0
    for chunk in chunks:
        count += index_chunk_phrase_mentions(session, chunk, lookup=lookup)
    await session.flush()
    logger.info(f"ebook_chunk_phrase_mentions_indexed {book_id=} {count=}")
    return count


async def load_phrase_lookup(
    session: AsyncSession,
    config: EbookSearchConfig,
    *,
    book_id: int | None = None,
    series_id: int | None = None,
) -> PhraseLookup:
    """Load protected phrases and aliases into RAM lookup maps.

    Args:
        session (AsyncSession): Active database session.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        book_id (int | None): Optional book scope to restrict loaded phrases.
        series_id (int | None): Optional series scope to restrict loaded phrases.

    Returns:
        PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test.
    """
    phrase_ids_by_norm: defaultdict[str, set[int]] = defaultdict(set)
    phrases_by_id: dict[int, EbookProtectedPhrase] = {}
    max_tokens = config.phrase_max_tokens

    statement = select(
        EbookProtectedPhrase,
        EbookPhraseAlias.alias_norm,
    ).outerjoin(EbookPhraseAlias, EbookPhraseAlias.phrase_id == EbookProtectedPhrase.id)
    scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
    if scope_filter is not None:
        statement = statement.where(scope_filter)

    for phrase, alias_norm in await session.execute(statement):
        phrases_by_id[phrase.id] = phrase
        phrase_ids_by_norm[phrase.phrase_norm].add(phrase.id)
        max_tokens = max(max_tokens, phrase.token_count)
        if alias_norm is not None:
            phrase_ids_by_norm[alias_norm].add(phrase.id)
            max_tokens = max(max_tokens, len(alias_norm.split()))

    return PhraseLookup(
        phrase_ids_by_norm={key: tuple(sorted(values)) for key, values in phrase_ids_by_norm.items()},
        phrases_by_id=phrases_by_id,
        min_tokens=config.phrase_min_tokens,
        max_tokens=max_tokens,
    )


def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None) -> object | None:
    """Build a SQLAlchemy filter for optional phrase book and series scope.

    Args:
        book_id (int | None): Optional book scope to include alongside global phrases.
        series_id (int | None): Optional series scope to include alongside global phrases.

    Returns:
        object | None: A combined SQLAlchemy filter clause, or ``None`` when no scope is given.
    """
    conditions = []
    if book_id is not None:
        conditions.append(or_(EbookProtectedPhrase.book_id.is_(None), EbookProtectedPhrase.book_id == book_id))
    if series_id is not None:
        conditions.append(or_(EbookProtectedPhrase.series_id.is_(None), EbookProtectedPhrase.series_id == series_id))
    if not conditions:
        return None
    return and_(*conditions)


def is_inside(child: PhraseMatch, parent: PhraseMatch) -> bool:
    """Return whether one token span is strictly inside another.

    Args:
        child (PhraseMatch): Candidate nested match.
        parent (PhraseMatch): Candidate enclosing match.

    Returns:
        bool: True when ``child`` lies within ``parent`` and is not the same span.
    """
    return (
        child.start_token >= parent.start_token
        and child.end_token <= parent.end_token
        and (child.start_token, child.end_token, child.phrase_id)
        != (parent.start_token, parent.end_token, parent.phrase_id)
    )


def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
    """Store protected phrase mentions for one chunk.

    Args:
        session (AsyncSession): Active database session.
        chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
        lookup (PhraseLookup): In-memory phrase and alias lookup maps.

    Returns:
        int: Number of phrase mentions stored for the chunk.
    """
    tokens_ = tokenize_with_offsets(chunk.text)
    token_texts = [token.text for token in tokens_]
    raw_matches: list[PhraseMatch] = []
    phrase_windows = generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens)
    for matched_norm, start, end in phrase_windows:
        for phrase_id in lookup.phrase_ids_by_norm.get(matched_norm, ()):
            phrase = lookup.phrases_by_id[phrase_id]
            raw_matches.append(
                PhraseMatch(
                    phrase_id=phrase_id,
                    matched_norm=matched_norm,
                    phrase_text=phrase.phrase_text,
                    phrase_norm=phrase.phrase_norm,
                    canonical_id=phrase.canonical_id,
                    phrase_type=phrase.phrase_type,
                    confidence=phrase.confidence,
                    importance=phrase.importance,
                    allow_nested=phrase.allow_nested,
                    suppress_children=phrase.suppress_children,
                    start_token=start,
                    end_token=end,
                    token_count=end - start,
                    start_char=tokens_[start].start_char,
                    end_char=tokens_[end - 1].end_char,
                    book_id=phrase.book_id,
                    series_id=phrase.series_id,
                )
            )

    matches = resolve_overlaps(raw_matches)
    for match in matches:
        session.add(
            EbookChunkPhraseMention(
                chunk_id=chunk.id,
                phrase_id=match.phrase_id,
                book_id=match.book_id if match.book_id is not None else chunk.source_id,
                series_id=match.series_id,
                start_char=match.start_char if match.start_char is not None else 0,
                end_char=match.end_char,
            )
        )
    return len(matches)


def resolve_overlaps(matches: Sequence[PhraseMatch]) -> list[PhraseMatch]:
    """Resolve overlapping phrase matches without relying only on longest match.

    Args:
        matches (Sequence[PhraseMatch]): Metadata-backed matches that may overlap.

    Returns:
        list[PhraseMatch]: The kept, non-suppressed matches.
    """
    sorted_matches = sorted(
        matches,
        key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
    )
    kept: list[PhraseMatch] = []
    for candidate in sorted_matches:
        if any(should_suppress(candidate, existing) for existing in kept):
            continue
        kept.append(candidate)
    return kept


def should_suppress(candidate: PhraseMatch, kept: PhraseMatch) -> bool:
    """Return whether an already-kept match should suppress a candidate.

    Args:
        candidate (PhraseMatch): Match being considered for keeping.
        kept (PhraseMatch): Match already kept that may suppress the candidate.

    Returns:
        bool: True when the candidate should be dropped in favor of the kept match.
    """
    if not overlaps(candidate, kept):
        return False
    if candidate.canonical_id == kept.canonical_id:
        return rank_match(kept) >= rank_match(candidate)
    if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
        return True
    return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)


def overlaps(first: PhraseMatch, second: PhraseMatch) -> bool:
    """Return whether two token spans overlap.

    Args:
        first (PhraseMatch): First match to compare.
        second (PhraseMatch): Second match to compare.

    Returns:
        bool: True when the two token spans share at least one token position.
    """
    return not (first.end_token <= second.start_token or first.start_token >= second.end_token)


def rank_match(match: PhraseMatch) -> tuple[float, float, int]:
    """Rank phrase matches by importance, confidence, then token count.

    Args:
        match (PhraseMatch): Match to build a sort key for.

    Returns:
        tuple[float, float, int]: A comparable key of importance, confidence, and token count.
    """
    return (match.importance, match.confidence, match.token_count)


def generate_query_ngrams(
    tokens_: Sequence[str],
    min_n: int,
    max_n: int,
) -> Iterator[tuple[str, int, int]]:
    """Generate normalized query windows from longest to shortest.

    Args:
        tokens_ (Sequence[str]): Normalized query tokens.
        min_n (int): Smallest window size to yield.
        max_n (int): Largest window size to yield, capped at the token count.

    Yields:
        tuple[str, int, int]: Normalized window text with its start and end token indices.
    """
    capped_max_n = min(max_n, len(tokens_))
    for ngram_size in range(capped_max_n, min_n - 1, -1):
        for start in range(len(tokens_) - ngram_size + 1):
            end = start + ngram_size
            phrase_norm = " ".join(tokens_[start:end])
            if phrase_norm in get_ignored_phrases():
                continue
            yield phrase_norm, start, end