import streamlit as st
from PyPDF2 import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate


st.set_page_config(page_title="PDF CHATBOT", layout="wide")

st.markdown("""
## Document Genie: Get instant insights from your Documents

This chatbot is built using the Retrieval-Augmented Generation (RAG) framework, leveraging Google's Generative AI model Gemini-PRO. It processes uploaded PDF documents by breaking them down into manageable chunks, creates a searchable vector store, and generates accurate answers to user queries. This advanced approach ensures high-quality, contextually relevant responses for an efficient and effective user experience.

### How It Works

Follow these simple steps to interact with the chatbot:

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 get_pdf(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text
    
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

def response_generate(text,query):
    text_splitter = RecursiveCharacterTextSplitter(
    # Set a really small chunk size, just to show.
    chunk_size=500,
    chunk_overlap=20,
    separators=["\n\n","\n"," ",".",","])
    chunks=text_splitter.split_text(text)
    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    db = Chroma.from_documents(chunks, embeddings)
    # Create retriever interface
    retriever = db.as_retriever()
    qa = RetrievalQA.from_chain_type(llm = GoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_API_KEY ), chain_type='stuff', retriever=retriever)
    return qa.run(query_text)

def main():
    st.header("Chat with your pdf💁")

    query = st.text_input("Ask a Question from the PDF Files", key="query")

    #if query: 
    #  user_call(query)

    st.title("Menu:")
    pdf_docs = 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"):
        with st.spinner("Processing..."):
            raw_text = get_pdf(pdf_docs)
            #text_chunks = text_splitter(raw_text)
            response = response_generate(raw_text,query)
            st.success("Done")
            st.write("Reply: ", response)

if __name__ == "__main__":
    main()