from typing import List
from sqlalchemy.orm import Session
from app.core.vector_store.faiss_store import FaissVectorStore
from app.db import models
from app.core.retrievers.base_retriever import Retriever

class FaissDBRetriever(Retriever):
    """
    A concrete retriever that uses a FAISS index and a local database
    to find and return relevant document text.
    """
    def __init__(self, vector_store: FaissVectorStore):
        self.vector_store = vector_store

    def retrieve_context(self, query: str, db: Session) -> List[str]:
        """
        Retrieves document text by first searching the FAISS index
        and then fetching the corresponding documents from the database.
        """
        # Pass db session explicitly to vector store method
        faiss_ids = self.vector_store.search_similar_documents(db_session=db, query_text=query, k=3)
        context_docs_text = []

        if faiss_ids:
            # self.print_all_vector_metadata(db)
            # Get document_ids using the FAISS IDs
            document_ids_from_vectors = db.query(models.VectorMetadata.document_id).filter(
                models.VectorMetadata.document_id.in_(faiss_ids)
            ).all()

            document_ids = [doc_id for (doc_id,) in document_ids_from_vectors]
            # Retrieve full documents by their IDs
            context_docs = db.query(models.Document).filter(
                models.Document.id.in_(document_ids)
            ).all()

            context_docs_text = [doc.text for doc in context_docs]

        return context_docs_text

    # def print_all_vector_metadata(self, db: Session):
    #     all_metadata = db.query(models.VectorMetadata).all()
    #     for metadata in all_metadata:
    #         print(f"VectorMetadata - ID: {metadata.id}, Document ID: {metadata.document_id}, Embedding Model: {metadata.embedding_model}")  