packages/core/src/Persistence/PersistPostgres.py

"""PostgreSQL-backed persistence adapter with pgvector + native FTS."""
from __future__ import annotations

import sys
from array import array
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Iterable

import numpy as np
from pgvector.sqlalchemy import Vector
from sqlalchemy import bindparam, create_engine, text
from sqlalchemy.engine import Engine

from Chunker.Chunk import Chunk
from constants import (
    DEFAULT_TABLE_NAME,
    EMBEDDING_DIMENSIONS,
    HYBRID_SEARCH_KEYWORD_WEIGHT,
    HYBRID_SEARCH_VECTOR_WEIGHT,
    POSTGRES_FTS_LANGUAGE,
)
from .Persist import DBConfig, PersistenceAdapter, _split_sql
from .persistence_registry import register_persistence_adapter


class PersistInPostgres(PersistenceAdapter):
    """Persistence adapter backed by PostgreSQL using pgvector and tsvector FTS."""

    def __init__(self, *, cfg: DBConfig, dim: int = EMBEDDING_DIMENSIONS, engine: Optional[Engine] = None,
                 engine_factory: Optional[Callable[[], Engine]] = None, **_: Any, ) -> None:
        if cfg.provider != "postgres":
            raise TypeError("cfg.provider must be 'postgres' for PersistInPostgres")
        self._cfg = cfg
        self._dim = dim
        self._table = cfg.table_map.get("chunks", DEFAULT_TABLE_NAME)
        self._engine = engine or (engine_factory or self._build_engine)()
        self._bootstrap()

    def _build_engine(self) -> Engine:
        connect_args: Dict[str, Any] = {}
        if self._cfg.auth_token:
            connect_args["password"] = self._cfg.auth_token
        return create_engine(self._normalized_url(), future=True, connect_args=connect_args)

    def _normalized_url(self) -> str:
        url = self._cfg.url.strip()
        if url.startswith("postgresql+"):
            return url
        if url.startswith("postgres://"):
            return "postgresql+psycopg2://" + url[len("postgres://"):]
        if url.startswith("postgresql://"):
            return "postgresql+psycopg2://" + url[len("postgresql://"):]
        return url

    def _bootstrap(self) -> None:
        sql_path = Path(__file__).parent.parent.parent.parent.parent / "provisioning" / "postgres" / "schema.sql"
        # Path from packages/core/src/Persistence/ up to repo root, then into provisioning
        # Adjust depth: Persistence -> src -> core -> packages -> repo root = 4 levels up
        sql_template = sql_path.read_text(encoding="utf-8")
        sql = sql_template.replace("{embedding_dim}", str(self._dim))

        with self._engine.begin() as conn:
            # Execute each statement individually (psycopg2 doesn't support multi-statement)
            for statement in _split_sql(sql):
                if statement.strip():
                    conn.execute(text(statement))

    @property
    def _upsert_sql(self) -> str:
        return (
            f"INSERT INTO {self._table} (id, repo, branch, path, language, start_row, start_col, end_row, end_col, "
            "start_bytes, end_bytes, chunk, status, mutation_id, embedding, search_vector) "
            f"VALUES (:id, :repo, :branch, :path, :language, :start_row, :start_col, :end_row, :end_col, :start_bytes, :end_bytes, :chunk, :status, :mutation_id, :embedding, to_tsvector('{POSTGRES_FTS_LANGUAGE}', :chunk)) "
            "ON CONFLICT(id) DO UPDATE SET repo=excluded.repo, branch=excluded.branch, path=excluded.path, language=excluded.language, "
            "start_row=excluded.start_row, start_col=excluded.start_col, end_row=excluded.end_row, end_col=excluded.end_col, "
            "start_bytes=excluded.start_bytes, end_bytes=excluded.end_bytes, chunk=excluded.chunk, status=excluded.status, "
            "mutation_id=excluded.mutation_id, embedding=excluded.embedding, search_vector=excluded.search_vector"
        )

    def persist_batch(self, chunks: Iterable[Chunk]) -> None:
        valid = [chunk for chunk in chunks if chunk is not None]
        if not valid:
            return

        stmt = text(self._upsert_sql).bindparams(bindparam("embedding", type_=Vector(self._dim)))
        with self._engine.begin() as conn:
            for chunk in valid:
                if chunk.embeddings is None:
                    raise ValueError(f"Chunk {chunk.path} missing embeddings")
                start_row, start_col = chunk.start_rc
                end_row, end_col = chunk.end_rc
                conn.execute(
                    stmt,
                    {
                        "id": chunk.id(),
                        "repo": chunk.repo,
                        "branch": chunk.branch,
                        "path": chunk.path,
                        "language": chunk.language,
                        "start_row": start_row,
                        "start_col": start_col,
                        "end_row": end_row,
                        "end_col": end_col,
                        "start_bytes": chunk.start_bytes,
                        "end_bytes": chunk.end_bytes,
                        "chunk": chunk.chunk,
                        "status": "committed",
                        "mutation_id": None,
                        "embedding": self._decode_embedding(chunk.embeddings),
                    },
                )

    def delete_batch(self, paths: Iterable[str], repo: str) -> None:
        if not paths:
            return
        where_clause = "path IN :paths"
        params: dict[str, Any] = {"paths": tuple(paths)}
        if repo is not None:
            where_clause += " AND repo = :repo"
            params["repo"] = repo

        delete_stmt = text(
            f"DELETE FROM {self._table} WHERE {where_clause}"
        ).bindparams(bindparam("paths", expanding=True))
        with self._engine.begin() as conn:
            conn.execute(delete_stmt, params)

    def search(self, query_embedding: Any, query_text: str, limit: int = 10, repo: str | None = None,
               branch: str | None = None, ) -> List[Chunk]:
        normalized_limit = max(1, int(limit))
        query = (query_text or "").strip()

        # We use a CTE approach for true hybrid search: 
        # 1. Fetch top candidates by vector similarity.
        # 2. Fetch top candidates by keyword (FTS).
        # 3. Combine and rerank them using the hybrid scoring formula.

        common_filters = []
        params: Dict[str, Any] = {
            "query_embedding": query_embedding,
            "query_text": query,
            "fts_lang": POSTGRES_FTS_LANGUAGE,
            "candidate_limit": normalized_limit * 5,
            "final_limit": normalized_limit,
            "vector_weight": HYBRID_SEARCH_VECTOR_WEIGHT,
            "keyword_weight": HYBRID_SEARCH_KEYWORD_WEIGHT,
        }

        if repo is not None:
            common_filters.append("repo = :repo")
            params["repo"] = repo
        if branch is not None:
            common_filters.append("branch = :branch")
            params["branch"] = branch

        where_clause = (" WHERE " + " AND ".join(common_filters)) if common_filters else ""

        sql = text(
            f"""
            WITH vector_candidates AS (
                SELECT id, 1 - (embedding <=> :query_embedding) AS v_score
                FROM {self._table}
                {where_clause}
                ORDER BY embedding <=> :query_embedding
                LIMIT :candidate_limit
            ),
            keyword_candidates AS (
                SELECT id, ts_rank_cd(search_vector, websearch_to_tsquery(:fts_lang, :query_text)) AS k_score
                FROM {self._table}
                {where_clause} AND search_vector @@ websearch_to_tsquery(:fts_lang, :query_text)
                ORDER BY k_score DESC
                LIMIT :candidate_limit
            ),
            combined_ids AS (
                SELECT id FROM vector_candidates
                UNION
                SELECT id FROM keyword_candidates
            )
            SELECT 
                c.id, c.path, c.repo, c.branch, c.chunk, c.embedding, c.language,
                c.start_row, c.start_col, c.end_row, c.end_col, c.start_bytes, c.end_bytes,
                COALESCE(vc.v_score, 1 - (c.embedding <=> :query_embedding)) AS vector_rank,
                COALESCE(kc.k_score, ts_rank_cd(c.search_vector, websearch_to_tsquery(:fts_lang, :query_text))) AS keyword_rank
            FROM combined_ids ci
            JOIN {self._table} c ON c.id = ci.id
            LEFT JOIN vector_candidates vc ON vc.id = ci.id
            LEFT JOIN keyword_candidates kc ON kc.id = ci.id
            ORDER BY (
                COALESCE(vc.v_score, 1 - (c.embedding <=> :query_embedding)) * :vector_weight +
                COALESCE(kc.k_score, ts_rank_cd(c.search_vector, websearch_to_tsquery(:fts_lang, :query_text))) * :keyword_weight
            ) DESC
            LIMIT :final_limit
            """
        ).bindparams(bindparam("query_embedding", type_=Vector(self._dim)))

        with self._engine.connect() as conn:
            # Handle empty query text for FTS (it would fail websearch_to_tsquery)
            if not query:
                # Vector-only fallback if no query text
                fallback_sql = text(
                    f"""
                    SELECT id, path, repo, branch, chunk, embedding, language,
                           start_row, start_col, end_row, end_col, start_bytes, end_bytes,
                           1 - (embedding <=> :query_embedding) AS vector_rank,
                           0.0 AS keyword_rank
                    FROM {self._table}
                    {where_clause}
                    ORDER BY embedding <=> :query_embedding
                    LIMIT :final_limit
                    """
                ).bindparams(bindparam("query_embedding", type_=Vector(self._dim)))
                rows = conn.execute(fallback_sql, params).mappings().all()
            else:
                rows = conn.execute(sql, params).mappings().all()

        return [self._row_to_chunk(dict(row)) for row in rows]

    def get_indexed_paths(self, repo: str | None = None) -> set[str]:
        where = ""
        params = {}
        if repo is not None:
            where = " WHERE repo = :repo"
            params["repo"] = repo

        sql = text(f"SELECT DISTINCT path FROM {self._table}{where}")
        with self._engine.connect() as conn:
            rows = conn.execute(sql, params).fetchall()
        return {row[0] for row in rows}

    def _decode_embedding(self, raw: bytes) -> np.ndarray:
        if len(raw) != self._dim * 4:
            raise ValueError(f"Embedding size mismatch: expected {self._dim * 4} bytes, got {len(raw)}")
        return np.frombuffer(raw, dtype=np.float32)

    def _row_to_chunk(self, row: Dict[str, Any]) -> Chunk:
        text_value = row.get("chunk") or ""
        embedding = row.get("embedding")
        embedding_bytes = None
        if embedding is not None:
            vals = array("f", embedding)
            embedding_bytes = vals.tobytes()
        return Chunk(
            chunk=text_value,
            repo=row.get("repo") or "",
            branch=row.get("branch"),
            path=row.get("path") or "",
            language=row.get("language") or "unknown",
            start_rc=(int(row.get("start_row") or 0), int(row.get("start_col") or 0)),
            end_rc=(int(row.get("end_row") or 0), int(row.get("end_col") or 0)),
            start_bytes=int(row.get("start_bytes") or 0),
            end_bytes=int(row.get("end_bytes") or len(text_value.encode("utf-8"))),
            embeddings=embedding_bytes,
        )


def _postgres_factory(
        *,
        cfg: DBConfig,
        dim: int,
        engine: Optional[Engine] = None,
        engine_factory: Optional[Callable[[], Engine]] = None,
        **kwargs: Any,
) -> PersistenceAdapter:
    return PersistInPostgres(
        cfg=cfg,
        dim=dim,
        engine=engine,
        engine_factory=engine_factory,
        **kwargs,
    )


register_persistence_adapter("postgres", _postgres_factory)

# Backward compatibility for legacy absolute import paths used by older tests.
sys.modules.setdefault("PersistPostgres", sys.modules[__name__])

__all__ = ["PersistInPostgres"]