import gradio as gr
import os
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain import FAISS
from gradio_pdf import PDF  # Import PDF from gradio_pdf

# Function to process uploaded PDF and generate responses based on the document and user input
def chat_with_pdf(pdf_file, api_key, user_question):
    # Set the Google API key
    os.environ["GOOGLE_API_KEY"] = api_key
    
    # Load the document
    loader = PyPDFLoader(pdf_file.name)
    pages = loader.load_and_split()

    # Create a vector db index
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    db = FAISS.from_documents(pages, embeddings)
    
    # Search relevant docs based on user question
    docs = db.similarity_search(user_question)
    
    # Prepare the context for the API request
    content = "\n".join([x.page_content for x in docs])
    qa_prompt = "Use the following pieces of context to answer the user's question. If you don't know the answer, just say that you don't know, don't try to make up an answer.----------------"
    input_text = qa_prompt + "\nContext:" + content + "\nUser question:\n" + user_question
    
    # Call Gemini API (ChatGoogleGenerativeAI) to generate a response
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
    result = llm.invoke(input_text)
    
    # Return the bot's response (without chat history)
    return result.content

# Create a Gradio interface with a split layout
with gr.Blocks() as iface:
    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="Upload PDF")  # Upload PDF file
            pdf_display = PDF(label="PDF Preview")  # PDF preview using gradio_pdf
        with gr.Column(scale=1):
            response_output = gr.Textbox(label="Bot Response")  # Output for the bot response
            question_box = gr.Textbox(label="Ask a question", placeholder="Enter your question here")
            api_key_box = gr.Textbox(label="API Key", type="password", placeholder="Enter your Google API Key here")
    
    # Directly display the PDF once uploaded without using the 'upload' method
    pdf_input.change(lambda pdf_file: pdf_file.name, inputs=pdf_input, outputs=pdf_display)
    
    # When the user submits a question, process it and return the bot's response
    question_box.submit(
        chat_with_pdf, 
        inputs=[pdf_input, api_key_box, question_box],
        outputs=response_output
    )
    
# Launch the Gradio app
iface.launch()