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, func, or_, select

from python.ebook_search.protected_phrases.config import get_ignored_phrases
from python.ebook_search.protected_phrases.models import (
    ChunkPhraseHit,
    HydratedPhraseMatch,
    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
    from python.ebook_search.protected_phrases.text_normalization import NormalizedToken

logger = logging.getLogger(__name__)


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.
    """
    norm_to_ids: defaultdict[str, list[int]] = defaultdict(list)
    alias_to_ids: defaultdict[str, list[int]] = defaultdict(list)
    max_tokens = config.phrase_max_tokens

    phrase_statement = select(
        EbookProtectedPhrase.id,
        EbookProtectedPhrase.phrase_norm,
        EbookProtectedPhrase.token_count,
    )
    scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
    if scope_filter is not None:
        phrase_statement = phrase_statement.where(scope_filter)

    for row in await session.execute(phrase_statement):
        phrase_id = int(row.id)
        phrase_norm = str(row.phrase_norm)
        norm_to_ids[phrase_norm].append(phrase_id)
        max_tokens = max(max_tokens, int(row.token_count))

    alias_statement = select(
        EbookPhraseAlias.alias_norm,
        EbookPhraseAlias.phrase_id,
    ).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id)
    if scope_filter is not None:
        alias_statement = alias_statement.where(scope_filter)

    for row in await session.execute(alias_statement):
        alias_norm = str(row.alias_norm)
        alias_to_ids[alias_norm].append(int(row.phrase_id))
        max_tokens = max(max_tokens, len(alias_norm.split()))

    return PhraseLookup(
        norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()},
        alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()},
        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 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


def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
    """Detect protected phrase windows in a user query using RAM hash lookups.

    Args:
        query_text (str): User query text to scan.
        lookup (PhraseLookup): In-memory phrase and alias lookup maps.

    Returns:
        list[PhraseMatch]: Unhydrated phrase matches found in the query.
    """
    return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup)


def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
    """Detect protected phrase windows in arbitrary text with character offsets.

    Args:
        text (str): Arbitrary text, such as a chunk, to scan.
        lookup (PhraseLookup): In-memory phrase and alias lookup maps.

    Returns:
        list[PhraseMatch]: Unhydrated phrase matches found in the text.
    """
    return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup)


def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]:
    """Detect protected phrase windows from already-normalized tokens.

    Args:
        tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets.
        lookup (PhraseLookup): In-memory phrase and alias lookup maps.

    Returns:
        list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans.
    """
    matches: list[PhraseMatch] = []
    seen: set[tuple[int | None, str, int, int]] = set()
    token_texts = [token.text for token in tokens_]
    for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens):
        phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ())
        alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ())
        for phrase_id in (*phrase_ids, *alias_ids):
            key = (phrase_id, phrase_norm, start, end)
            if key in seen:
                continue
            seen.add(key)
            matches.append(
                PhraseMatch(
                    phrase_norm=phrase_norm,
                    phrase_id=phrase_id,
                    start_token=start,
                    end_token=end,
                    token_count=end - start,
                    start_char=tokens_[start].start_char,
                    end_char=tokens_[end - 1].end_char,
                )
            )
    return matches


async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
    """Fetch protected phrase metadata for raw phrase matches.

    Args:
        session (AsyncSession): Active database session.
        matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.

    Returns:
        list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached.
    """
    if not matches:
        return []

    phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None})
    if not phrase_ids:
        return []

    rows = {
        row.id: row
        for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
    }
    hydrated: list[HydratedPhraseMatch] = []
    for match in matches:
        if match.phrase_id is None:
            continue
        phrase = rows.get(match.phrase_id)
        if phrase is None:
            continue
        hydrated.append(
            HydratedPhraseMatch(
                phrase_id=phrase.id,
                matched_norm=match.phrase_norm,
                phrase_text=phrase.phrase_text,
                phrase_norm=phrase.phrase_norm,
                canonical_id=phrase.canonical_id,
                phrase_type=phrase.phrase_type,
                token_count=match.token_count,
                confidence=phrase.confidence,
                importance=phrase.importance,
                allow_nested=phrase.allow_nested,
                suppress_children=phrase.suppress_children,
                start_token=match.start_token,
                end_token=match.end_token,
                start_char=match.start_char,
                end_char=match.end_char,
                book_id=phrase.book_id,
                series_id=phrase.series_id,
            )
        )
    return hydrated


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

    Args:
        first (HydratedPhraseMatch): First match to compare.
        second (HydratedPhraseMatch): 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 is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool:
    """Return whether one token span is strictly inside another.

    Args:
        child (HydratedPhraseMatch): Candidate nested match.
        parent (HydratedPhraseMatch): 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 rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]:
    """Rank phrase matches by importance, confidence, then token count.

    Args:
        match (HydratedPhraseMatch): 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 should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool:
    """Return whether an already-kept match should suppress a candidate.

    Args:
        candidate (HydratedPhraseMatch): Match being considered for keeping.
        kept (HydratedPhraseMatch): 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 resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]:
    """Resolve overlapping phrase matches without relying only on longest match.

    Args:
        matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap.

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


async def detect_protected_phrases_for_query(
    session: AsyncSession,
    query_text: str,
    config: EbookSearchConfig,
    *,
    lookup: PhraseLookup | None = None,
    book_id: int | None = None,
    series_id: int | None = None,
) -> list[HydratedPhraseMatch]:
    """Run the full online protected-phrase query-detection pipeline.

    Args:
        session (AsyncSession): Active database session.
        query_text (str): User query text to detect phrases in.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
        book_id (int | None): Optional book scope for lookup loading.
        series_id (int | None): Optional series scope for lookup loading.

    Returns:
        list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
    """
    active_lookup = (
        lookup
        if lookup is not None
        else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
    )
    return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))


async def index_chunk_phrase_mentions_for_book(
    session: AsyncSession,
    book_id: int,
    config: EbookSearchConfig,
    *,
    series_id: int | None = None,
    lookup: PhraseLookup | None = None,
) -> 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.
        series_id (int | None): Optional series scope for lookup loading.
        lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.

    Returns:
        int: Total number of chunk phrase mentions indexed for the book.
    """
    active_lookup = (
        lookup
        if lookup is not None
        else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_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 += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
    await session.flush()
    logger.info(f"ebook_chunk_phrase_mentions_indexed {book_id=} {count=}")
    return count


async 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.
    """
    raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
    hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
    for match in hydrated:
        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(hydrated)


async def phrase_hits_for_chunks(
    session: AsyncSession,
    *,
    chunk_ids: Sequence[int],
    phrase_ids: Sequence[int],
) -> dict[int, tuple[ChunkPhraseHit, ...]]:
    """Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.

    Args:
        session (AsyncSession): Active database session.
        chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
        phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.

    Returns:
        dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
    """
    if not chunk_ids or not phrase_ids:
        return {}

    mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
    statement = (
        select(
            EbookChunkPhraseMention.chunk_id,
            EbookProtectedPhrase.id.label("phrase_id"),
            EbookProtectedPhrase.phrase_text,
            mention_count,
        )
        .join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
        .where(
            EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
            EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
        )
        .group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
        .order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
    )
    hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
    for row in await session.execute(statement):
        hits[int(row.chunk_id)].append(
            ChunkPhraseHit(
                phrase_id=int(row.phrase_id),
                phrase_text=str(row.phrase_text),
                mention_count=int(row.mention_count),
            )
        )
    return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}


async def phrase_hit_counts_for_chunks(
    session: AsyncSession,
    *,
    chunk_ids: Sequence[int],
    phrase_ids: Sequence[int],
) -> dict[int, int]:
    """Return phrase-hit counts by chunk id using indexed chunk mentions.

    Args:
        session (AsyncSession): Active database session.
        chunk_ids (Sequence[int]): Chunk ids to count mentions for.
        phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.

    Returns:
        dict[int, int]: Total mention count per chunk id.
    """
    hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
    return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}