from utils.settings import configure_settings
from utils.index import load_index
from utils.constant import INDEX_PATH, TOP_K_RETRIEVAL, TOP_N_RERANKER

from llama_index.core import PromptTemplate
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.postprocessor import LLMRerank
from llama_index.core.query_engine import RetrieverQueryEngine
import Stemmer

class QueryEngineManager:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(QueryEngineManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self._initialized = True
            self.index = None
            self.retriever = None
            self.reranker = None
            self.query_engine = None
            self._configure()
        
    def initialize_index(self):
        self.index = load_index(path=INDEX_PATH)
        self.nodes = list(self.index.docstore.docs.values())
        
    def _configure(self):
        configure_settings()
        self.initialize_index()
        self.reranker = LLMRerank(top_n=TOP_N_RERANKER)
        
    def get_engine(self, bm25: bool = False, semantic: bool = False):
        
        if bm25:
            self.retriever = BM25Retriever.from_defaults(
                nodes=self.nodes,
                stemmer=Stemmer.Stemmer("english"),
                similarity_top_k=TOP_K_RETRIEVAL,
                language="english"
            )
        elif semantic:
            self.retriever = self.index.as_retriever(similarity_top_k=TOP_K_RETRIEVAL)
            
        qa_template = PromptTemplate(
            """Given the following context and question, provide a detailed response.
            Context: {context_str}
            Question: {query_str}
            Let me explain this in detail:""",
            prompt_type="text_qa"
        )
        
        self.query_engine = RetrieverQueryEngine.from_args(
            retriever=self.retriever,
            text_qa_template=qa_template,
            #node_postprocessors=[self.reranker]
        )
        
        return self.query_engine

def get_engine(bm25: bool = False, semantic: bool = False):
    engine_manager = QueryEngineManager()
    return engine_manager.get_engine(bm25, semantic), engine_manager