import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
from langchain_community.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv
from langchain_openai import OpenAI, ChatOpenAI
from langchain_openai import OpenAIEmbeddings

load_dotenv()
os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
os.environ["LANGCHAIN_TRACING_V2"]="true"
os.environ["LANGCHAIN_API_KEY"] = st.secrets["LANGCHAIN_API_KEY"]

def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            try:
                page_text = page.extract_text()
                if page_text:
                    text += page_text
            except Exception as e:
                print(f"Error reading page: {e}")
    return text

def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2500, chunk_overlap=750)
    chunks = text_splitter.split_text(text)
    return chunks

def get_vector_store(text_chunks):
    vector_store = FAISS.from_texts(text_chunks, OpenAIEmbeddings())
    vector_store.save_local("faiss_index")

def get_conversational_chain():
    prompt_template = """You are an assistant for teachers. Your objective is to provide
    comprehensive and accurate responses based on the context provided. Make sure that 
    you generate whole output.
    context: {context}
    question: {question}
    """

    model = ChatOpenAI(model="gpt-3.5-turbo")
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)

    return chain

def user_input(user_question):
    embeddings = OpenAIEmbeddings()
    new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
    docs = new_db.similarity_search(user_question)

    chain = get_conversational_chain()
    result = ""
    with st.spinner("Processing..."):
        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
        result = response["output_text"]
    
    st.session_state.chat_history.append({"role": "assistant", "content": result})

def main():
    st.set_page_config("Chat PDF")
    st.header("AI-powered EduPlanner💁")

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    with st.sidebar:
        #st.image("pic123.png")
        st.title("Menu:")
        pdf_docs = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button", accept_multiple_files=True)
        if st.button("Submit & Process"):
            if pdf_docs:
                with st.spinner("Processing..."):
                    raw_text = get_pdf_text(pdf_docs)
                    text_chunks = get_text_chunks(raw_text)
                    get_vector_store(text_chunks)
                    st.success("Done")
            else:
                st.warning("Please upload PDF files first before submitting.")

    # Display chat history
    for idx, chat in enumerate(st.session_state.chat_history):
        with st.chat_message(chat["role"]):
            st.write(chat["content"])
            if chat["role"] == "assistant":
                st.download_button(
                    label="Download",
                    data=chat["content"],
                    file_name=f"response_{idx}.txt",
                    mime="text/plain",
                    key=f"download_{idx}",
                )

    user_question = st.chat_input("Ask a Question from the PDF Files")

    if user_question:
        if not pdf_docs:
            st.warning("Please upload PDF files and process first before asking questions.")
        else:
            st.session_state.chat_history.append({"role": "user", "content": user_question})
            st.chat_message("user").write(user_question)
            user_input(user_question)
            st.experimental_rerun()

if __name__ == "__main__":
    main()