import os
from pathlib import Path
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain.prompts import PromptTemplate

BASE_DIR = Path(__file__).resolve().parent.parent
DB_DIR = BASE_DIR / "db"

def build_general_qa_chain(model_name=None):
    if not DB_DIR.exists():
        print("📦 No DB found. Building vectorstore...")
        from scripts import load_documents, chunk_and_embed, setup_vectorstore
        load_documents.main()
        chunk_and_embed.main()
        setup_vectorstore.main()

    embedding = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma(persist_directory=str(DB_DIR), embedding_function=embedding)

    template = """Use the following context to answer the question. 
    If the answer isn't found in the context, use your general knowledge but say so.
    Do not answer questions that are completely irrelevant to the main points of the context.
    Always cite your sources at the end with 'Source: <filename>' when using course materials.

    Context: {context}
    Question: {question}
    Helpful Answer:"""

    QA_PROMPT = PromptTemplate(
        template=template,
        input_variables=["context", "question"]
    )

    llm = ChatOpenAI(model_name=model_name or "gpt-4o-mini", temperature=0.0) 
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
        chain_type_kwargs={"prompt": QA_PROMPT},
        return_source_documents=True
    )
    return qa_chain