import streamlit as st
import os
from model import load_vectorstore, ask_question

st.set_page_config(page_title="Simple RAG Q&A", layout="centered")
st.title("RAG Q&A with Mistral AI")
st.write("Upload a PDF and ask questions about its content.")

# PDF upload
uploaded_file = st.file_uploader("Upload PDF", type=["pdf"])
pdf_path = "/app/data/document.pdf"

if uploaded_file:
    os.makedirs("/app/data", exist_ok=True)
    with open(pdf_path, "wb") as f:
        f.write(uploaded_file.read())
    st.success("PDF uploaded!")
    
    with st.spinner("Indexing document..."):
        try:
            load_vectorstore(pdf_path)
            st.success("Document indexed!")
        except Exception as e:
            st.error(f"Indexing failed: {str(e)}")

# Query input
query = st.text_input("Enter your question", 
                      "How many articles are there in the Selenium webdriver python course?")
if st.button("Ask") and query:
    if not os.path.exists(pdf_path):
        st.error("Please upload a PDF first.")
    else:
        with st.spinner("Generating answer..."):
            try:
                result = ask_question(query, pdf_path)
                st.subheader("Answer")
                st.write(result["answer"])
                
                st.subheader("Retrieved Contexts")
                for i, context in enumerate(result["contexts"], 1):
                    with st.expander(f"Context {i}"):
                        st.write(context)
            except Exception as e:
                st.error(f"Failed to generate answer: {str(e)}")