webmail-client/backend/app/search/pg.py

"""Postgres full-text SearchBackend (SPEC §5.2, §8; M0-M1).

Indexing populates ``messages.body_indexed`` (a ``tsvector``) from the doc's
searchable fields. Querying translates the parsed §8.2 groups (AND of ORs) into
SQL over messages joined to mailboxes/accounts, always filtered by ``user_id``
(invariant I-7). Relevance ranking / typo tolerance is Meilisearch's job (M2);
here results are date-ordered with keyset pagination.
"""

from __future__ import annotations

import base64
import datetime as dt
import html
import re
from uuid import UUID

from sqlalchemy import (
    Text,
    any_,
    case,
    cast,
    exists,
    func,
    literal,
    not_,
    or_,
    select,
    tuple_,
    update,
)
from sqlalchemy.sql.elements import ColumnElement

from app.db import get_sessionmaker
from app.models import Attachment, Mailbox, Message
from app.search.base import SearchBackend, SearchDoc, SearchHit, SearchQuery

_TS_CONFIG = "simple"  # language-neutral; avoids wrong-language stemming (M0)


def _prefix_tsquery(value: str) -> ColumnElement[object]:
    """A prefix ``tsquery`` for search-as-you-type: ``invo`` finds ``invoice``.

    We can't tokenize in Python and match the index: Postgres' parser keeps an
    email/host/URL as a *single* lexeme (``lukas.remis@gmail.com``), so splitting
    on word chars would never match it. Instead we re-parse the term with the
    same ``plainto_tsquery`` used at index time, then append ``:*`` to each
    lexeme to make it a prefix. ``plainto_tsquery`` strips tsquery operators and
    quotes, so this is injection- and error-safe on hostile input; an all-symbol
    term yields the empty query (matches nothing).
    """
    lexemes = cast(func.plainto_tsquery(_TS_CONFIG, value), Text)
    prefixed = func.regexp_replace(lexemes, r"'([^']+)'", r"'\1':*", "g")
    return func.to_tsquery(_TS_CONFIG, prefixed)
_RELATIVE_MULT = {"d": 1, "w": 7, "m": 30, "y": 365}


def _email_parts(email: str) -> str:
    """Split an address on @/./-/+/_ so its domain and local part are searchable
    on their own — a bare "wfdif" should find dyrektor@wfdif.com.pl. Postgres
    indexes the whole address as a single lexeme, so the parts are invisible to a
    word search otherwise."""
    return re.sub(r"[@._+\-]+", " ", email)


def _doc_text_blob(doc: SearchDoc) -> str:
    parts = [doc.subject, doc.body_text]
    for field in (doc.from_, doc.to, doc.cc):
        for addr in field:
            parts.append(addr.get("name", ""))
            email = addr.get("email", "")
            if email:
                parts.append(email)  # whole address (one lexeme, for exact/prefix)
                parts.append(_email_parts(email))  # + its parts as separate tokens
    parts.extend(doc.attachment_names)
    return " ".join(p for p in parts if p)


def _parse_date(value: str) -> dt.datetime | None:
    try:
        return dt.datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=dt.UTC)
    except ValueError:
        return None


def _parse_relative(value: str) -> dt.timedelta | None:
    m = re.fullmatch(r"(\d+)([dwmy])", value.strip().lower())
    if not m:
        return None
    n = int(m.group(1)) * _RELATIVE_MULT[m.group(2)]
    return dt.timedelta(days=n)


def _parse_size(value: str) -> int | None:
    m = re.fullmatch(r"(\d+)([km])", value.strip().lower())
    if not m:
        return None
    return int(m.group(1)) * (1024 if m.group(2) == "k" else 1024 * 1024)


def _flag_condition(flag: str) -> ColumnElement[bool]:
    # flag = ANY(messages.flags) — the array membership operator, not func any().
    return literal(flag) == any_(Message.flags)


class PgSearchBackend(SearchBackend):
    async def index(self, doc: SearchDoc) -> None:
        blob = _doc_text_blob(doc)
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message)
                .where(Message.id == doc.message_id, Message.user_id == doc.user_id)
                .values(body_indexed=func.to_tsvector(_TS_CONFIG, blob))
            )
            await session.commit()

    async def delete(self, message_id: UUID) -> None:
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message).where(Message.id == message_id).values(body_indexed=None)
            )
            await session.commit()

    async def delete_user(self, user_id: UUID) -> None:
        # User deletion cascades message rows (FK ON DELETE CASCADE); clearing the
        # tsvector here keeps the backend contract satisfied for partial re-index.
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message).where(Message.user_id == user_id).values(body_indexed=None)
            )
            await session.commit()

    # --- querying ------------------------------------------------------------

    def _clause_condition(self, clause: dict[str, object]) -> ColumnElement[bool] | None:
        value = str(clause.get("value", ""))
        if clause.get("type") == "text":
            if clause.get("phrase"):
                # Quoted phrase: keep exact-adjacency semantics, no prefix.
                return Message.body_indexed.op("@@")(func.websearch_to_tsquery(_TS_CONFIG, value))
            like = f"%{value}%"
            addr_match = or_(
                cast(Message.from_addr, Text).ilike(like),
                cast(Message.to_addr, Text).ilike(like),
                cast(Message.cc_addr, Text).ilike(like),
            )
            if "@" in value:
                # A term with '@' is an address search: match the address substring
                # only (like from:/to:/cc:), so it narrows predictably as you type —
                # not the split-word body matching that made a partial term over-hit.
                return addr_match
            fts = Message.body_indexed.op("@@")(_prefix_tsquery(value))
            if "." in value:
                # A dotted term (e.g. "wfdif.com") stays one host lexeme in FTS, so
                # also match it as an address substring. Bare words don't need this:
                # address parts are indexed as their own tokens (see _doc_text_blob).
                return or_(fts, addr_match)
            return fts
        key = clause.get("key")
        cond = self._filter_condition(str(key), value)
        return cond

    def _filter_condition(self, key: str, value: str) -> ColumnElement[bool] | None:
        like = f"%{value}%"
        if key == "from":
            return cast(Message.from_addr, Text).ilike(like)
        if key == "to":
            return cast(Message.to_addr, Text).ilike(like)
        if key == "cc":
            return cast(Message.cc_addr, Text).ilike(like)
        if key == "subject":
            return Message.subject.ilike(like)
        if key == "in":
            return or_(Mailbox.special_use == value.lower(), Mailbox.path.ilike(like))
        if key == "is":
            v = value.lower()
            if v == "unread":
                return not_(_flag_condition("\\Seen"))
            if v == "read":
                return _flag_condition("\\Seen")
            if v == "starred":
                return _flag_condition("\\Flagged")
            if v == "answered":
                return _flag_condition("\\Answered")
            return None
        if key == "has":
            return Message.has_attachments.is_(True) if value.lower() == "attachment" else None
        if key == "filename":
            return exists().where(
                Attachment.message_id == Message.id, Attachment.filename.ilike(like)
            )
        if key == "before":
            d = _parse_date(value)
            return Message.internal_date < d if d else None
        if key == "after":
            d = _parse_date(value)
            return Message.internal_date >= d if d else None
        if key == "newer_than":
            delta = _parse_relative(value)
            return Message.internal_date >= func.now() - delta if delta else None
        if key == "older_than":
            delta = _parse_relative(value)
            return Message.internal_date < func.now() - delta if delta else None
        if key == "larger":
            n = _parse_size(value)
            return Message.size_bytes > n if n is not None else None
        if key == "smaller":
            n = _parse_size(value)
            return Message.size_bytes < n if n is not None else None
        return None

    def _build_conditions(self, groups: list[list[dict[str, object]]]) -> list[ColumnElement[bool]]:
        conditions: list[ColumnElement[bool]] = []
        for group in groups:
            ors: list[ColumnElement[bool]] = []
            for clause in group:
                cond = self._clause_condition(clause)
                if cond is None:
                    continue
                if clause.get("negate"):
                    cond = not_(cond)
                ors.append(cond)
            if ors:
                conditions.append(or_(*ors))
        return conditions

    async def query(
        self, q: SearchQuery, *, limit: int, cursor: str | None
    ) -> tuple[list[SearchHit], str | None]:
        groups = q.filters.get("groups", []) if isinstance(q.filters, dict) else []
        assert isinstance(groups, list)
        conditions = self._build_conditions(groups)
        terms = _positive_text_terms(groups)

        # Rank a match in a sender/recipient name or address above a body/subject-only
        # match, so "who" results come first; ties break by recency.
        addr_hit = _address_hit_expr(terms)
        rank = (
            func.ts_rank_cd(Message.body_indexed, _prefix_tsquery(q.text))
            if q.text
            else literal(0.0)
        )
        stmt = (
            select(
                Message.id,
                Message.thread_id,
                Message.subject,
                Message.snippet,
                Message.internal_date,
                rank.label("score"),
                addr_hit.label("addr_hit"),
            )
            .select_from(Message)
            .join(Mailbox, Mailbox.id == Message.mailbox_id)
            .where(Message.user_id == q.user_id, Message.thread_id.is_not(None))
        )
        for cond in conditions:
            stmt = stmt.where(cond)

        # Exclude Trash/Junk from search by default (Gmail-style), so deleting a
        # message from search results removes it. An explicit `in:trash`/`in:junk`
        # overrides this to search there on purpose.
        if not _targets_removal_folder(groups):
            stmt = stmt.where(
                or_(
                    Mailbox.special_use.is_(None),
                    Mailbox.special_use.notin_(_REMOVAL_FOLDERS),
                )
            )

        decoded = _decode_cursor(cursor)
        if decoded is not None:
            cur_hit, cur_date, cur_id = decoded
            # Keyset over the full ordering (addr_hit, date, id), all descending.
            stmt = stmt.where(
                tuple_(addr_hit, Message.internal_date, Message.id)
                < tuple_(literal(cur_hit), literal(cur_date), literal(cur_id))
            )
        stmt = stmt.order_by(
            addr_hit.desc(), Message.internal_date.desc(), Message.id.desc()
        ).limit(limit + 1)

        async with get_sessionmaker()() as session:
            rows = (await session.execute(stmt)).all()

        has_more = len(rows) > limit
        rows = rows[:limit]
        hits = [
            SearchHit(
                message_id=row.id,
                thread_id=row.thread_id,
                score=float(row.score or 0.0),
                highlights=_highlights(row.subject, row.snippet, terms),
            )
            for row in rows
        ]
        next_cursor = (
            _encode_cursor(int(rows[-1].addr_hit), rows[-1].internal_date, rows[-1].id)
            if has_more and rows
            else None
        )
        return hits, next_cursor


_REMOVAL_FOLDERS = ("trash", "junk")


def _targets_removal_folder(groups: list[list[dict[str, object]]]) -> bool:
    """True if the query explicitly scopes to Trash/Junk via an ``in:`` filter, in
    which case the default removal-folder exclusion must not apply."""
    for group in groups:
        for clause in group:
            if (
                clause.get("type") == "filter"
                and clause.get("key") == "in"
                and str(clause.get("value", "")).strip().lower() in _REMOVAL_FOLDERS
            ):
                return True
    return False


def _positive_text_terms(groups: list[list[dict[str, object]]]) -> list[str]:
    terms: list[str] = []
    for group in groups:
        for clause in group:
            if clause.get("type") == "text" and not clause.get("negate"):
                terms.append(str(clause.get("value", "")))
    return terms


def _address_hit_expr(terms: list[str]) -> ColumnElement[int]:
    """1 if any positive free-text term appears in a from/to/cc name or address,
    else 0. Used to sort address/name matches ahead of body/subject-only ones."""
    conds: list[ColumnElement[bool]] = []
    for t in terms:
        if t:
            like = f"%{t}%"
            conds.append(cast(Message.from_addr, Text).ilike(like))
            conds.append(cast(Message.to_addr, Text).ilike(like))
            conds.append(cast(Message.cc_addr, Text).ilike(like))
    if not conds:
        return literal(0)
    return case((or_(*conds), 1), else_=0)


def _highlights(subject: str, snippet: str, terms: list[str]) -> dict[str, str]:
    out: dict[str, str] = {}
    if not terms:
        return out
    marked_subject = _mark(subject, terms)
    if "<mark>" in marked_subject:
        out["subject"] = marked_subject
    marked_body = _mark(snippet, terms)
    if "<mark>" in marked_body:
        out["body"] = marked_body
    return out


def _mark(value: str, terms: list[str]) -> str:
    """HTML-escape ``value`` and wrap term matches in <mark> (SPEC §5.2 highlights)."""
    escaped = html.escape(value)
    for term in terms:
        if not term:
            continue
        pattern = re.compile(re.escape(html.escape(term)), re.IGNORECASE)
        escaped = pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", escaped)
    return escaped


def _encode_cursor(addr_hit: int, date: dt.datetime, message_id: UUID) -> str:
    raw = f"{addr_hit}|{date.isoformat()}|{message_id}"
    return base64.urlsafe_b64encode(raw.encode()).decode()


def _decode_cursor(cursor: str | None) -> tuple[int, dt.datetime, UUID] | None:
    if not cursor:
        return None
    try:
        raw = base64.urlsafe_b64decode(cursor.encode()).decode()
        hit_str, date_str, id_str = raw.split("|", 2)
        return int(hit_str), dt.datetime.fromisoformat(date_str), UUID(id_str)
    except (ValueError, TypeError):
        return None