import gradio as gr
from PyPDF2 import PdfReader
import os
import openai

# Set OpenAI key
openai.api_key = os.getenv("OPENAI_API_KEY")

pdf_text = ""

def extract_text_from_pdf(pdf_file):
    reader = PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() or ""
    return text

def process_pdf(pdf):
    global pdf_text
    pdf_text = extract_text_from_pdf(pdf)
    return "PDF loaded! Ask anything about it."

def chat_with_pdf(question):
    if not pdf_text:
        return "Please upload and process a PDF first."
    
    prompt = f"""You are a helpful assistant. The user uploaded a PDF document. Here's its content:

--- BEGIN DOCUMENT ---
{pdf_text}
--- END DOCUMENT ---

Now, answer the following question based on the document:
Q: {question}
A:"""

    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",  # or "gpt-4"
        messages=[
            {"role": "system", "content": "You are a helpful assistant that answers questions about uploaded PDFs."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=500,
        temperature=0.3,
    )

    return response.choices[0].message["content"]

with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Chat with your PDF (No Chunking, No Embeddings)")

    with gr.Row():
        pdf_file = gr.File(label="Upload your PDF", file_types=[".pdf"])
        load_button = gr.Button("Load PDF")

    status = gr.Textbox(label="Status")

    with gr.Row():
        question = gr.Textbox(label="Your Question")
        answer = gr.Textbox(label="Answer", lines=10)
        ask_button = gr.Button("Ask")

    load_button.click(process_pdf, inputs=pdf_file, outputs=status)
    ask_button.click(chat_with_pdf, inputs=question, outputs=answer)

demo.launch()