import os
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from langchain_core.prompts import PromptTemplate
# Modern LangChain imports
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_groq import ChatGroq

class ChatbotManager:
    def __init__(self):
        self.embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en", model_kwargs={"device": "cpu"})

        # Using Groq for fast text generation
        self.llm = ChatGroq(
            model_name="llama-3.1-8b-instant",
            temperature=0.1,
            groq_api_key=os.getenv("GROQ_API_KEY")
        )

        # ✅ Initialize Qdrant Client (local path)
        self.client = QdrantClient(path="./qdrant_db")
        self.collection_name = "vector_db"

        # ✅ Load the stored vectors from Qdrant
        self.db = QdrantVectorStore(client=self.client, embeddings=self.embeddings, collection_name=self.collection_name)

        self.prompt = PromptTemplate(
            template="""You are an AI assistant answering questions based on the latest uploaded document only.
            
            Document Context:
            {context}
            
            User Question:
            {input}
            
            Helpful and precise answer:""",
            input_variables=["context", "input"]
        )

        self.retriever = self.db.as_retriever(search_kwargs={"k": 3})
        
        # Modern way to create a retrieval chain
        combine_docs_chain = create_stuff_documents_chain(self.llm, self.prompt)
        self.qa = create_retrieval_chain(self.retriever, combine_docs_chain)

    def close(self):
        if hasattr(self, 'client'):
            self.client.close()

    def get_response(self, query: str) -> str:
        try:
            # Modern chain returns a dict with an 'answer' key
            response = self.qa.invoke({"input": query})
            return response.get("answer", "No answer found.")
        except Exception as e:
            return f"⚠️ Error: {e}"