import os
import tempfile
os.environ["TRANSFORMERS_CACHE"] = "/cache"

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

# App setup
st.set_page_config(
    page_title="PrepPal",
    page_icon="📘",
    layout="wide"
)

@st.cache_resource
def load_model():
    return pipeline(
        "summarization",
        model="Falconsai/text_summarization",
        device=-1  # CPU only
    )

def get_pdf_text(uploaded_file):
    with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp:
        tmp.write(uploaded_file.getbuffer())
        with fitz.open(tmp.name) as doc:
            return " ".join([page.get_text() for page in doc])

def create_summary(text, model):
    chunks = [text[i:i+800] for i in range(0, len(text), 800)]
    return " ".join([
        model(chunk, max_length=120, min_length=30)[0]['summary_text']
        for chunk in chunks
    ])

def main():
    st.title("📘 PrepPal - Study Assistant")
    
    # All three tabs
    tab1, tab2, tab3 = st.tabs(["📄 Summarize", "❓ Ask a Doubt", "💬 Feedback"])
    
    with tab1:
        st.header("PDF Summarizer")
        uploaded_file = st.file_uploader(
            "Upload PDF (max 5MB)",
            type=["pdf"]
        )
        
        if uploaded_file and uploaded_file.size <= 5_000_000:
            if st.button("Generate Summary"):
                with st.spinner("Extracting text..."):
                    text = get_pdf_text(uploaded_file)
                
                if text:  # Only proceed if text extraction succeeded
                    with st.spinner("Loading AI model..."):
                        model = load_model()
                    
                    with st.spinner("Generating summary..."):
                        summary = create_summary(text, model)
                    
                    st.subheader("Summary")
                    st.write(summary)
                else:
                    st.info("Could not extract text from PDF")
    
    with tab2:
        st.header("Ask a Question")
        st.write("Feature coming soon!")
    
    with tab3:
        st.header("Share Feedback")
        feedback = st.text_area("Your thoughts about PrepPal")
        if st.button("Submit Feedback"):
            st.success("Thank you for your feedback!")

if __name__ == "__main__":
    main()