import numpy as np
from transformers import pipeline

def retrieve_top_k_documents(vector_store, query, top_k=5):
    documents = vector_store.similarity_search(query, k=top_k)
    documents = rerank_documents(query, documents)
    return documents 

# Reranking: Cross-Encoder for refining top-k results
def rerank_documents(query, documents, reranker_model_name="cross-encoder/ms-marco-electra-base"):
    """
    Re-rank documents using a cross-encoder model.

    Parameters:
        query (str): The user's query.
        documents (list): List of LangChain Document objects.
        reranker_model_name (str): Hugging Face model name for re-ranking.

    Returns:
        list: Re-ranked list of Document objects with updated scores.
    """
    # Initialize the cross-encoder model
    reranker = pipeline("text-classification", model=reranker_model_name, return_all_scores=False)

    # Pair the query with each document's text
    rerank_inputs = [{"text": query, "text_pair": doc.page_content} for doc in documents]

    # Get relevance scores for each query-document pair
    scores = reranker(rerank_inputs)

    # Attach the new scores to the documents
    for doc, score in zip(documents, scores):
        doc.metadata["rerank_score"] = score["score"]  # Add score to document metadata

    # Sort documents by the rerank_score in descending order
    documents = sorted(documents, key=lambda x: x.metadata.get("rerank_score", 0), reverse=True)
    return documents


# Query Handling: Retrieve top-k candidates using FAISS with IVF index not used only for learning
def retrieve_top_k_documents_manual(vector_store, query, top_k=5):
    """
    Retrieve top-k documents using FAISS index and optionally rerank them.

    Parameters:
        vector_store (FAISS): The vector store containing the FAISS index and docstore.
        query (str): The user's query string.
        top_k (int): The number of top results to retrieve.
        reranker_model_name (str): The Hugging Face model name for cross-encoder reranking.

    Returns:
        list: Top-k retrieved and reranked documents.
    """
    # Encode the query into a dense vector
    embedding_model = vector_store.embedding_function
    query_vector = embedding_model.embed_query(query)  # Encode the query
    query_vector = np.array([query_vector]).astype('float32')
    
    # Search the FAISS index for top_k results
    distances, indices = vector_store.index.search(query_vector, top_k)

    # Retrieve documents from the docstore
    documents = []
    for idx in indices.flatten():
        if idx == -1:  # FAISS can return -1 for invalid indices
            continue
        doc_id = vector_store.index_to_docstore_id[idx]

        # Access the internal dictionary of InMemoryDocstore
        internal_docstore = getattr(vector_store.docstore, "_dict", None)
        if internal_docstore and doc_id in internal_docstore:  # Check if doc_id exists
            document = internal_docstore[doc_id]
            documents.append(document)

    # Rerank the documents 
    documents = rerank_documents(query, documents)
    
    return documents