import os
from dotenv import load_dotenv
import streamlit as st
import os
from langchain_openai import AzureChatOpenAI
from langchain_openai import AzureOpenAIEmbeddings
from langchain.chains import RetrievalQA

from langchain_community.document_loaders import PyPDFLoader

import streamlit as st
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
import tempfile

from langchain_community.vectorstores import Chroma

load_dotenv()

azure_config = {
    "AZURE_OPENAI_ENDPOINT": os.getenv("AZURE_OPENAI_ENDPOINT"),
    "AZURE_OPENAI_MODEL_DEPLOYMENT_NAME": os.getenv("AZURE_OPENAI_MODEL_DEPLOYMENT_NAME"),
    "AZURE_OPENAI_MODEL_NAME": os.getenv("AZURE_OPENAI_MODEL_NAME"),
    "AZURE_OPENAI_API_KEY": os.getenv("AZURE_OPENAI_API_KEY"),
    "AZURE_OPENAI_API_VERSION": os.getenv("AZURE_OPENAI_API_VERSION"),
    "AZURE_OPENAI_EMBEDDING_ENDPOINT": os.getenv("AZURE_OPENAI_API_VERSION"),
    "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME" : os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"),
    "AZURE_OPENAI_EMBEDDING_NAME": os.getenv("AZURE_OPENAI_EMBEDDING_NAME"),
    "AZURE_OPENAI_EMBEDDING_API_KEY": os.getenv("AZURE_OPENAI_EMBEDDING_API_KEY"),
    "AZURE_OPENAI_EMBEDDING_API_VERSION" : os.getenv("AZURE_OPENAI_EMBEDDING_API_VERSION")
    }

CHROMA_PATH = os.path.join(os.getcwd(), "chroma_db")

st.set_page_config(page_title="PDF chatbot", layout="wide")
st.title("RAG Demo - PDF Q&A")

st.markdown("""
1. **Upload Your Documents**: The system accepts multiple PDF files at once, analyzing the content to provide comprehensive insights.

2. **Ask a Question**: After processing the documents, ask any question related to the content of your uploaded documents for a precise answer.
""")

def main():
    st.header("Ask a question")

    # Initialize vector store
    embeddings = AzureOpenAIEmbeddings(
        model="text-embedding-ada-002",
        api_key=azure_config["AZURE_OPENAI_API_KEY"],
        api_version="2024-02-01",
        azure_endpoint="https://ey-openai.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2023-05-15"
    )
    vectorstore = Chroma(
        persist_directory=CHROMA_PATH,
        embedding_function=embeddings
    )

    with st.sidebar:
        st.title("Menu:")
        uploaded_files = st.file_uploader(
            "Upload your PDF Files and Click on the Submit & Process Button", 
            accept_multiple_files=True, 
            key="pdf_uploader")
        if st.button("Submit & Process", key="process_button") and uploaded_files:  # Check if API key is provided before processing
            with st.spinner("Processing..."):
                for uploaded_file in uploaded_files:
                    try:
                        # Create a temporary file to store the uploaded PDF
                        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                            tmp_file.write(uploaded_file.getvalue())
                            tmp_file_path = tmp_file.name
                        
                        # Load PDF using PyPDFLoader
                        loader = PyPDFLoader(tmp_file_path)
                        pages = loader.load()
                        
                        # Add metadata to each page
                        for page in pages:
                            page.metadata["page_number"] = pages.index(page) + 1
                        
                        # Split into chunks
                        text_splitter = RecursiveCharacterTextSplitter(
                            chunk_size=1000,
                            chunk_overlap=200,
                            length_function=len,
                            separators=["\n\n", "\n", " ", ""]
                        )
                        chunks = text_splitter.split_documents(pages)
                        # Clean up temporary file
                        os.unlink(tmp_file_path)                        
                    
                        vectorstore.add_documents(chunks)
                        vectorstore.persist()
                    
                    except Exception as e:
                        st.error(f"Error processing {uploaded_file.name}: {str(e)}")
                        continue
                print("vector store created")
                st.success("Vector store created")

    user_question = st.text_input("Ask a Question from the PDF Files", key="user_question")
    if user_question:
        print("Question: ", user_question)

        retriever = vectorstore.as_retriever(search_type='similarity', search_kwargs={'k': 5})

        llm = AzureChatOpenAI(temperature=0,
                        api_key=azure_config["AZURE_OPENAI_API_KEY"],
                        openai_api_version=azure_config["AZURE_OPENAI_API_VERSION"],
                        azure_endpoint=azure_config["AZURE_OPENAI_ENDPOINT"],
                        model=azure_config["AZURE_OPENAI_MODEL_DEPLOYMENT_NAME"],
                        validate_base_url=False)

        chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)
        answer = chain.invoke(user_question)

        print(answer)
        st.write("Reply: ", answer['result'])


if __name__ == "__main__":
    main()