import os
import torch
from sentence_transformers import SentenceTransformer

# Resilient LangChain Imports
try:
    from langchain_community.retrievers import BM25Retriever
    from langchain.retrievers import EnsembleRetriever
    from langchain_core.documents import Document
except ImportError:
    # Fallback for older versions or missing community package
    from langchain.retrievers import BM25Retriever, EnsembleRetriever
    from langchain.schema import Document

class AdvancedRAG:
    """Advanced RAG system with Hybrid Search (BM25 + Vector) and Reranking."""
    
    def __init__(self, vector_db_collection, embedder_model='all-MiniLM-L6-v2'):
        self.collection = vector_db_collection
        self.embedder = SentenceTransformer(embedder_model, device='cuda' if torch.cuda.is_available() else 'cpu')
        self.bm25_retriever = None

    def initialize_bm25(self, texts):
        """Initialize BM25 with a corpus of texts."""
        if not texts:
            return
        documents = [Document(page_content=t) for t in texts]
        self.bm25_retriever = BM25Retriever.from_documents(documents)
        self.bm25_retriever.k = 5

    def hybrid_search(self, query, top_k=5):
        """Perform hybrid search combining vector similarity and keyword matching."""
        # 1. Vector Search
        query_embedding = self.embedder.encode(query).tolist()
        vector_results = self.collection.query(query_embeddings=[query_embedding], n_results=top_k)
        vector_docs = vector_results['documents'][0] if vector_results['documents'] else []
        
        # 2. BM25 Search (if initialized)
        bm25_docs = []
        if self.bm25_retriever:
            try:
                results = self.bm25_retriever.get_relevant_documents(query)
                bm25_docs = [doc.page_content for doc in results]
            except: pass

        # 3. Ensemble (Simple deduplication and ranking)
        combined = list(dict.fromkeys(vector_docs + bm25_docs))
        return combined[:top_k]

    def add_intelligence(self, new_text):
        """Add new information to the brain."""
        embedding = self.embedder.encode(new_text).tolist()
        self.collection.add(
            documents=[new_text],
            embeddings=[embedding],
            ids=[str(hash(new_text))]
        )