app/vectorstore/faiss_wrapper.py

"""LangChain-backed FAISS vector store (default backend).

Uses ``langchain_community.vectorstores.FAISS`` — open source, runs locally,
no separate vector database process. Embeddings and persistence use the
configured LangChain embedder and ``settings.faiss_index_path``.

Tenant isolation: FAISS has no server-side metadata filters; we
over-fetch and post-filter by ``tenant_id`` (same pattern as the legacy
custom FAISS wrapper).

Concurrent writes are serialized with a lock so parallel ingest workers do
not corrupt the in-memory index or on-disk files.
"""

import logging
import shutil
import tempfile
import threading
import time
from pathlib import Path
from typing import Any, cast

from langchain_community.vectorstores.faiss import FAISS
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

from app.config import settings
from app.models.schemas import SearchResult
from app.vectorstore.base import VectorStore

logger = logging.getLogger(__name__)

_FETCH_MULTIPLIER = 20  # over-fetch factor to compensate for post-tenant filtering
# When a tenant has few chunks in a large shared index, global top-k never includes them.
_TENANT_SCOPED_SEARCH_MAX = 2500


class FAISSVectorStore(VectorStore):
    """VectorStore backed by LangChain FAISS with post-hoc tenant filtering.

    Args:
        embedding: LangChain-compatible embeddings instance.

    Example::

        vs = FAISSVectorStore(embedding=get_embedding_client())
        vs.add_documents(docs)
        results = vs.search("Victorian terrace", tenant_id="t-abc", k=5)
    """

    def __init__(self, embedding: Embeddings) -> None:
        self._embedding = embedding
        self._index_path = Path(settings.faiss_index_path)
        self._index_path.mkdir(parents=True, exist_ok=True)
        self._store: FAISS | None = None
        self._lock = threading.RLock()
        self._index_disk_mtime: float = 0.0
        self._load()

    def _load(self) -> None:
        """Load a persisted LangChain FAISS index from disk if present."""
        index_file = self._index_path / "index.faiss"
        if index_file.exists():
            try:
                self._store = FAISS.load_local(
                    folder_path=str(self._index_path),
                    embeddings=self._embedding,
                    allow_dangerous_deserialization=True,
                )
                self._ensure_index_compatible()
                if self._store is not None:
                    self._index_disk_mtime = index_file.stat().st_mtime
                    logger.info("Loaded LangChain FAISS index from %s", self._index_path)
            except Exception as exc:
                logger.warning("Could not load FAISS index: %s", exc)
                self._quarantine_corrupt_files(str(exc))
                self._store = None
        else:
            logger.info("No FAISS index found at %s — will create on first add", self._index_path)

    def _quarantine_corrupt_files(self, reason: str) -> None:
        """Move unreadable index files aside so a fresh index can be built."""
        stamp = time.strftime("%Y%m%d-%H%M%S")
        for name in ("index.faiss", "index.pkl"):
            path = self._index_path / name
            if not path.exists():
                continue
            dest = self._index_path / f"{name}.corrupt.{stamp}"
            try:
                path.rename(dest)
                logger.warning("Quarantined corrupt FAISS file %s → %s (%s)", path, dest, reason)
            except OSError as move_exc:
                logger.warning("Could not quarantine %s: %s", path, move_exc)

    def _save(self) -> None:
        """Persist the current FAISS index atomically (temp dir → rename)."""
        if self._store is None:
            return
        self._index_path.mkdir(parents=True, exist_ok=True)
        tmp_dir = Path(tempfile.mkdtemp(prefix="faiss_save_", dir=self._index_path.parent))
        try:
            self._store.save_local(str(tmp_dir))
            for name in ("index.faiss", "index.pkl"):
                src = tmp_dir / name
                if not src.is_file():
                    continue
                dst = self._index_path / name
                tmp_dst = dst.with_suffix(dst.suffix + ".tmp")
                shutil.copy2(src, tmp_dst)
                tmp_dst.replace(dst)
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)
        index_file = self._index_path / "index.faiss"
        if index_file.is_file():
            self._index_disk_mtime = index_file.stat().st_mtime

    def _reload_if_disk_newer(self) -> None:
        """Reload index when another process has written newer files to disk."""
        index_file = self._index_path / "index.faiss"
        if not index_file.is_file():
            return
        disk_mtime = index_file.stat().st_mtime
        if disk_mtime <= self._index_disk_mtime:
            return
        with self._lock:
            if not index_file.is_file():
                return
            disk_mtime = index_file.stat().st_mtime
            if disk_mtime <= self._index_disk_mtime:
                return
            logger.info("FAISS index on disk is newer than in-memory copy — reloading")
            self._load()

    def total_vectors(self) -> int:
        """Return FAISS ``ntotal`` (0 when no index is loaded)."""
        with self._lock:
            if self._store is None:
                return 0
            return int(self._store.index.ntotal)

    @property
    def index_loaded(self) -> bool:
        """True when an on-disk index was loaded or documents have been added."""
        return self._store is not None

    def _embedding_dim(self) -> int:
        return len(self._embedding.embed_query("dimension probe"))

    def _index_dim(self) -> int | None:
        if self._store is None:
            return None
        return int(self._store.index.d)

    def _reset_persisted_index(self, reason: str) -> None:
        """Drop in-memory and on-disk index (e.g. embedding model / dimension change)."""
        logger.warning(
            "Resetting FAISS index at %s: %s. Re-upload tenant documents to rebuild the index.",
            self._index_path,
            reason,
        )
        self._store = None
        for name in ("index.faiss", "index.pkl"):
            path = self._index_path / name
            if path.exists():
                path.unlink()

    def _ensure_index_compatible(self) -> None:
        """Invalidate a persisted index built with a different embedding dimension."""
        if self._store is None:
            return
        idx_d = self._index_dim()
        emb_d = self._embedding_dim()
        if idx_d is not None and idx_d != emb_d:
            self._reset_persisted_index(
                f"index dimension {idx_d} != current embeddings {emb_d}",
            )

    def add_documents(self, documents: list[Document], *, persist: bool = True) -> None:
        """Embed and insert ``documents`` into the FAISS index.

        Args:
            documents: LangChain Documents with full metadata.
            persist: When false, skip the atomic disk write (call :meth:`flush` later).
        """
        if not documents:
            return

        self._reload_if_disk_newer()
        with self._lock:
            self._ensure_index_compatible()
            try:
                if self._store is None:
                    self._store = FAISS.from_documents(documents, self._embedding)
                else:
                    self._store.add_documents(documents)
            except AssertionError as exc:
                idx_d = self._index_dim()
                emb_d = self._embedding_dim()
                raise RuntimeError(
                    f"FAISS add failed (index dim={idx_d}, embeddings dim={emb_d}): {exc!r}"
                ) from exc

            if persist:
                self._save()
        logger.debug("Added %d documents to LangChain FAISS (persist=%s)", len(documents), persist)

    def flush(self) -> None:
        """Persist the in-memory index to disk (no-op when empty)."""
        with self._lock:
            self._save()

    def search(
        self,
        query: str,
        tenant_id: str,
        k: int = 10,
        *,
        hierarchy_level: str | None = None,
        doc_id_in: frozenset[str] | None = None,
        purpose_in: frozenset[str] | None = None,
        exclude_purpose: frozenset[str] | None = None,
    ) -> list[SearchResult]:
        """Search FAISS and post-filter by ``tenant_id`` (and optional metadata filters).

        Over-fetches by ``_FETCH_MULTIPLIER`` to ensure enough tenant-matching
        results are available after filtering.

        Args:
            query: Plain-text search query.
            tenant_id: Only return chunks belonging to this tenant.
            k: Maximum number of results after filtering.
            hierarchy_level: Restrict to this ``hierarchy_level`` when set.
            doc_id_in: Restrict to these ``doc_id`` values when set.
            purpose_in: If set, only return chunks whose ``document_purpose``
                is in this set (e.g. ``{"style_corpus"}`` for style-only
                retrieval). Missing metadata is treated as ``report_source``.
            exclude_purpose: If set, drop chunks whose ``document_purpose`` is
                in this set (e.g. ``{"style_corpus"}`` to keep facts from past
                reports out of new reports' factual retrieval).

        Returns:
            Filtered list of :class:`~app.models.schemas.SearchResult`.
        """
        try:
            self._reload_if_disk_newer()
            if self._store is None:
                return []
            # Embed outside the index lock so concurrent reads do not serialize on embedding.
            query_vector = self._embedding.embed_query(query)
            with self._lock:
                if self._store is None:
                    return []
                tenant_n = self.count(tenant_id)
                fetch_k = max(k * _FETCH_MULTIPLIER, k * 4)
                if tenant_n > 0 and tenant_n <= _TENANT_SCOPED_SEARCH_MAX:
                    pairs = self._similarity_pairs_for_tenant(
                        query_vector,
                        tenant_id,
                        fetch_k=fetch_k,
                    )
                else:
                    if tenant_n > 0:
                        fetch_k = min(
                            self.total_vectors(),
                            max(fetch_k, tenant_n * 30),
                        )
                    pairs = self._store.similarity_search_with_score_by_vector(
                        query_vector,
                        k=fetch_k,
                    )
        except Exception as exc:
            logger.warning("FAISS search failed: %s", exc)
            return []

        results: list[SearchResult] = []
        for doc, distance in pairs:
            meta = doc.metadata
            if meta.get("tenant_id") != tenant_id:
                continue
            hl = str(meta.get("hierarchy_level") or "paragraph")
            if hierarchy_level is not None and hl != hierarchy_level:
                continue
            did = str(meta.get("doc_id", ""))
            if doc_id_in is not None and did not in doc_id_in:
                # Always allow knowledge-base (KB) documents through — these are
                # the firm's approved templates / training uploads that should be
                # retrievable regardless of which report is being generated.
                if not meta.get("kb"):
                    continue
            # ``document_purpose`` discriminator: ``style_corpus`` chunks
            # (the user's past reports) must never appear as factual
            # evidence for a different property. KB chunks are exempt from
            # the purpose filter since they are template/boilerplate, not
            # tenant-private past reports.
            doc_purpose = str(meta.get("document_purpose") or "report_source")
            if not meta.get("kb"):
                if purpose_in is not None and doc_purpose not in purpose_in:
                    continue
                if exclude_purpose is not None and doc_purpose in exclude_purpose:
                    continue
            d = float(distance)
            sim = 1.0 / (1.0 + d)
            st = meta.get("section_title")
            sid = meta.get("section_id")
            pi = meta.get("paragraph_index")
            src = meta.get("source")
            kb = meta.get("kb")
            kb_path = meta.get("kb_path")
            chunk_role = meta.get("chunk_role")
            pidx: int | None
            if isinstance(pi, int):
                pidx = pi
            else:
                try:
                    pidx = int(pi) if pi is not None else None
                except (TypeError, ValueError):
                    pidx = None
            results.append(
                SearchResult(
                    chunk_id=str(meta.get("chunk_id", "")),
                    doc_id=did,
                    tenant_id=str(meta.get("tenant_id", "")),
                    text=doc.page_content,
                    score=sim,
                    section_type=str(meta.get("section_type", "paragraph")),
                    hierarchy_level=hl,
                    section_title=str(st) if st is not None else None,
                    section_id=str(sid) if sid is not None else None,
                    paragraph_index=pidx,
                    parent_chunk_id=str(meta.get("parent_chunk_id"))
                    if meta.get("parent_chunk_id")
                    else None,
                    source=str(src) if src is not None else None,
                    kb=bool(kb) if kb is not None else None,
                    kb_path=str(kb_path) if kb_path is not None else None,
                    document_purpose=doc_purpose,
                    chunk_role=str(chunk_role) if chunk_role is not None else None,
                )
            )
            if len(results) >= k:
                break

        return results

    def _similarity_pairs_for_tenant(
        self,
        query_vector: list[float],
        tenant_id: str,
        *,
        fetch_k: int,
    ) -> list[tuple[Document, float]]:
        """Score only this tenant's vectors (avoids KB/global chunks crowding top-k)."""
        import numpy as np

        store = self._store
        if store is None:
            return []

        q = np.asarray(query_vector, dtype=np.float32)
        id_to_idx = {doc_id: int(idx) for idx, doc_id in store.index_to_docstore_id.items()}
        scored: list[tuple[Document, float]] = []
        for docstore_id, doc in store.docstore._dict.items():
            if doc.metadata.get("tenant_id") != tenant_id:
                continue
            faiss_idx = id_to_idx.get(docstore_id)
            if faiss_idx is None:
                continue
            vec = np.asarray(store.index.reconstruct(faiss_idx), dtype=np.float32)
            scored.append((doc, float(np.linalg.norm(q - vec))))

        scored.sort(key=lambda item: item[1])
        return scored[:fetch_k]

    async def search_async(
        self,
        query: str,
        tenant_id: str,
        k: int = 10,
        *,
        hierarchy_level: str | None = None,
        doc_id_in: frozenset[str] | None = None,
        purpose_in: frozenset[str] | None = None,
        exclude_purpose: frozenset[str] | None = None,
    ) -> list[SearchResult]:
        """Non-blocking search via thread pool (embedding no longer under global lock)."""
        from functools import partial

        from app.async_executor import run_sync_in_executor

        return await run_sync_in_executor(
            partial(
                self.search,
                query,
                tenant_id,
                k=k,
                hierarchy_level=hierarchy_level,
                doc_id_in=doc_id_in,
                purpose_in=purpose_in,
                exclude_purpose=exclude_purpose,
            ),
        )

    def delete_document(self, doc_id: str) -> None:
        """Remove all chunks for ``doc_id``.

        LangChain FAISS supports deletion by document ID if the index was
        built with ``docstore`` support (the default).

        Args:
            doc_id: Document identifier to remove.
        """
        with self._lock:
            if self._store is None:
                return

            ds = cast(Any, self._store.docstore)._dict
            ids_to_delete = [
                doc_id_key for doc_id_key, doc in ds.items() if doc.metadata.get("doc_id") == doc_id
            ]
            if ids_to_delete:
                self._store.delete(ids_to_delete)
                self._save()
                logger.info("Deleted %d chunks for doc_id=%s from FAISS", len(ids_to_delete), doc_id)

    def count(self, tenant_id: str) -> int:
        """Count live chunks for ``tenant_id``.

        Args:
            tenant_id: Tenant identifier.

        Returns:
            Integer chunk count.
        """
        self._reload_if_disk_newer()
        with self._lock:
            if self._store is None:
                return 0
            ds = cast(Any, self._store.docstore)._dict
            return len([d for d in ds.values() if d.metadata.get("tenant_id") == tenant_id])

    def count_for_doc(self, doc_id: str) -> int:
        """Count live chunks for a specific document.

        Args:
            doc_id: Document identifier.

        Returns:
            Integer chunk count for that document.
        """
        self._reload_if_disk_newer()
        with self._lock:
            if self._store is None:
                return 0
            ds = cast(Any, self._store.docstore)._dict
            return len([d for d in ds.values() if d.metadata.get("doc_id") == doc_id])

    def indexed_doc_ids(self) -> set[str]:
        """Set of ``doc_id`` values present in the index (one scan)."""
        self._reload_if_disk_newer()
        with self._lock:
            if self._store is None:
                return set()
            ds = cast(Any, self._store.docstore)._dict
            out: set[str] = set()
            for doc in ds.values():
                did = doc.metadata.get("doc_id")
                if did:
                    out.add(str(did))
            return out