import streamlit as st
from transformers import pipeline
import PyPDF2

# Load Hugging Face Models
@st.cache_resource
def load_models():
    question_generator = pipeline("text2text-generation", model="valhalla/t5-small-qg-prepend")
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    return question_generator, summarizer

question_generator, summarizer = load_models()

# Extract text from a PDF file
def extract_text_from_pdf(pdf_file):
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    extracted_text = ""
    for page in pdf_reader.pages:
        extracted_text += page.extract_text()
    return extracted_text

# Summarize text
def summarize_text(text):
    summary = summarizer(text, max_length=130, min_length=30, do_sample=False)
    return summary[0]["summary_text"]

# Generate MCQs
def generate_questions(text, num_questions=5):
    clean_text = text[:500]  # Limit input to first 500 characters
    prompt = (
        f"Create {num_questions} diverse multiple-choice questions focusing on key concepts. "
        f"Use the text: {clean_text}"
    )
    questions = question_generator(
        prompt,
        max_length=128,
        num_beams=5,
        num_return_sequences=min(num_questions, 3),  # Generate up to 3 per request
    )
    return [q["generated_text"] for q in questions]

# Streamlit App
st.title("AI-Powered Quiz Generator")
st.subheader("Create quizzes and summaries from any text or PDF file using AI.")

# Input options
st.write("### Step 1: Provide Input")
input_option = st.radio("Choose input method:", ("Text Input", "Upload PDF"))
input_text = ""

if input_option == "Text Input":
    input_text = st.text_area("Enter your text:", height=200)
elif input_option == "Upload PDF":
    pdf_file = st.file_uploader("Upload a PDF file:", type=["pdf"])
    if pdf_file:
        input_text = extract_text_from_pdf(pdf_file)
        st.write("Extracted Text Preview:")
        st.write(input_text[:1000])  # Show a preview of the extracted text

# Generate quiz and summary
if st.button("Generate Quiz"):
    if input_text.strip():
        # Summarize content
        st.write("### Summary")
        with st.spinner("Summarizing content..."):
            summary = summarize_text(input_text)
            st.write(summary)

        # Generate MCQs
        st.write("### Generated Questions")
        with st.spinner("Generating questions..."):
            questions = generate_questions(summary, num_questions=5)
            for i, question in enumerate(questions):
                st.write(f"**{i+1}. {question}**")
    else:
        st.error("Please enter some text or upload a PDF to generate a quiz!")

# Footer
st.markdown("---")
st.markdown("**Created by Engr Muhammad Arsalan**")