app.py

import streamlit as st
import uuid
from datetime import datetime
from models import NoteSession, NoteSettings, Quiz
from text_processor import TextProcessor
from ai_handler import AIHandler
from config import MAX_TEXT_LENGTH, MAX_FILE_SIZE_MB

st.set_page_config(
    page_title="NoteAI - AI Note Summarizer",
    page_icon="📝",
    layout="wide"
)

if 'current_session' not in st.session_state:
    st.session_state.current_session = None
if 'ai_handler' not in st.session_state:
    try:
        st.session_state.ai_handler = AIHandler()
    except ValueError as e:
        st.error(f"⚠️ {str(e)}")
        st.info("Please set up your Google Gemini API key in a .env file. Get one FREE at https://makersuite.google.com/app/apikey")
        st.stop()

text_processor = TextProcessor()

def create_session(text: str, settings: NoteSettings) -> NoteSession:
    session = NoteSession(
        id=str(uuid.uuid4()),
        timestamp=datetime.now(),
        original_text=text,
        settings=settings
    )
    return session


def process_notes(session: NoteSession):
    with st.spinner("🧹 Cleaning text..."):
        cleaned_text = text_processor.clean_text(session.original_text)
    
    with st.spinner("✂️ Chunking text..."):
        chunks = text_processor.chunk_text(cleaned_text)
        st.info(f"Processing {len(chunks)} chunk(s)...")
    
    with st.spinner("📝 Generating summary..."):
        try:
            summary = st.session_state.ai_handler.generate_summary_for_chunks(
                chunks,
                session.settings.summary_length,
                session.settings.output_format,
                session.settings.tone
            )
            session.summary = summary
        except Exception as e:
            st.error(f"Summary generation failed: {str(e)}")
            return False
    
    if session.settings.quiz_enabled:
        with st.spinner("❓ Generating quiz..."):
            try:
                quiz_data = st.session_state.ai_handler.generate_quiz(
                    cleaned_text,
                    session.settings.quiz_difficulty
                )
                session.quiz = Quiz(
                    multiple_choice=quiz_data.get("multiple_choice", []),
                    short_answer=quiz_data.get("short_answer", [])
                )
                st.success(f"✅ Quiz generated: {len(session.quiz.multiple_choice)} MC, {len(session.quiz.short_answer)} SA")
            except Exception as e:
                st.error(f"Quiz generation failed: {str(e)}")
                session.quiz = None
    else:
        session.quiz = None
    
    return True


st.title("NoteAI")
st.caption("AI-powered note summarizer")

col1, col2 = st.columns([1, 1])

with col1:
    st.header("Input")
    
    input_method = st.radio(
        "Choose input method:",
        ["Paste Text", "Upload File"],
        horizontal=True
    )
    
    input_text = ""
    
    if input_method == "Paste Text":
        input_text = st.text_area(
            "Paste your notes here:",
            height=300,
            placeholder="Enter or paste your notes here...",
            help=f"Maximum {MAX_TEXT_LENGTH} characters"
        )
        
        if input_text:
            word_count = text_processor.count_words(input_text)
            st.caption(f" {len(input_text)} characters | {word_count} words")
    
    else:
        uploaded_file = st.file_uploader(
            "Upload a text or PDF file:",
            type=["txt", "pdf"],
            help=f"Maximum file size: {MAX_FILE_SIZE_MB}MB"
        )
        
        if uploaded_file:
            file_size_mb = len(uploaded_file.getvalue()) / (1024 * 1024)
            
            if file_size_mb > MAX_FILE_SIZE_MB:
                st.error(f"❌ File size ({file_size_mb:.1f}MB) exceeds maximum of {MAX_FILE_SIZE_MB}MB")
            else:
                try:
                    file_type = uploaded_file.name.split('.')[-1].lower()
                    input_text = text_processor.extract_text_from_file(
                        uploaded_file.getvalue(),
                        file_type
                    )
                    
                    word_count = text_processor.count_words(input_text)
                    st.success(f"✅ File loaded: {word_count} words")
                
                except Exception as e:
                    st.error(f"❌ Error reading file: {str(e)}")
    
    if st.button(" Clear Input"):
        st.rerun()

with col2:
    st.header("Settings")
    
    summary_length = st.select_slider(
        "Summary Length:",
        options=["short", "medium", "long"],
        value="medium"
    )
    
    output_format = st.selectbox(
        "Output Format:",
        ["bullets", "outline", "flashcards"],
        index=0
    )
    
    tone = st.selectbox(
        "Tone:",
        ["simple", "academic"],
        index=0,
        help="Simple: Easy to understand | Academic: Formal and precise"
    )
    
    st.divider()
    
    quiz_enabled = st.checkbox("Generate Quiz", value=False)
    
    quiz_difficulty = "medium"
    if quiz_enabled:
        quiz_difficulty = st.selectbox(
            "Difficulty:",
            ["easy", "medium", "hard"],
            index=1
        )

st.divider()
generate_col1, generate_col2, generate_col3 = st.columns([1, 2, 1])

with generate_col2:
    generate_button = st.button(
        " Generate Summary",
        type="primary",
        use_container_width=True
    )

if generate_button:
    is_valid, error_msg = text_processor.validate_input(input_text, MAX_TEXT_LENGTH)
    
    if not is_valid:
        st.error(f"❌ {error_msg}")
    else:
        settings = NoteSettings(
            summary_length=summary_length,
            output_format=output_format,
            tone=tone,
            quiz_enabled=quiz_enabled,
            quiz_difficulty=quiz_difficulty
        )
        
        session = create_session(input_text, settings)
        
        success = process_notes(session)
        
        if success:
            st.session_state.current_session = session
            st.success("✅ Processing complete!")
            st.rerun()

if st.session_state.current_session:
    st.divider()
    
    session = st.session_state.current_session
    
    st.header("Summary")
    if session.summary:
        st.markdown(session.summary)
        
        from reportlab.lib.pagesizes import letter
        from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
        from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
        from reportlab.lib.units import inch
        from io import BytesIO
        
        pdf_buffer = BytesIO()
        doc = SimpleDocTemplate(pdf_buffer, pagesize=letter, topMargin=0.75*inch, bottomMargin=0.75*inch)
        styles = getSampleStyleSheet()
        story = []
        
        title_style = ParagraphStyle('CustomTitle', parent=styles['Heading1'], fontSize=18, spaceAfter=12)
        
        story.append(Paragraph("NoteAI Summary", title_style))
        story.append(Paragraph(f"Generated: {session.timestamp.strftime('%Y-%m-%d %H:%M:%S')}", styles['Normal']))
        story.append(Spacer(1, 0.3*inch))
        
        for line in session.summary.split('\n'):
            if line.strip():
                story.append(Paragraph(line, styles['Normal']))
        
        doc.build(story)
        pdf_data = pdf_buffer.getvalue()
        pdf_buffer.close()
        
        st.download_button(
            label="Export Summary as PDF",
            data=pdf_data,
            file_name=f"noteai_summary_{session.id[:8]}.pdf",
            mime="application/pdf"
        )
    else:
        st.info("No summary available")
    
    st.divider()
    
    if session.settings.quiz_enabled:
        st.header("Quiz")
        
        if session.quiz and hasattr(session.quiz, 'multiple_choice') and session.quiz.multiple_choice:
            for i, q in enumerate(session.quiz.multiple_choice, 1):
                st.markdown(f"**{i}. {q['question']}**")
                for key, value in q.get('options', {}).items():
                    st.markdown(f"   {key}. {value}")
                
                with st.expander(f"Show answer"):
                    st.success(f"Answer: {q.get('correct_answer', 'N/A')}")
                st.markdown("")
        
            if hasattr(session.quiz, 'short_answer') and session.quiz.short_answer:
                st.markdown("**Short Answer Questions**")
                for i, q in enumerate(session.quiz.short_answer, 1):
                    st.markdown(f"**{i}. {q['question']}**")
                    with st.expander(f"Show sample answer"):
                        st.info(q.get('sample_answer', 'N/A'))
                    st.markdown("")
            
            from reportlab.lib.pagesizes import letter
            from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
            from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
            from reportlab.lib.units import inch
            from io import BytesIO
            
            pdf_buffer = BytesIO()
            doc = SimpleDocTemplate(pdf_buffer, pagesize=letter, topMargin=0.75*inch, bottomMargin=0.75*inch)
            styles = getSampleStyleSheet()
            story = []
            
            title_style = ParagraphStyle('CustomTitle', parent=styles['Heading1'], fontSize=18, spaceAfter=12)
            
            story.append(Paragraph("NoteAI Quiz", title_style))
            story.append(Paragraph(f"Generated: {session.timestamp.strftime('%Y-%m-%d %H:%M:%S')}", styles['Normal']))
            story.append(Spacer(1, 0.3*inch))
            
            for i, q in enumerate(session.quiz.multiple_choice, 1):
                story.append(Paragraph(f"<b>{i}. {q['question']}</b>", styles['Normal']))
                for key, value in q.get('options', {}).items():
                    story.append(Paragraph(f"&nbsp;&nbsp;&nbsp;{key}. {value}", styles['Normal']))
                story.append(Paragraph(f"&nbsp;&nbsp;&nbsp;Answer: {q.get('correct_answer', 'N/A')}", styles['Normal']))
                story.append(Spacer(1, 0.1*inch))
            
            doc.build(story)
            pdf_data = pdf_buffer.getvalue()
            pdf_buffer.close()
            
            st.download_button(
                label="Export Quiz as PDF",
                data=pdf_data,
                file_name=f"noteai_quiz_{session.id[:8]}.pdf",
                mime="application/pdf"
            )
            
            st.divider()
            
            pdf_buffer_both = BytesIO()
            doc_both = SimpleDocTemplate(pdf_buffer_both, pagesize=letter, topMargin=0.75*inch, bottomMargin=0.75*inch)
            story_both = []
            
            story_both.append(Paragraph("NoteAI Summary", title_style))
            story_both.append(Paragraph(f"Generated: {session.timestamp.strftime('%Y-%m-%d %H:%M:%S')}", styles['Normal']))
            story_both.append(Spacer(1, 0.3*inch))
            
            for line in session.summary.split('\n'):
                if line.strip():
                    story_both.append(Paragraph(line, styles['Normal']))
            
            story_both.append(Spacer(1, 0.3*inch))
            story_both.append(Paragraph("Quiz", title_style))
            
            for i, q in enumerate(session.quiz.multiple_choice, 1):
                story_both.append(Paragraph(f"<b>{i}. {q['question']}</b>", styles['Normal']))
                for key, value in q.get('options', {}).items():
                    story_both.append(Paragraph(f"&nbsp;&nbsp;&nbsp;{key}. {value}", styles['Normal']))
                story_both.append(Paragraph(f"&nbsp;&nbsp;&nbsp;Answer: {q.get('correct_answer', 'N/A')}", styles['Normal']))
                story_both.append(Spacer(1, 0.1*inch))
            
            doc_both.build(story_both)
            pdf_data_both = pdf_buffer_both.getvalue()
            pdf_buffer_both.close()
            
            st.download_button(
                label="Export Summary + Quiz as PDF",
                data=pdf_data_both,
                file_name=f"noteai_complete_{session.id[:8]}.pdf",
                mime="application/pdf"
            )
        elif not session.quiz:
            st.warning("Quiz generation failed or is in progress...")
        else:
            st.info("No quiz questions were generated.")