backend/core/rag_store.py

"""Per-tenant, two-tier FAISS vector store.

Tiers:

* ``TIER_MASTER`` — operator master-template paragraphs. **Not** scrubbed at
  ingest (the master is property-agnostic boilerplate); it carries the firm's
  approved wording and is the structural authority.
* ``TIER_REFERENCE`` — past completed reports. **Always** scrubbed at ingest;
  used only for style/terminology. A read-time scrub backstop guards against any
  chunk that slipped through unscrubbed.

Each (tenant, tier) pair is an independent ``IndexFlatIP`` over L2-normalised
embeddings (so inner product == cosine similarity), persisted to disk alongside a
JSON metadata sidecar.
"""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

import numpy as np

from backend.config import settings
from backend.core import pii_scrubber
from backend.core.embeddings import Embedder, get_embedder
from backend.core.reference_filter import build_reference_allowlist, meta_matches_allowlist
from backend.utils import tenant_store

logger = logging.getLogger(__name__)

TIER_MASTER = "master"
TIER_REFERENCE = "reference"


@dataclass
class Chunk:
    text: str
    section_id: str = ""
    doc_id: str = ""
    tier: str = TIER_MASTER
    is_scrubbed: bool = False
    chunk_id: str = ""
    source_filename: str = ""
    paragraph_index: int = 0


@dataclass
class SearchHit:
    text: str
    section_id: str
    doc_id: str
    tier: str
    score: float
    is_scrubbed: bool
    source_filename: str = ""
    paragraph_index: int = 0
    chunk_id: str = ""


def _search_hit_from_meta(
    meta: dict,
    *,
    tier: str,
    rank: float,
    text: str,
) -> SearchHit:
    return SearchHit(
        text=text,
        section_id=meta.get("section_id", ""),
        doc_id=meta.get("doc_id", ""),
        tier=tier,
        score=rank,
        is_scrubbed=bool(meta.get("is_scrubbed")),
        source_filename=meta.get("source_filename", ""),
        paragraph_index=int(meta.get("paragraph_index") or 0),
        chunk_id=meta.get("chunk_id", ""),
    )


@dataclass
class _TierIndex:
    index: object  # faiss.IndexFlatIP
    meta: list[dict] = field(default_factory=list)


class RagStore:
    """Lazy, thread-safe per-tenant/tier FAISS store."""

    def __init__(self, embedder: Embedder | None = None) -> None:
        self._embedder = embedder or get_embedder()
        self._cache: dict[tuple[str, str], _TierIndex] = {}
        self._lock = threading.RLock()

    # ── persistence ──────────────────────────────────────────────────────────
    def _paths(self, tenant_id: str, tier: str) -> tuple[Path, Path]:
        d = tenant_store.faiss_dir(tenant_id, tier)
        return d / "index.faiss", d / "meta.json"

    def _new_index(self):
        import faiss

        return faiss.IndexFlatIP(self._embedder.embed_dim)

    def _load(self, tenant_id: str, tier: str) -> _TierIndex:
        import faiss

        idx_path, meta_path = self._paths(tenant_id, tier)
        if idx_path.is_file() and meta_path.is_file():
            try:
                # io_flags=0: load into RAM (avoid mmap — Windows cannot replace mmap'd files).
                index = faiss.read_index(str(idx_path), 0)
                meta = json.loads(meta_path.read_text(encoding="utf-8"))
                return _TierIndex(index=index, meta=meta)
            except Exception as exc:  # noqa: BLE001
                logger.warning("Failed to load FAISS %s/%s (%s); rebuilding empty.",
                               tenant_id, tier, exc)
        return _TierIndex(index=self._new_index(), meta=[])

    def _persist(self, tenant_id: str, tier: str, ti: _TierIndex) -> None:
        import faiss
        import os
        import tempfile

        from backend.utils.runtime_paths import ensure_data_drive_runtime_dirs

        ensure_data_drive_runtime_dirs()

        idx_path, meta_path = self._paths(tenant_id, tier)
        idx_path = idx_path.resolve()
        meta_path = meta_path.resolve()
        idx_path.parent.mkdir(parents=True, exist_ok=True)

        payload = json.dumps(ti.meta, ensure_ascii=False, separators=(",", ":"))
        scratch_root = settings.data_dir_path / "tmp" / "faiss_persist"
        scratch_root.mkdir(parents=True, exist_ok=True)

        def _atomic_write() -> None:
            with tempfile.TemporaryDirectory(prefix="faiss_", dir=str(scratch_root)) as td:
                tmp_dir = Path(td)
                tmp_idx = tmp_dir / "index.faiss"
                tmp_meta = tmp_dir / "meta.json"
                faiss.write_index(ti.index, os.fspath(tmp_idx))
                tmp_meta.write_text(payload, encoding="utf-8")
                for final, tmp in ((idx_path, tmp_idx), (meta_path, tmp_meta)):
                    final_s = os.fspath(final)
                    tmp_s = os.fspath(tmp)
                    if os.path.exists(final_s):
                        os.remove(final_s)
                    os.replace(tmp_s, final_s)

        try:
            _atomic_write()
        except OSError as exc:
            logger.error(
                "FAISS persist failed for %s/%s (%s); quarantining and retrying once.",
                tenant_id,
                tier,
                exc,
            )
            for stale in idx_path.parent.glob("*.write.*"):
                stale.unlink(missing_ok=True)
            for final in (idx_path, meta_path):
                if final.is_file():
                    bad = final.with_suffix(final.suffix + ".bad")
                    if bad.is_file():
                        bad.unlink(missing_ok=True)
                    os.replace(os.fspath(final), os.fspath(bad))
            _atomic_write()

    def _get(self, tenant_id: str, tier: str) -> _TierIndex:
        key = (tenant_id, tier)
        with self._lock:
            if key not in self._cache:
                self._cache[key] = self._load(tenant_id, tier)
            return self._cache[key]

    # ── ingest ─────────────────────────────────────────────────────────────--
    def ingest_document(
        self,
        tenant_id: str,
        doc_id: str,
        chunks: list[Chunk],
        *,
        tier: str,
        source_filename: str = "",
    ) -> int:
        """Embed and add chunks. REFERENCE-tier chunks are PII-scrubbed before storage."""
        prepared: list[Chunk] = []
        audit_entries: list[dict] = []
        # One session per document so the same name/address/reference is masked
        # to one stable token across every chunk of this file (referential
        # integrity), while distinct values remain distinguishable.
        scrub_session = (
            pii_scrubber.ScrubSession() if tier == TIER_REFERENCE else None
        )
        for c in chunks:
            text = c.text or ""
            scrubbed = c.is_scrubbed
            chunk_audit: dict[str, int] = {}
            if tier == TIER_REFERENCE and not scrubbed:
                result = pii_scrubber.scrub_reference_for_ingest(
                    text, session=scrub_session
                )
                if result is None:
                    logger.warning(
                        "Skipping REFERENCE chunk from %s (PII could not be fully redacted)",
                        doc_id,
                    )
                    continue
                text = result.text
                chunk_audit = dict(result.audit)
                scrubbed = True
            if not text.strip():
                continue
            if chunk_audit:
                audit_entries.append({
                    "redactions": [
                        {"type": k, "count": v} for k, v in sorted(chunk_audit.items())
                    ],
                    "is_scrubbed": True,
                })
            prepared.append(
                Chunk(
                    text=text,
                    section_id=c.section_id,
                    doc_id=c.doc_id or doc_id,
                    tier=tier,
                    is_scrubbed=scrubbed,
                    chunk_id=c.chunk_id,
                    source_filename=c.source_filename or source_filename,
                    paragraph_index=c.paragraph_index,
                )
            )
        if not prepared:
            return 0

        if tier == TIER_REFERENCE and audit_entries:
            self._append_scrub_audit(
                tenant_id,
                doc_id=doc_id,
                chunks_processed=len(audit_entries),
                chunk_audits=audit_entries,
            )

        vecs = self._embedder.embed_documents([c.text for c in prepared])
        arr = np.asarray(vecs, dtype="float32")
        with self._lock:
            ti = self._get(tenant_id, tier)
            ti.index.add(arr)
            base = len(ti.meta)
            for i, c in enumerate(prepared):
                cid = c.chunk_id or f"{doc_id}:{base + i}"
                src = c.source_filename or source_filename or doc_id.split(":", 1)[-1]
                ti.meta.append({
                    "text": c.text,
                    "paragraph_text": c.text,
                    "section_id": c.section_id,
                    "doc_id": c.doc_id or doc_id,
                    "tier": tier,
                    "is_scrubbed": c.is_scrubbed,
                    "chunk_id": cid,
                    "source_filename": src,
                    "paragraph_index": c.paragraph_index or 0,
                })
            self._persist(tenant_id, tier, ti)
        logger.info("Ingested %d chunks into %s/%s (doc=%s)",
                    len(prepared), tenant_id, tier, doc_id)
        return len(prepared)

    @staticmethod
    def _append_scrub_audit(
        tenant_id: str,
        *,
        doc_id: str,
        chunks_processed: int,
        chunk_audits: list[dict],
    ) -> None:
        """Persist scrub audit metadata (types and counts only, never originals)."""
        record = {
            "document": doc_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "chunks_processed": chunks_processed,
            "chunks": chunk_audits,
        }
        path = tenant_store.scrub_audit_path(tenant_id)
        with path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    def clear_tier(self, tenant_id: str, tier: str) -> None:
        """Drop all vectors for a (tenant, tier) — used by admin override."""
        with self._lock:
            ti = _TierIndex(index=self._new_index(), meta=[])
            self._cache[(tenant_id, tier)] = ti
            self._persist(tenant_id, tier, ti)

    def evict_tier_cache(self, tenant_id: str, tier: str) -> None:
        """Drop in-memory index without persisting (used before disk purge)."""
        with self._lock:
            self._cache.pop((tenant_id, tier), None)

    def remove_document(
        self,
        tenant_id: str,
        tier: str,
        *,
        source_filename: str | None = None,
        doc_id: str | None = None,
    ) -> int:
        """Drop all chunks for a reference upload and rebuild the FAISS index."""
        if not source_filename and not doc_id:
            return 0

        with self._lock:
            ti = self._get(tenant_id, tier)
            kept: list[dict] = []
            removed = 0
            for row in ti.meta:
                sf = str(row.get("source_filename") or "").strip()
                did = str(row.get("doc_id") or "").strip()
                match = False
                if source_filename and sf == source_filename:
                    match = True
                if doc_id and did == doc_id:
                    match = True
                if source_filename and did == f"reference:{source_filename}":
                    match = True
                if match:
                    removed += 1
                else:
                    kept.append(row)

            if removed == 0:
                return 0

            if not kept:
                ti = _TierIndex(index=self._new_index(), meta=[])
            else:
                texts = [r["text"] for r in kept]
                vecs = self._embedder.embed_documents(texts)
                arr = np.asarray(vecs, dtype="float32")
                index = self._new_index()
                index.add(arr)
                ti = _TierIndex(index=index, meta=kept)

            self._cache[(tenant_id, tier)] = ti
            self._persist(tenant_id, tier, ti)
            logger.info(
                "Removed %d chunk(s) from tenant=%s tier=%s (source=%s doc_id=%s)",
                removed,
                tenant_id,
                tier,
                source_filename,
                doc_id,
            )
            return removed

    def count(self, tenant_id: str, tier: str) -> int:
        return len(self._get(tenant_id, tier).meta)

    def list_source_filenames(self, tenant_id: str, tier: str) -> list[str]:
        """Unique source filenames ingested for a tier."""
        seen: set[str] = set()
        out: list[str] = []
        for row in self._get(tenant_id, tier).meta:
            name = str(row.get("source_filename") or row.get("doc_id") or "").strip()
            if name.startswith("reference:"):
                name = name.split(":", 1)[-1]
            if name and name not in seen:
                seen.add(name)
                out.append(name)
        return out

    def sample_chunk_texts(self, tenant_id: str, tier: str, *, limit: int = 40) -> list[str]:
        """Return up to ``limit`` chunk texts from a tier (for style analysis)."""
        texts: list[str] = []
        for row in self._get(tenant_id, tier).meta:
            text = str(row.get("text") or "").strip()
            if text:
                texts.append(text)
            if len(texts) >= limit:
                break
        return texts

    def fetch_section_chunks(
        self,
        tenant_id: str,
        *,
        tier: str,
        section_id: str,
        source_key: str | None = None,
        allowed_doc_keys: frozenset[str] | None = None,
    ) -> list[SearchHit]:
        """Return EVERY stored chunk for a ``(tenant, tier, section_id)`` in document
        order — a metadata scan, NOT a similarity search.

        This backs section-complete baseline assembly: the full past-report section is
        mapped, not only the top-K semantically nearest chunks. REFERENCE chunks pass
        the same read-time scrub backstop as :meth:`search` (unscrubbed dropped, text
        re-sanitised). When ``source_key`` is given, only chunks from that one document
        (matched on ``source_filename`` or ``doc_id``) are returned, pinning assembly to
        a single report. When ``allowed_doc_keys`` is set, the upload allowlist applies.
        """
        want = (section_id or "").strip().upper()
        if not want:
            return []
        ti = self._get(tenant_id, tier)
        out: list[SearchHit] = []
        for m in ti.meta:
            if (m.get("section_id") or "").strip().upper() != want:
                continue
            if allowed_doc_keys is not None and not meta_matches_allowlist(m, allowed_doc_keys):
                continue
            if source_key is not None:
                key = (m.get("source_filename") or m.get("doc_id") or "").strip()
                if key != source_key:
                    continue
            if tier == TIER_REFERENCE and not m.get("is_scrubbed"):
                continue
            text = m.get("text") or ""
            if tier == TIER_REFERENCE:
                text = pii_scrubber.sanitize_for_generation_context(text)
            if not text.strip():
                continue
            # Uniform rank: ordering is by document position, not similarity.
            out.append(_search_hit_from_meta(m, tier=tier, rank=1.0, text=text))
        out.sort(key=lambda h: (h.paragraph_index or 0, h.chunk_id or ""))
        return out

    # ── search ─────────────────────────────────────────────────────────────--
    def search(
        self,
        tenant_id: str,
        query: str,
        *,
        tier: str | None = None,
        section_id: str | None = None,
        top_k: int = 5,
        section_strict: bool = False,
        allowed_doc_keys: frozenset[str] | None = None,
        reference_document_ids: list[str] | None = None,
        strict_uploaded_only: bool = False,
    ) -> list[SearchHit]:
        """Search one or both tiers, ranked by similarity.

        When ``tier`` is None, MASTER hits are ranked ahead of REFERENCE hits at
        equal relevance (master is authoritative). When ``section_id`` is given,
        chunks tagged with that section are boosted.

        When ``section_strict`` is True, only chunks whose ``section_id`` equals
        ``section_id`` are returned (used to pin mapping to the correct paragraph).

        When ``allowed_doc_keys`` is set, only chunks whose ``doc_id`` or
        ``source_filename`` appears in the allowlist are returned.

        When ``strict_uploaded_only`` is True and ``reference_document_ids`` is
        non-empty, builds an allowlist via :func:`build_reference_allowlist` and
        drops hits that fail :func:`meta_matches_allowlist` on ``doc_id`` or
        ``source_filename``.
        """
        if allowed_doc_keys is None and strict_uploaded_only:
            allowed_doc_keys = build_reference_allowlist(
                tenant_id,
                reference_document_ids,
                strict_uploaded_only=True,
            )
        tiers = [tier] if tier else [TIER_MASTER, TIER_REFERENCE]
        hits: list[SearchHit] = []
        qvec = np.asarray([self._embedder.embed_query(query)], dtype="float32")
        pool_mult = 20 if allowed_doc_keys else 3

        for t in tiers:
            ti = self._get(tenant_id, t)
            n = len(ti.meta)
            if n == 0:
                continue
            k = min(max(top_k * pool_mult, top_k), n)
            scores, idxs = ti.index.search(qvec, k)
            for score, i in zip(scores[0], idxs[0]):
                if i < 0:
                    continue
                m = ti.meta[i]
                if allowed_doc_keys is not None and not meta_matches_allowlist(m, allowed_doc_keys):
                    continue
                if section_strict and section_id and m.get("section_id") != section_id:
                    continue
                # Reference backstop: never surface unscrubbed reference chunks.
                if t == TIER_REFERENCE and not m.get("is_scrubbed"):
                    continue
                text = m["text"]
                if t == TIER_REFERENCE:
                    text = pii_scrubber.sanitize_for_generation_context(text)
                    if not text.strip():
                        continue
                rank = float(score)
                if t == TIER_MASTER:
                    rank += 0.05  # authoritative tier nudge
                if section_id and m.get("section_id") == section_id:
                    rank += settings.retrieval_section_boost
                hits.append(_search_hit_from_meta(m, tier=t, rank=rank, text=text))

        hits.sort(key=lambda h: h.score, reverse=True)
        return hits[:top_k]

    def search_for_generation(
        self,
        tenant_id: str,
        query: str,
        *,
        section_id: str | None = None,
        top_k: int = 5,
        section_strict: bool = False,
    ) -> list[SearchHit]:
        """Retrieve paragraphs for report mapping — MASTER boilerplate only.

        Past completed reports (REFERENCE tier) are never passed to the mapping
        step, so another property's sensitive data cannot enter the current report.
        """
        hits = self.search(
            tenant_id,
            query,
            tier=TIER_MASTER,
            section_id=section_id,
            top_k=top_k,
            section_strict=section_strict,
        )
        safe: list[SearchHit] = []
        for h in hits:
            text = pii_scrubber.sanitize_for_generation_context(h.text)
            if not text.strip():
                continue
            safe.append(
                _search_hit_from_meta(
                    {
                        "section_id": h.section_id,
                        "doc_id": h.doc_id,
                        "is_scrubbed": h.is_scrubbed,
                        "source_filename": h.source_filename,
                        "paragraph_index": h.paragraph_index,
                        "chunk_id": h.chunk_id,
                    },
                    tier=h.tier,
                    rank=h.score,
                    text=text,
                )
            )
        return safe

    def search_for_reference_mapping(
        self,
        tenant_id: str,
        query: str,
        *,
        section_id: str | None = None,
        top_k: int = 5,
        section_strict: bool = False,
        allowed_doc_keys: frozenset[str] | None = None,
    ) -> list[SearchHit]:
        """Retrieve scrubbed past-report excerpts for minimum-AI mapping."""
        hits = self.search(
            tenant_id,
            query,
            tier=TIER_REFERENCE,
            section_id=section_id,
            top_k=top_k,
            section_strict=section_strict,
            allowed_doc_keys=allowed_doc_keys,
        )
        safe: list[SearchHit] = []
        for h in hits:
            text = pii_scrubber.sanitize_for_generation_context(h.text)
            if not text.strip():
                continue
            safe.append(
                _search_hit_from_meta(
                    {
                        "section_id": h.section_id,
                        "doc_id": h.doc_id,
                        "is_scrubbed": h.is_scrubbed,
                        "source_filename": h.source_filename,
                        "paragraph_index": h.paragraph_index,
                        "chunk_id": h.chunk_id,
                    },
                    tier=h.tier,
                    rank=h.score,
                    text=text,
                )
            )
        return safe


_instance: RagStore | None = None


def get_rag_store() -> RagStore:
    global _instance
    if _instance is None:
        _instance = RagStore()
    return _instance


def reset_rag_store() -> None:
    global _instance
    _instance = None