import streamlit as st
import os
import json
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
pdf_path = "/app/data/document.pdf"
uploaded_file = st.file_uploader("Upload PDF", type=["pdf"])

if uploaded_file:
    os.makedirs("/app/data", exist_ok=True)
    try:
        with open(pdf_path, "wb") as f:
            f.write(uploaded_file.read())
        st.success("PDF uploaded!")
        
        with st.spinner("Indexing document..."):
            load_vectorstore(pdf_path)
            st.success("Document indexed!")
    except Exception as e:
        st.error(f"Failed to upload/index PDF: {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)}")

# Query endpoint for testing
if "query" in st.experimental_get_query_params():
    query = st.experimental_get_query_params().get("query", [""])[0]
    if query and os.path.exists(pdf_path):
        try:
            result = ask_question(query, pdf_path)
            st.json(result)
        except Exception as e:
            st.json({"error": str(e)})