import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA

def main():
    load_dotenv()
    
    if not os.getenv("OPENAI_API_KEY"):
        print("Please set OPENAI_API_KEY in .env file")
        return

    print("Loading documents...")
    loader = DirectoryLoader("./data", glob="**/*.pdf", loader_cls=PyPDFLoader)
    docs = loader.load()
    if not docs:
        print("No documents found in ./data folder.")
        return

    print(f"Loaded {len(docs)} documents.")

    print("Splitting documents...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)

    print("Creating vector store...")
    vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())

    print("Initializing QA chain...")
    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
    qa_chain = RetrievalQA.from_chain_type(
        llm,
        retriever=vectorstore.as_retriever(),
        return_source_documents=True
    )

    while True:
        query = input("\nAsk a question (or 'q' to quit): ")
        if query.lower() == 'q':
            break
        
        result = qa_chain.invoke({"query": query})
        print(f"\nAnswer: {result['result']}")
        print("\nSources:")
        for source in result['source_documents']:
            print(f"- {source.metadata.get('source', 'Unknown')}")

if __name__ == "__main__":
    main()
