backend/core/grants/postgres_fts.py

"""PostgreSQL full-text search (to_tsvector) + SQLite FTS5 — bez python token overlap."""
from __future__ import annotations

import logging
import re
from typing import Any, Dict, List, Optional, Set

from sqlalchemy import inspect, text
from sqlalchemy.orm import Session

from core.grants.catalog_service import _search_hay
from core.grants.completeness import grant_dict_from_row
from core.grants.models import Grant
from core.subscription.db import engine

logger = logging.getLogger(__name__)

_FTS5_READY = False


def get_db_dialect() -> str:
    return engine.dialect.name


def get_fts_backend() -> str:
    dialect = get_db_dialect()
    if dialect == "postgresql":
        return "postgresql_tsvector"
    if dialect == "sqlite":
        return "sqlite_fts5"
    return "python_fallback"


def is_real_fts_available() -> bool:
    return get_fts_backend() in ("postgresql_tsvector", "sqlite_fts5")


def build_search_document(item: Dict[str, Any]) -> str:
    return _search_hay(item)


def _query_tokens(query: str) -> List[str]:
    return [t for t in re.split(r"\W+", (query or "").lower()) if len(t) > 2]


def _fts5_query(query: str) -> str:
    tokens = _query_tokens(query)
    if not tokens:
        return ""
    return " OR ".join(f'"{t}"' for t in tokens[:12])


def _postgres_tsquery(query: str) -> str:
    """OR-tokeny jak FTS5 — plainto_tsquery wymaga AND i psuje multi-word PL."""
    tokens = _query_tokens(query)[:12]
    if not tokens:
        return ""
    parts: List[str] = []
    for token in tokens:
        safe = re.sub(r"[':\\&|!()]", " ", token).strip()
        if safe:
            parts.append(safe)
    return " | ".join(parts) if parts else ""


def ensure_fts_schema(db: Session) -> Dict[str, Any]:
    """Tworzy FTS5 (sqlite) lub kolumnę search_document (postgres)."""
    global _FTS5_READY
    backend = get_fts_backend()
    stats: Dict[str, Any] = {"backend": backend, "created": False}

    if backend == "sqlite_fts5":
        exists = db.execute(
            text("SELECT name FROM sqlite_master WHERE type='table' AND name='grants_fts'")
        ).fetchone()
        if not exists:
            db.execute(
                text(
                    "CREATE VIRTUAL TABLE grants_fts USING fts5("
                    "grant_id UNINDEXED, document, tokenize='unicode61')"
                )
            )
            db.commit()
            stats["created"] = True
        _FTS5_READY = True
        return stats

    if backend == "postgresql_tsvector":
        inspector = inspect(engine)
        cols = {c["name"] for c in inspector.get_columns("grants")}
        if "search_document" not in cols:
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE grants ADD COLUMN search_document TEXT"))
            stats["created"] = True
        return stats

    return stats


def sync_fts_index(db: Session, *, limit: int = 5000) -> Dict[str, Any]:
    """Indeksuje dokumenty w FTS5 lub kolumnie search_document."""
    ensure_fts_schema(db)
    backend = get_fts_backend()
    rows = db.query(Grant).limit(limit).all()
    indexed = 0

    if backend == "sqlite_fts5":
        db.execute(text("DELETE FROM grants_fts"))
        for row in rows:
            doc = build_search_document(grant_dict_from_row(row))
            if not doc.strip():
                continue
            db.execute(
                text("INSERT INTO grants_fts(grant_id, document) VALUES (:gid, :doc)"),
                {"gid": row.source_id, "doc": doc},
            )
            indexed += 1
        db.commit()
        return {"backend": backend, "indexed": indexed}

    if backend == "postgresql_tsvector":
        for row in rows:
            doc = build_search_document(grant_dict_from_row(row))
            row.search_document = doc
            if doc.strip():
                indexed += 1
        db.commit()
        return {"backend": backend, "indexed": indexed}

    return {"backend": backend, "indexed": 0, "reason": "no_fts_backend"}


def postgres_fts_scores(
    db: Session,
    query: str,
    candidate_ids: Optional[List[str]] = None,
    *,
    limit: int = 50,
) -> Dict[str, float]:
    """
    Zwraca source_id -> score 0..1 z prawdziwego FTS (tsvector lub FTS5).
    candidate_ids ogranicza wyniki do przefiltrowanego zbioru.
    """
    if not query.strip() or not is_real_fts_available():
        return {}

    ensure_fts_schema(db)
    backend = get_fts_backend()
    allowed: Optional[Set[str]] = set(candidate_ids) if candidate_ids else None
    scores: Dict[str, float] = {}

    if backend == "sqlite_fts5":
        fts_q = _fts5_query(query)
        if not fts_q:
            return {}
        try:
            rows = db.execute(
                text(
                    "SELECT grant_id, bm25(grants_fts) AS rank "
                    "FROM grants_fts WHERE grants_fts MATCH :q "
                    "ORDER BY rank LIMIT :lim"
                ),
                {"q": fts_q, "lim": max(limit * 4, 80)},
            ).fetchall()
        except Exception as e:
            logger.warning("[PostgresFTS] FTS5 query failed: %s", e)
            return {}

        if not rows:
            return {}

        raw_ranks = [float(r[1]) for r in rows]
        min_r, max_r = min(raw_ranks), max(raw_ranks)
        span = max(max_r - min_r, 1e-6)
        for rank, (gid, bm) in enumerate(rows, start=1):
            if allowed is not None and gid not in allowed:
                continue
            norm = (float(bm) - min_r) / span
            scores[str(gid)] = max(scores.get(str(gid), 0), min(1.0, norm), 1.0 / rank)
            if len(scores) >= limit:
                break
        return scores

    if backend == "postgresql_tsvector":
        ts_q = _postgres_tsquery(query)
        if not ts_q:
            return {}
        try:
            rows = db.execute(
                text(
                    "SELECT source_id, ts_rank("
                    "  to_tsvector('simple', coalesce(search_document, '')),"
                    "  to_tsquery('simple', :q)"
                    ") AS rank "
                    "FROM grants "
                    "WHERE search_document IS NOT NULL AND search_document != '' "
                    "  AND to_tsvector('simple', search_document) @@ to_tsquery('simple', :q) "
                    "ORDER BY rank DESC LIMIT :lim"
                ),
                {"q": ts_q, "lim": max(limit * 4, 80)},
            ).fetchall()
        except Exception as e:
            logger.warning("[PostgresFTS] tsvector query failed: %s", e)
            return {}

        if not rows:
            return {}
        max_rank = max(float(r[1]) for r in rows) or 1.0
        for rank, (gid, ts_rank) in enumerate(rows, start=1):
            if allowed is not None and gid not in allowed:
                continue
            scores[str(gid)] = max(
                scores.get(str(gid), 0),
                min(1.0, float(ts_rank) / max_rank),
                1.0 / rank,
            )
            if len(scores) >= limit:
                break
        return scores

    return {}