import streamlit as st
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain.memory import ConversationBufferWindowMemory
import os
from dotenv import load_dotenv
import base64

load_dotenv()
os.environ['GROQ_API_KEY'] = os.getenv("GROQ_API_KEY")

def pdf_rag_page():
    st.title('SAM GPT')
    st.sidebar.header("Upload your PDFs")
    uploaded_files = st.sidebar.file_uploader("Choose PDF files", type=["pdf"], accept_multiple_files=True)

    if uploaded_files:
        with st.sidebar:
            st.write("### PDF Previews")
            # Display the PDFs in the sidebar using iframes
            for uploaded_file in uploaded_files:
                pdf_data = uploaded_file.read()
                pdf_base64 = base64.b64encode(pdf_data).decode('utf-8')
                pdf_display = f'<iframe src="data:application/pdf;base64,{pdf_base64}" width="100%" height="400px"></iframe>'
                st.markdown(pdf_display, unsafe_allow_html=True)
                st.write(f"**{uploaded_file.name}**")

    st.write("### Instructions 📜")
    st.write("""
        1. Upload one or more PDF files using the sidebar.
        2. Enter your query in the text input below.
        3. Click the 'Get Results' button to process the PDFs and provide relevant answers based on the content.
    """)

    def load_and_process_pdfs(files):
        all_docs = []
        for file in files:
            with open(file.name, "wb") as f:
                f.write(file.getbuffer())
            loader = PyPDFLoader(file.name)
            docs = loader.load()
            all_docs.extend(docs)

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(all_docs)
        vectorstore = FAISS.from_documents(documents=splits, embedding=OpenAIEmbeddings())
        retriever = vectorstore.as_retriever()
        system_prompt = (
            "You are an assistant for question-answering tasks. "
            "Use the following pieces of retrieved context to answer "
            "the question. If you don't know the answer, say that you "
            "don't know."
            "\n\n"
            "{context}"
        )
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            ("human", "{input}"),
        ])
        llm = ChatGroq(model="llama-3.1-70b-versatile")
        question_answer_chain = create_stuff_documents_chain(llm, prompt)
        rag_chain = create_retrieval_chain(retriever, question_answer_chain)
        return rag_chain

    rag_chain = None
    if uploaded_files:
        rag_chain = load_and_process_pdfs(uploaded_files)

    input_text = st.text_input("Please feel free to ask any doubts! 📝")

    if st.button("Get Results"):
        if input_text and rag_chain:
            with st.spinner('Processing...'):
                try:
                    response = rag_chain.invoke({"input": input_text})
                    st.write(response["answer"])
                except Exception as e:
                    st.error(f"An error occurred: {e}")
        elif input_text:
            st.warning("Please upload PDF files to ask questions.")

if __name__ == "__main__":
    pdf_rag_page()