import os
from langchain_chroma import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.callbacks import StdOutCallbackHandler
from langchain.memory import ConversationBufferMemory
from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings
from pydantic import SecretStr

CHROMA_DB_DIR = "./chroma_db"

# Load environment variables from .env file
load_dotenv()

# Get API credentials from environment variables
OPENAI_API_KEY = os.getenv("GROQ_API_KEY")
OPENAI_API_BASE = os.getenv("GROQ_API_BASE")

if not OPENAI_API_KEY:
    raise ValueError(
        "OPENAI_API_KEY not found in environment variables. Please check your .env file."
    )
if not OPENAI_API_BASE:
    raise ValueError(
        "OPENAI_API_BASE not found in environment variables. Please check your .env file."
    )

model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "cpu"}
encode_kwargs = {"normalize_embeddings": False}
embeddings = HuggingFaceEmbeddings(
    model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
)


def get_qa_chain():
    vectordb = Chroma(
        persist_directory=None,
        embedding_function=embeddings,
        collection_name="docs_collection",
    )

    print(f"Number of embedded documents: {vectordb._collection.count()}")
    retriever = vectordb.as_retriever(search_kwargs={"k": 3})

    llm = ChatOpenAI(
        model="llama-3.1-8b-instant",
        api_key=SecretStr(OPENAI_API_KEY) if OPENAI_API_KEY else None,
        base_url=OPENAI_API_BASE,
        temperature=0.2,
    )

    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True, output_key="answer"
    )
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm, retriever=retriever, memory=memory, callbacks=[StdOutCallbackHandler()]
    )

    return conversation_chain


def answer_question(question):
    qa_chain = get_qa_chain()
    result = qa_chain.invoke({"question": question})

    answer = result["answer"]

    # Format the answer for better markdown display
    formatted_answer = f"""
## Answer

{answer}

---
*Generated using AI-powered document analysis*
"""

    return formatted_answer