import streamlit as st
from pypdf import PdfReader
from transformers import pipeline

# -----------------------------
# PAGE CONFIG
# -----------------------------
st.set_page_config(page_title="AI Study Companion", layout="wide")

st.title("🎓 AI Study Companion")
st.write("Your all-in-one AI tool for reading, understanding, and practicing.")

# -----------------------------
# LOAD MODEL
# -----------------------------
@st.cache_resource
def load_model():
    return pipeline("text2text-generation", model="google/flan-t5-base")

generator = load_model()

# -----------------------------
# INPUT SECTION
# -----------------------------
uploaded_file = st.file_uploader("📄 Upload PDF", type=["pdf"])
text_input = st.text_area("✍️ Or paste your text here:")

text_data = ""

if uploaded_file:
    reader = PdfReader(uploaded_file)
    for page in reader.pages:
        if page.extract_text():
            text_data += page.extract_text()

if text_input:
    text_data = text_input

# -----------------------------
# MAIN APP
# -----------------------------
if text_data:

    st.subheader("📖 Text Preview")
    st.write(text_data[:1500])

    # -----------------------------
    # FEATURE SELECTION
    # -----------------------------
    option = st.selectbox(
        "Choose what you want to do:",
        ["Simplify Text", "Summarize", "Ask Question", "Generate Quiz"]
    )

    # -----------------------------
    # SIMPLIFY
    # -----------------------------
    if option == "Simplify Text":
        if st.button("✨ Simplify"):
            with st.spinner("Simplifying..."):
                prompt = f"Explain this in very simple English:\n{text_data[:500]}"
                response = generator(prompt, max_length=150)
                st.success(response[0]['generated_text'])

    # -----------------------------
    # SUMMARIZE
    # -----------------------------
    elif option == "Summarize":
        if st.button("📝 Summarize"):
            with st.spinner("Summarizing..."):
                prompt = f"Summarize this text:\n{text_data[:500]}"
                response = generator(prompt, max_length=120)
                st.success(response[0]['generated_text'])

    # -----------------------------
    # Q&A
    # -----------------------------
    elif option == "Ask Question":
        question = st.text_input("Enter your question:")

        if question:
            with st.spinner("Thinking..."):
                prompt = f"Answer the question based on the text below:\n\nText:\n{text_data[:700]}\n\nQuestion:\n{question}"
                response = generator(prompt, max_length=120)
                st.success(response[0]['generated_text'])

    # -----------------------------
    # QUIZ GENERATOR
    # -----------------------------
    elif option == "Generate Quiz":
        if st.button("🧠 Generate Quiz"):
            with st.spinner("Creating quiz..."):
                prompt = f"Create 5 multiple choice questions (MCQs) with answers from this text:\n{text_data[:500]}"
                response = generator(prompt, max_length=300)
                st.success(response[0]['generated_text'])

else:
    st.info("Upload a PDF or paste text to start.")

# -----------------------------
# FOOTER
# -----------------------------
st.markdown("---")
st.caption("🚀 Built with Streamlit + Hugging Face | AI Study Companion")