app/retrieval/retriever.py

"""Tenant-scoped retriever backed by the LangChain vector store.

In the LangChain experiment branch, the retriever no longer computes
query embeddings manually.  Instead, it delegates directly to
``VectorStore.search(query_text, tenant_id, k)`` — the LangChain store
handles embedding the query internally.

Compared to master branch:
- ✅ Simpler interface — no ``get_embedding_client()`` call
- ✅ Single responsibility — retrieval only
- ✅ Tenant isolation guaranteed by the vector store (FAISS post-filter by ``tenant_id``)
"""

import logging
from functools import partial

from app.async_executor import run_sync_in_executor
from app.config import settings
from app.models.schemas import SearchResult
from app.retrieval.chunk_role import role_priority
from app.retrieval.semantic_cache import get_semantic_cache, is_semantic_cache_active
from app.retrieval.vector_search import async_vs_search
from app.vectorstore.factory import get_vectorstore

logger = logging.getLogger(__name__)


def _scoped_cache_query(
    query: str,
    *,
    hierarchy_level: str | None = None,
    doc_id_in: frozenset[str] | None = None,
) -> str:
    """Distinct cache key when hierarchy or document scope differs."""
    if not hierarchy_level and not doc_id_in:
        return query
    doc_sig = ",".join(sorted(doc_id_in)) if doc_id_in else ""
    return f"{query}\n#scope:{hierarchy_level or ''}:{doc_sig}"


async def retrieve_scoped_async(
    query: str,
    tenant_id: str,
    k: int,
    *,
    hierarchy_level: str | None = None,
    doc_id_in: frozenset[str] | None = None,
) -> list[SearchResult]:
    """Scoped async retrieval (hierarchy / doc filter) with optional semantic cache."""
    cache_query = _scoped_cache_query(
        query, hierarchy_level=hierarchy_level, doc_id_in=doc_id_in
    )
    cache = get_semantic_cache()
    if cache is not None:
        cached = await cache.get(cache_query, tenant_id)
        if cached is not None:
            logger.debug(
                "Semantic cache hit (scoped) tenant=%s k=%d",
                tenant_id,
                min(k, len(cached)),
            )
            return cached[:k]

    vs = get_vectorstore()
    results = await async_vs_search(
        vs,
        query,
        tenant_id,
        k=k,
        hierarchy_level=hierarchy_level,
        doc_id_in=doc_id_in,
    )
    if cache is not None and results:
        await cache.put(cache_query, tenant_id, results)
    return results


async def retrieve_async(
    query: str,
    tenant_id: str,
    k: int | None = None,
) -> list[SearchResult]:
    """Async retrieval with optional semantic cache (Qdrant when enabled)."""
    top_k = k if k is not None else settings.retrieval_top_k
    cache = get_semantic_cache()
    if cache is not None:
        cached = await cache.get(query, tenant_id)
        if cached is not None:
            logger.debug(
                "Semantic cache hit tenant=%s k=%d", tenant_id, min(top_k, len(cached))
            )
            return cached[:top_k]

    vs = get_vectorstore()
    results = await async_vs_search(vs, query, tenant_id, k=top_k)
    if cache is not None and results:
        await cache.put(query, tenant_id, results)
    logger.debug(
        "Retrieved %d results (async) for tenant=%s (k=%d)",
        len(results),
        tenant_id,
        top_k,
    )
    return results


def use_async_retrieval_path() -> bool:
    """Use async retrieval (and semantic cache when Qdrant cache is configured)."""
    return bool(settings.enable_async_pipeline or is_semantic_cache_active())


def _use_async_retrieval_path() -> bool:
    return use_async_retrieval_path()


async def retrieve_unified(
    query: str,
    tenant_id: str,
    k: int | None = None,
) -> list[SearchResult]:
    """Async-safe retrieval: cache/async path or executor-wrapped sync search."""
    top_k = k if k is not None else settings.retrieval_top_k
    if _use_async_retrieval_path():
        return await retrieve_async(query=query, tenant_id=tenant_id, k=top_k)
    return await run_sync_in_executor(
        partial(retrieve, query=query, tenant_id=tenant_id, k=top_k)
    )


async def retrieve_for_report_unified(
    query: str,
    tenant_id: str,
    primary_document_id: str | None,
    k: int | None = None,
    secondary_document_ids: list[str] | None = None,
) -> list[SearchResult]:
    """Async-safe :func:`retrieve_for_report` with optional semantic cache."""
    if _use_async_retrieval_path():
        return await retrieve_for_report_async(
            query=query,
            tenant_id=tenant_id,
            primary_document_id=primary_document_id,
            k=k,
            secondary_document_ids=secondary_document_ids,
        )
    return await run_sync_in_executor(
        partial(
            retrieve_for_report,
            query=query,
            tenant_id=tenant_id,
            primary_document_id=primary_document_id,
            k=k,
            secondary_document_ids=secondary_document_ids,
        )
    )


def retrieve(
    query: str,
    tenant_id: str,
    k: int | None = None,
) -> list[SearchResult]:
    """Retrieve the top-k chunks nearest to ``query`` for ``tenant_id``.

    The query is embedded internally by the LangChain vector store —
    no manual vector computation is required by the caller.

    Tenant isolation is enforced by post-filtering FAISS hits so only the
    requested ``tenant_id`` is returned.

    Args:
        query: Plain-text retrieval query (typically concatenated fact bullets).
        tenant_id: Restrict results to this tenant's documents.
        k: Number of candidates to retrieve.  Defaults to
            ``settings.retrieval_top_k``.

    Returns:
        List of :class:`~app.models.schemas.SearchResult` ordered by
        descending similarity score.

    Example::

        results = retrieve("Victorian terrace, 95 sqm", tenant_id="t-abc")
    """
    top_k = k if k is not None else settings.retrieval_top_k
    vs = get_vectorstore()
    results = vs.search(query=query, tenant_id=tenant_id, k=top_k)
    logger.debug(
        "Retrieved %d results for tenant=%s (k=%d)", len(results), tenant_id, top_k
    )
    return results


def reorder_by_chunk_role(
    results: list[SearchResult],
    boost_boilerplate: bool = True,
) -> list[SearchResult]:
    """Stable-sort retrieval hits to bring ``chunk_role='boilerplate'`` first.

    Used at assembly / low AI-involvement tiers so the firm's approved
    standard wording is preferentially placed in the LLM context (and in the
    deterministic stitcher's input). Original similarity order is preserved
    within each role bucket.
    """
    if not boost_boilerplate or not results:
        return results
    return sorted(results, key=lambda r: role_priority(r.chunk_role))


def retrieve_for_report(
    query: str,
    tenant_id: str,
    primary_document_id: str | None,
    k: int | None = None,
    secondary_document_ids: list[str] | None = None,
) -> list[SearchResult]:
    """Retrieve chunks for ``tenant_id``, prioritising the report's source upload.

    All chunks remain tenant-scoped (privacy). When ``primary_document_id`` is
    the UUID of the document the user attached when creating this report, chunks
    from that file are listed first so RAG favours the user's RICS / survey
    upload over any other files they may have ingested under the same tenant.

    Optional ``secondary_document_ids`` (e.g. exemplar / Beh Rang reference PDFs)
    are ordered after the primary document and before the rest of the library.

    Args:
        query: Semantic query (e.g. expanded bullets).
        tenant_id: Isolated tenant identifier.
        primary_document_id: ``Report.document_id`` or ``None`` to use global ordering.
        k: Result cap (defaults to ``settings.retrieval_top_k``).
        secondary_document_ids: Additional document UUIDs to prioritise after primary.

    Returns:
        Up to ``k`` :class:`SearchResult` rows, primary-document chunks first.
    """
    if k is None:
        k = settings.retrieval_top_k

    fetch_n = max(k * 4, min(48, k * 6))
    pool = retrieve(query=query, tenant_id=tenant_id, k=fetch_n)
    if not primary_document_id:
        return pool[:k]

    secondary = set(secondary_document_ids or [])
    preferred = [r for r in pool if r.doc_id == primary_document_id]
    second = [r for r in pool if r.doc_id in secondary and r.doc_id != primary_document_id]
    rest = [
        r
        for r in pool
        if r.doc_id != primary_document_id and r.doc_id not in secondary
    ]
    merged = preferred + second + rest
    logger.debug(
        "retrieve_for_report: primary_doc=%s secondary=%d preferred=%d total_returned=%d",
        primary_document_id[:8] if primary_document_id else None,
        len(secondary),
        len(preferred),
        min(k, len(merged)),
    )
    return merged[:k]


def retrieve_document_level_context(
    *,
    template_id: str,
    skeleton_excerpt: str,
    tenant_id: str,
    primary_document_id: str | None,
    reference_document_ids: list[str] | None,
    product_label: str | None = None,
) -> list[SearchResult]:
    """Broad semantic pull for whole-document narrative (survey + exemplar PDFs).

    Uses a section-aware query over the template skeleton so hits span report-wide
    phrasing, not only bullet-level detail. Results are capped per document so both
    the raw context PDF and reference RICS exemplars can contribute.
    """
    if not primary_document_id and not (reference_document_ids or []):
        return []

    sk = (skeleton_excerpt or "").strip().replace("\n", " ")[:600]
    pl = (product_label or "RICS Home Survey Level 3 (Building Survey)").strip()
    doc_query = (
        f"{pl} report section {template_id}. "
        f"Professional surveyor narrative and condition description. {sk}"
    )
    fetch_n = max(48, settings.rag_doc_context_max_chunks * 12)
    pool = retrieve(query=doc_query, tenant_id=tenant_id, k=fetch_n)

    ref_set = {str(x) for x in (reference_document_ids or []) if x}
    allowed: set[str] = set()
    if primary_document_id:
        allowed.add(str(primary_document_id))
    allowed.update(ref_set)

    filtered = [r for r in pool if r.doc_id in allowed]
    if not filtered:
        return []

    primary_max = settings.rag_doc_chunks_primary
    per_ref = settings.rag_doc_chunks_per_reference
    total_max = settings.rag_doc_context_max_chunks
    ref_counts: dict[str, int] = {}
    primary_count = 0
    out: list[SearchResult] = []
    for r in filtered:
        if len(out) >= total_max:
            break
        did = r.doc_id
        if did == primary_document_id:
            if primary_count < primary_max:
                out.append(r)
                primary_count += 1
        elif did in ref_set:
            n = ref_counts.get(did, 0)
            if n < per_ref:
                out.append(r)
                ref_counts[did] = n + 1
    logger.debug(
        "retrieve_document_level_context: section=%s chunks=%d primary=%d",
        template_id,
        len(out),
        primary_count,
    )
    return out


async def retrieve_for_report_async(
    query: str,
    tenant_id: str,
    primary_document_id: str | None,
    k: int | None = None,
    secondary_document_ids: list[str] | None = None,
) -> list[SearchResult]:
    """Async variant of :func:`retrieve_for_report`."""
    if k is None:
        k = settings.retrieval_top_k
    fetch_n = max(k * 4, min(48, k * 6))
    pool = await retrieve_async(query=query, tenant_id=tenant_id, k=fetch_n)
    if not primary_document_id:
        return pool[:k]
    secondary = set(secondary_document_ids or [])
    preferred = [r for r in pool if r.doc_id == primary_document_id]
    second = [r for r in pool if r.doc_id in secondary and r.doc_id != primary_document_id]
    rest = [
        r
        for r in pool
        if r.doc_id != primary_document_id and r.doc_id not in secondary
    ]
    return (preferred + second + rest)[:k]


async def retrieve_document_level_context_async(
    *,
    template_id: str,
    skeleton_excerpt: str,
    tenant_id: str,
    primary_document_id: str | None,
    reference_document_ids: list[str] | None,
    product_label: str | None = None,
) -> list[SearchResult]:
    """Async variant of :func:`retrieve_document_level_context`."""
    if not primary_document_id and not (reference_document_ids or []):
        return []
    sk = (skeleton_excerpt or "").strip().replace("\n", " ")[:600]
    pl = (product_label or "RICS Home Survey Level 3 (Building Survey)").strip()
    doc_query = (
        f"{pl} report section {template_id}. "
        f"Professional surveyor narrative and condition description. {sk}"
    )
    fetch_n = max(48, settings.rag_doc_context_max_chunks * 12)
    pool = await retrieve_async(query=doc_query, tenant_id=tenant_id, k=fetch_n)
    ref_set = {str(x) for x in (reference_document_ids or []) if x}
    allowed: set[str] = set()
    if primary_document_id:
        allowed.add(str(primary_document_id))
    allowed.update(ref_set)
    filtered = [r for r in pool if r.doc_id in allowed]
    if not filtered:
        return []
    primary_max = settings.rag_doc_chunks_primary
    per_ref = settings.rag_doc_chunks_per_reference
    total_max = settings.rag_doc_context_max_chunks
    ref_counts: dict[str, int] = {}
    primary_count = 0
    out: list[SearchResult] = []
    for r in filtered:
        if len(out) >= total_max:
            break
        did = r.doc_id
        if did == primary_document_id:
            if primary_count < primary_max:
                out.append(r)
                primary_count += 1
        elif did in ref_set:
            n = ref_counts.get(did, 0)
            if n < per_ref:
                out.append(r)
                ref_counts[did] = n + 1
    return out