"""
Retrieval and answer generation logic using LangChain.
"""

import os
from typing import Any

from chromadb import Client
from chromadb.config import Settings

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms.openai import OpenAIChat
from langchain.chains import RetrievalQA
from langchain.vectorstores import Chroma

def get_answer(question: str, client: Client, collection_name: str, k: int = 3) -> str:
    """
    Retrieve relevant FAQ chunks and generate an answer using OpenAIChat.
    """
    # Set up embeddings and LLM
    embedding = OpenAIEmbeddings()
    llm = OpenAIChat(temperature=0)

    # Load vector store
    vectorstore = Chroma(
        client=client,
        collection_name=collection_name,
        embedding_function=embedding
    )

    # Build RetrievalQA chain
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(search_kwargs={"k": k}),
        return_source_documents=True
    )

    # Run chain
    result = qa_chain({"question": question})
    answer = result.get("answer", "")
    return answer.strip()