import streamlit as st
from pypdf import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
import google.generativeai as genai

# Configure Gemini API
api_key = st.secrets["GOOGLE_API_KEY"]
genai.configure(api_key=api_key)

# Streamlit page config
st.set_page_config(page_title="PDF RAG Chatbot")

st.header("📄 Chat with PDF using Gemini AI")

# Extract text from PDFs
def get_pdf_text(pdf_docs):
    text = ""

    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)

        for page in pdf_reader.pages:
            extracted = page.extract_text()
            if extracted:
                text += extracted

    return text

# Split text into chunks
def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )

    chunks = text_splitter.split_text(text)

    return chunks

# Create vector store
def get_vector_store(text_chunks):
    embeddings = GoogleGenerativeAIEmbeddings(
        model="models/embedding-001",
        google_api_key=api_key
    )

    vector_store = FAISS.from_texts(
        text_chunks,
        embedding=embeddings
    )

    vector_store.save_local("faiss_index")

# Create QA chain
def get_conversational_chain():

    model = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        temperature=0.3,
        google_api_key=api_key
    )

    chain = load_qa_chain(
        model,
        chain_type="stuff"
    )

    return chain

# Handle user question
def user_input(user_question):

    embeddings = GoogleGenerativeAIEmbeddings(
        model="models/embedding-001",
        google_api_key=api_key
    )

    new_db = FAISS.load_local(
        "faiss_index",
        embeddings,
        allow_dangerous_deserialization=True
    )

    docs = new_db.similarity_search(user_question)

    chain = get_conversational_chain()

    response = chain(
        {
            "input_documents": docs,
            "question": user_question
        },
        return_only_outputs=True
    )

    st.subheader("🤖 Answer")
    st.write(response["output_text"])

# Sidebar for PDF upload
with st.sidebar:
    st.title("📚 Upload Your PDF Files")

    pdf_docs = st.file_uploader(
        "Upload PDF Files",
        accept_multiple_files=True
    )

    if st.button("Submit & Process"):

        with st.spinner("Processing PDFs..."):

            raw_text = get_pdf_text(pdf_docs)

            text_chunks = get_text_chunks(raw_text)

            get_vector_store(text_chunks)

            st.success("PDF Processing Complete!")

# User question input
user_question = st.text_input(
    "Ask a Question from the Uploaded PDF Files"
)

if user_question:
    user_input(user_question)