from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain_groq import ChatGroq
from langchain.chains import ConversationalRetrievalChain
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from config import VECTOR_DIR, EMBED_MODEL, GROQ_API_KEY, GROQ_MODEL

# ---- Memory ----
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

def reset_memory():
    memory.clear()

def build_chain():
    embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
    vectordb = FAISS.load_local(VECTOR_DIR, embeddings, allow_dangerous_deserialization=True)
    retriever = vectordb.as_retriever(search_kwargs={"k": 3})

    llm = ChatGroq(model=GROQ_MODEL, api_key=GROQ_API_KEY, temperature=0.1)

    # Simple combined prompt
    chat_prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template(
            "You are a helpful assistant. Use the context to answer user questions."
        ),
        HumanMessagePromptTemplate.from_template(
            "Context:\n{context}\n\nQuestion:\n{question}\nAnswer clearly and concisely:"
        )
    ])

    chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        combine_docs_chain_kwargs={"prompt": chat_prompt}
    )
    return chain

def answer(question: str) -> str:
    chain = build_chain()
    result = chain.invoke({"question": question})
    return result.get("answer") or result.get("result") or "No answer."