# retrieve_raw.py
# Verbatim chunk retrieval: vector search + cross-encoder re-ranking, no LLM.
#
# Returns the top re-ranked chunks with their full text, file metadata, and
# scores. Useful for browsing source material directly and verifying what
# the RAG pipeline retrieves before LLM synthesis.
#
# Uses the same vector store, embedding model, and re-ranker as
# query_topk_prompt_engine_v3.py, but skips the LLM step entirely.
#
# E.M.F. February 2026

# Environment vars must be set before importing huggingface/transformers
# libraries, because huggingface_hub.constants evaluates HF_HUB_OFFLINE
# at import time.
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "./models"
os.environ["HF_HUB_OFFLINE"] = "1"

from llama_index.core import (
    StorageContext,
    load_index_from_storage,
    Settings,
)
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.postprocessor import SentenceTransformerRerank
import sys
import textwrap

#
# Globals
#

# Embedding model (must match build_exp_claude.py)
EMBED_MODEL = HuggingFaceEmbedding(cache_folder="./models", model_name="BAAI/bge-large-en-v1.5", local_files_only=True)

# Cross-encoder model for re-ranking (cached in ./models/)
RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-12-v2"
RERANK_TOP_N = 15
RETRIEVE_TOP_K = 30

# Output formatting
WRAP_WIDTH = 80


def main():
    # No LLM needed -- set embed model only
    Settings.embed_model = EMBED_MODEL

    # Load persisted vector store
    storage_context = StorageContext.from_defaults(persist_dir="./storage_exp")
    index = load_index_from_storage(storage_context)

    # Build retriever (vector search only, no query engine / LLM)
    retriever = index.as_retriever(similarity_top_k=RETRIEVE_TOP_K)

    # Cross-encoder re-ranker
    reranker = SentenceTransformerRerank(
        model=RERANK_MODEL,
        top_n=RERANK_TOP_N,
    )

    # Query
    if len(sys.argv) < 2:
        print("Usage: python retrieve_raw.py QUERY_TEXT")
        sys.exit(1)
    q = " ".join(sys.argv[1:])

    # Retrieve and re-rank
    nodes = retriever.retrieve(q)
    reranked = reranker.postprocess_nodes(nodes, query_str=q)

    # Output
    print(f"\nQuery: {q}")
    print(f"Retrieved {len(nodes)} chunks, re-ranked to top {len(reranked)}\n")

    for i, node in enumerate(reranked, 1):
        meta = getattr(node, "metadata", None) or node.node.metadata
        score = getattr(node, "score", None)
        file_name = meta.get("file_name", "unknown")
        text = node.get_content()

        print("="*WRAP_WIDTH)
        print(f"=== [{i}] {file_name}  (score: {score:.3f}) ")
        print("="*WRAP_WIDTH)
        # Wrap text for readability
        for line in text.splitlines():
            if line.strip():
                print(textwrap.fill(line, width=WRAP_WIDTH))
            else:
                print()
        print()


if __name__ == "__main__":
    main()