import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableParallel

class ProjectRAGEngine:
    def __init__(self, api_key):
        self.embeddings = OpenAIEmbeddings(openai_api_key=api_key)
        self.llm = ChatOpenAI(model="gpt-4o", openai_api_key=api_key, temperature=0)
        self.vector_store = None

    def process_documents(self, pdf_paths):
        all_docs = []
        for path in pdf_paths:
            try:
                loader = PyPDFLoader(path)
                docs = loader.load()
                all_docs.extend(docs)
            except Exception as e:
                print(f"Error loading {path}: {e}")
        
        # Splitting logic to handle large reports [cite: 10]
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(all_docs)
        self.vector_store = FAISS.from_documents(splits, self.embeddings)

    def _format_docs(self, docs):
        return "\n\n".join(doc.page_content for doc in docs)

    def get_answer(self, query):
        if not self.vector_store:
            return "Please upload documents first.", []

        # System prompt ensuring grounded responses [cite: 18, 25]
        template = """
        You are a professional Project Analyst. Answer strictly based on the provided context.
        If the answer is not in the context, say you don't know. 
        Cite document names and page numbers for every answer. Include direct quotes.
        
        Context: {context}
        Question: {question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        retriever = self.vector_store.as_retriever(search_kwargs={"k": 5})

        # Pure LCEL Chain composition
        rag_chain_from_docs = (
            RunnablePassthrough.assign(context=(lambda x: self._format_docs(x["context"])))
            | prompt
            | self.llm
            | StrOutputParser()
        )

        rag_chain_with_source = RunnableParallel(
            {"context": retriever, "question": RunnablePassthrough()}
        ).assign(answer=rag_chain_from_docs)

        result = rag_chain_with_source.invoke(query)
        
        sources = [{"content": doc.page_content, "metadata": doc.metadata} for doc in result["context"]]
        return result["answer"], sources