import streamlit as st
import PyPDF2
from groq import Groq

def extract_text_from_pdf(pdf_file):
    """
    Extracts text from an uploaded PDF file and returns it as a single string.
    """
    text = ""
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

def ask_question_from_pdf(pdf_text, user_question, api_key):
    """
    Uses the Groq model to answer a question based on the provided PDF text.
    """
    client = Groq(api_key=api_key)  # Initialize Groq client with the API key.

    # Truncate the PDF text if it exceeds the model's input limit.
    max_input_tokens = 2048  # Model input token limit.
    truncated_pdf_text = pdf_text[:max_input_tokens]

    # Stream the response for the user's question.
    stream = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant. Use the provided PDF content to answer the user's question."
            },
            {
                "role": "user",
                "content": f"PDF content:\n{truncated_pdf_text}\n\nQuestion: {user_question}"
            }
        ],
        model="llama-3.3-70b-versatile",
        temperature=0.5,
        max_completion_tokens=512,
        top_p=1,
        stream=True,
    )

    # Collect the streamed response.
    answer = ""
    for chunk in stream:
        if chunk.choices[0].delta.content:  # Ensure content is not None
            answer += chunk.choices[0].delta.content
    return answer

# Streamlit app
st.title("PDF Question Answering with Groq")
st.write("Upload a PDF file and ask questions about its content.")

# Input for Groq API key
api_key = st.text_input("Enter your Groq API key", type="password")

# File uploader
uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")

if uploaded_file is not None:
    pdf_text = extract_text_from_pdf(uploaded_file)
    st.write("PDF content successfully extracted!")

    # Question input
    user_question = st.text_input("Ask a question about the PDF content")

    if st.button("Get Answer") and user_question and api_key:
        with st.spinner("Fetching the answer..."):
            try:
                answer = ask_question_from_pdf(pdf_text, user_question, api_key)
                st.success("Answer fetched!")
                st.write("**Answer:**", answer)
            except Exception as e:
                st.error(f"An error occurred: {e}")
    elif not api_key:
        st.error("Please provide a valid Groq API key.")
    elif not user_question:
        st.error("Please enter a question.")

st.write("Note: The Groq API key is required for the app to function.")