src/streamlit_app.py

import os
os.environ["TRANSFORMERS_CACHE"] = "/cache"  # Special directory in Spaces

import streamlit as st
import fitz  # PyMuPDF
from transformers import pipeline

# Set page config
st.set_page_config(
    page_title="PrepPal",
    page_icon="📘",
    layout="wide",
    menu_items={
        'About': "PrepPal - AI-powered study assistant"
    }
)

# Load summarizer model (using a smaller, faster model)
@st.cache_resource
def load_summarizer():
    try:
        return pipeline(
            "summarization", 
            model="Falconsai/text_summarization",  # Smaller model
            device=-1  # Use CPU (more reliable in Spaces)
        )
    except Exception as e:
        st.error(f"❌ Failed to load model: {str(e)}")
        return None

# Extract text from PDF with size limit
def extract_text_from_pdf(uploaded_file):
    text = ""
    try:
        # Check file size (max 5MB)
        if uploaded_file.size > 5_000_000:
            st.error("File too large (max 5MB)")
            return ""
            
        with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:
            for page in doc:
                text += page.get_text()
        return text.strip()
    except Exception as e:
        st.error(f"❌ Error extracting text: {str(e)}")
        return ""

# Summarize text in chunks
def summarize_text(text, summarizer, max_chunk=1024):
    if not text:
        return ""
        
    try:
        chunks = [text[i:i+max_chunk] for i in range(0, len(text), max_chunk)]
        summary = ""
        for chunk in chunks:
            result = summarizer(
                chunk,
                max_length=150,
                min_length=30,
                do_sample=False
            )
            summary += result[0]['summary_text'] + " "
        return summary.strip()
    except Exception as e:
        st.error(f"❌ Summarization failed: {str(e)}")
        return ""

# Load model
summarizer = load_summarizer()

# UI Layout
st.title("📘 PrepPal - Study Assistant")
st.markdown("Upload your notes and get an AI-powered summary")

tab1, tab2 = st.tabs(["📄 Summarize Notes", "💬 Feedback"])

with tab1:
    st.header("PDF Summarizer")
    uploaded_file = st.file_uploader(
        "Choose a PDF file (max 5MB)",
        type=["pdf"],
        accept_multiple_files=False
    )

    if uploaded_file and summarizer:
        with st.spinner("Extracting text..."):
            text = extract_text_from_pdf(uploaded_file)

        if text:
            st.subheader("Extracted Text Preview")
            st.text_area("", text[:500] + "...", height=150, disabled=True)

            if st.button("Generate Summary"):
                with st.spinner("Summarizing..."):
                    summary = summarize_text(text, summarizer)
                
                if summary:
                    st.subheader("AI Summary")
                    st.write(summary)
                    
                    st.download_button(
                        "Download Summary",
                        data=summary,
                        file_name="summary.txt",
                        mime="text/plain"
                    )
                else:
                    st.warning("No summary generated")

with tab2:
    st.header("Feedback")
    st.write("We'd love to hear your thoughts!")
    feedback = st.text_area("Your feedback")
    if st.button("Submit Feedback"):
        st.success("Thank you! Your feedback has been recorded.")

# Add footer
st.markdown("---")
st.caption("PrepPal v1.0 | AI-powered study assistant")