import os
import gradio as gr
from langchain.document_loaders import PyPDFLoader
from langchain.llms import OpenAI
from langchain.chains import load_qa_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Setup OpenAI API key (replace with your OpenAI API key)
openai.api_key = os.getenv("OPENAI_API_KEY")
# Initialize GPT-4o-mini model
llm = OpenAI(model_name="gpt-4o-mini")

# Function to load PDF document
def load_pdf(file):
    loader = PyPDFLoader(file.name)
    documents = loader.load()
    return documents

# Summarization using the "stuff" method
def summarize_pdf(file):
    docs = load_pdf(file)

    # Define a prompt for summarization
    prompt = ChatPromptTemplate.from_messages([("system", "Write a concise summary of the following:\n\n{context}")])

    # Create summarization chain
    chain = create_stuff_documents_chain(llm, prompt)
    summary = chain.invoke({"context": docs})
    return summary['result']

# Question Answering on PDF
def qa_on_pdf(file, question):
    docs = load_pdf(file)

    # Create QA chain
    qa_chain = load_qa_chain(llm)
    response = qa_chain.run({"input_documents": docs, "question": question})
    return response

# Gradio interface
def create_interface():
    # Define the interface components
    with gr.Blocks() as interface:
        gr.Markdown("# Research Paper Summarization and QA")

        with gr.Row():
            with gr.Column():
                pdf_file = gr.File(label="Upload Research Paper PDF")
                summarize_btn = gr.Button("Summarize PDF")
                summary_output = gr.Textbox(label="Summary", lines=10)

                summarize_btn.click(summarize_pdf, inputs=[pdf_file], outputs=[summary_output])

            with gr.Column():
                qa_question = gr.Textbox(label="Ask a Question about the PDF", placeholder="Type your question here...")
                qa_btn = gr.Button("Submit Question")
                qa_output = gr.Textbox(label="Answer", lines=10)

                qa_btn.click(qa_on_pdf, inputs=[pdf_file, qa_question], outputs=[qa_output])

    return interface

# Launch Gradio app
interface = create_interface()
interface.launch()