src/chatbot/rag_pipeline.py

import threading
from typing import Optional

# Vector search backend — prefer ChromaDB, fall back to FAISS, then keyword
try:
    import chromadb
    from chromadb.utils import embedding_functions
    VECTOR_BACKEND = "chroma"
except ImportError:
    try:
        import faiss
        import numpy as np
        from sentence_transformers import SentenceTransformer
        VECTOR_BACKEND = "faiss"
    except ImportError:
        VECTOR_BACKEND = "none"

from src.chatbot.document_hub import get_document


# Configuration
TOP_K = 5
EMBED_MODEL = "all-MiniLM-L6-v2"


# Backend state (lazy-initialised)
_lock = threading.Lock()

# ChromaDB
_chroma_client = None
_chroma_collection = None

# FAISS
_faiss_index = None
_faiss_metadata: list[dict] = []
_embed_model = None


def _get_chroma_collection():
    global _chroma_client, _chroma_collection
    if _chroma_collection is None:
        _chroma_client = chromadb.Client()
        ef = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=EMBED_MODEL
        )
        _chroma_collection = _chroma_client.get_or_create_collection(
            name="rag_chunks",
            embedding_function=ef,
            metadata={"hnsw:space": "cosine"},
        )
    return _chroma_collection


def _get_faiss_model():
    global _embed_model
    if _embed_model is None:
        _embed_model = SentenceTransformer(EMBED_MODEL)
    return _embed_model



# Public API
def index_document(doc_id: str) -> dict:
    """
    Embed all chunks from a processed document and add them to the vector
    index.  Chunks are now dicts {"text", "page"} — page number is stored
    as metadata so it survives retrieval.

    Returns
    -------
    dict: { status, doc_id, chunks_indexed }
    """
    doc = get_document(doc_id)
    if not doc:
        return {"status": "error", "message": f"Document '{doc_id}' not found in DocumentHub."}

    chunks = doc["chunks"]        # list[dict]  {"text": str, "page": int}
    filename = doc["filename"]

    with _lock:
        if VECTOR_BACKEND == "chroma":
            _index_chroma(doc_id, filename, chunks)
        elif VECTOR_BACKEND == "faiss":
            _index_faiss(doc_id, filename, chunks)
        else:
            print("[RAGPipeline] No vector backend — keyword search will be used.")

    print(f"[RAGPipeline] Indexed {len(chunks)} chunks for '{filename}' (backend={VECTOR_BACKEND})")
    return {"status": "success", "doc_id": doc_id, "chunks_indexed": len(chunks)}


def retrieve_context(query: str, doc_id: Optional[str] = None, top_k: int = TOP_K) -> dict:
    """
    Find the most relevant chunks for *query*.

    Returns
    -------
    dict:
        context_str  — formatted text block ready for LLM injection
        source       — human-readable source reference, e.g. "page 4"
                       (taken from the top-ranked chunk)
        chunks       — raw list of chunk dicts for callers that need more detail
    """
    with _lock:
        if VECTOR_BACKEND == "chroma":
            chunks = _query_chroma(query, doc_id, top_k)
        elif VECTOR_BACKEND == "faiss":
            chunks = _query_faiss(query, doc_id, top_k)
        else:
            chunks = _keyword_search(query, doc_id, top_k)

    if not chunks:
        return {
            "context_str": "No relevant context found in the uploaded documents.",
            "source": None,
            "chunks": [],
        }

    lines = ["--- RELEVANT DOCUMENT CONTEXT ---"]
    for i, chunk in enumerate(chunks, 1):
        filename = chunk.get("filename", "Unknown")
        page = chunk.get("page")
        page_label = f"Page {page}" if page else "unknown page"
        lines.append(f"\n[Excerpt {i} — {filename}, {page_label}]\n{chunk.get('text', '').strip()}")
    lines.append("\n--- END OF CONTEXT ---")

    # Source = top chunk's page (most relevant result)
    top = chunks[0]
    source = f"page {top['page']}" if top.get("page") else top.get("filename", "unknown")

    return {
        "context_str": "\n".join(lines),
        "source": source,
        "chunks": chunks,
    }


def build_rag_prompt(user_query: str, context_str: str, base_system_prompt: str) -> str:
    """
    Inject *context_str* into *base_system_prompt* so the LLM answers are
    grounded in document content.
    """
    return (
        f"{base_system_prompt.strip()}\n\n"
        "You have access to the following excerpts retrieved from the user's documents. "
        "Use them to answer accurately. If the answer is not in the excerpts, say so.\n\n"
        f"{context_str}"
    )



# ChromaDB helpers
def _index_chroma(doc_id: str, filename: str, chunks: list[dict]) -> None:
    col = _get_chroma_collection()
    texts = [c["text"] for c in chunks]
    ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))]
    metadatas = [
        {"doc_id": doc_id, "filename": filename, "page": c.get("page", 0), "chunk_index": i}
        for i, c in enumerate(chunks)
    ]
    col.upsert(documents=texts, ids=ids, metadatas=metadatas)


def _query_chroma(query: str, doc_id: Optional[str], top_k: int) -> list[dict]:
    col = _get_chroma_collection()
    where = {"doc_id": doc_id} if doc_id else None

    results = col.query(
        query_texts=[query],
        n_results=top_k,
        where=where,
        include=["documents", "metadatas", "distances"],
    )

    chunks = []
    for text, meta in zip(results["documents"][0], results["metadatas"][0]):
        chunks.append({
            "text": text,
            "filename": meta.get("filename", ""),
            "doc_id": meta.get("doc_id", ""),
            "page": meta.get("page"),
        })
    return chunks



# FAISS helpers
def _index_faiss(doc_id: str, filename: str, chunks: list[dict]) -> None:
    global _faiss_index, _faiss_metadata
    model = _get_faiss_model()
    texts = [c["text"] for c in chunks]
    embeddings = model.encode(texts, convert_to_numpy=True)

    if _faiss_index is None:
        _faiss_index = faiss.IndexFlatL2(embeddings.shape[1])

    _faiss_index.add(embeddings)
    for c in chunks:
        _faiss_metadata.append({
            "text": c["text"],
            "doc_id": doc_id,
            "filename": filename,
            "page": c.get("page"),
        })


def _query_faiss(query: str, doc_id: Optional[str], top_k: int) -> list[dict]:
    if _faiss_index is None or _faiss_index.ntotal == 0:
        return []

    model = _get_faiss_model()
    q_vec = model.encode([query], convert_to_numpy=True)
    _, indices = _faiss_index.search(q_vec, min(top_k * 3, _faiss_index.ntotal))

    results = []
    for idx in indices[0]:
        if idx == -1:
            continue
        meta = _faiss_metadata[idx]
        if doc_id and meta["doc_id"] != doc_id:
            continue
        results.append(meta)
        if len(results) >= top_k:
            break
    return results



# Keyword fallback
def _keyword_search(query: str, doc_id: Optional[str], top_k: int) -> list[dict]:
    """Bag-of-words overlap search — used when no vector library is installed."""
    from src.chatbot.document_hub import documents_db, doc_lock

    query_tokens = set(query.lower().split())
    scored: list[tuple[int, dict]] = []

    with doc_lock:
        targets = (
            {doc_id: documents_db[doc_id]}.items()
            if doc_id and doc_id in documents_db
            else documents_db.items()
        )
        for did, doc in targets:
            for chunk in doc["chunks"]:           # chunk is now a dict
                chunk_tokens = set(chunk["text"].lower().split())
                overlap = len(query_tokens & chunk_tokens)
                if overlap > 0:
                    scored.append((overlap, {
                        "text": chunk["text"],
                        "page": chunk.get("page"),
                        "doc_id": did,
                        "filename": doc["filename"],
                    }))

    scored.sort(key=lambda x: x[0], reverse=True)
    return [item for _, item in scored[:top_k]]