import streamlit as st
from backend.ocr import extract_text_from_pdf
from backend.rag_model import setup_retriever_and_qa, get_answer

st.set_page_config(page_title="StudyMate - PDF Q&A", layout="wide")
st.title("📚 StudyMate: AI-Powered PDF-Based Q&A System")

# Initialize retriever and RAG model
retriever, rag_pipeline = setup_retriever_and_qa()

uploaded_file = st.file_uploader("Upload your study material (PDF)", type="pdf")

if uploaded_file:
    with st.spinner("Extracting content from PDF..."):
        full_text = extract_text_from_pdf(uploaded_file)
        st.success("PDF content extracted!")

    if full_text:
        st.text_area("📄 Extracted Text Preview", full_text[:2000], height=300)

        query = st.text_input("💬 Ask a question based on the PDF")

        if query:
            with st.spinner("Thinking..."):
                answer = get_answer(full_text, query, retriever, rag_pipeline)
            st.markdown(f"**🧠 Answer:** {answer}")
    else:
        st.error("Failed to extract text from the PDF.")
else:
    st.info("Please upload a PDF to get started.")