import gradio as gr
import PyPDF2
import openai
from config import OPENAI_API_KEY
import os
openai.api_key = os.getenv("OPENAI_API_KEY")

def extract_text_from_pdf(pdf_file):
    text = ""
    with open(pdf_file.name, "rb") as file:
        reader = PyPDF2.PdfReader(file)
        for page in reader.pages:
            text += page.extract_text() + "\n"
    return text

def answer_question(pdf_file, question):
    # Extract text from the PDF
    text = extract_text_from_pdf(pdf_file)
    
    # Define the assistant's behavior
    assistant_prompt = f"""
    You are a helpful assistant that answers questions based on the content of the PDF provided.
    Here is the content of the PDF:
    {text}

    User question: {question}
    """

    # Call OpenAI API to get the answer using GPT-4 Turbo
    response = openai.ChatCompletion.create(
        model="gpt-4-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": assistant_prompt}
        ]
    )
    
    answer = response.choices[0].message['content']
    return answer

# Create Gradio interface using the updated input/output classes
iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.File(label="Upload PDF"),
        gr.Textbox(label="Ask a question about the PDF", placeholder="What do you want to know?")
    ],
    outputs="text",
    title="PDF Q&A with OpenAI Assistant",
    description="Upload a PDF document and ask questions about its content. The assistant will provide answers based on the PDF.",
    examples=[
        ["renesas-ra6m1-group-datasheet.pdf", "Which Renesas products are mentioned in this PDF?"]
    ]
)

# Launch the interface
iface.launch()