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

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

st.title("📚 AI Reading Assistant")
st.write("Upload a PDF or paste text. Get simple explanations, summaries, and answers.")

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

generator = load_model()

# -----------------------------
# FILE UPLOAD
# -----------------------------
uploaded_file = st.file_uploader("Upload PDF", type=["pdf"])

text_data = ""

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

# -----------------------------
# TEXT INPUT
# -----------------------------
text_input = st.text_area("Or paste your text here:")

if text_input:
    text_data = text_input

# -----------------------------
# MAIN FUNCTIONALITY
# -----------------------------
if text_data:

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

    # -----------------------------
    # SIMPLIFY TEXT
    # -----------------------------
    if st.button("✨ Simplify Paragraph"):
        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 TEXT
    # -----------------------------
    if st.button("📝 Summarize Text"):
        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'])

    # -----------------------------
    # QUESTION ANSWERING
    # -----------------------------
    st.subheader("❓ Ask a 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'])

    # -----------------------------
    # WORD EXPLANATION
    # -----------------------------
    st.subheader("🔍 Word Explanation")
    word = st.text_input("Enter a difficult word:")

    if word:
        with st.spinner("Explaining..."):
            prompt = f"Explain the word '{word}' in simple English and give an example."
            response = generator(prompt, max_length=80)
            st.success(response[0]['generated_text'])

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

# -----------------------------
# FOOTER
# -----------------------------
st.markdown("---")
st.caption("Built with ❤️ using Streamlit + Hugging Face (FLAN-T5)")