import streamlit as st
import os
from utils import generate_mcq, load_docs, load_llama, downlode_hugging_face_embeddings, split_docs, create_vectorstore
import tempfile
from langchain_chroma import Chroma


st.set_page_config(page_title="GenAI MCQ Generator", layout="wide")
st.title("GenAI based MCQ Quiz App")

# session states initialization
if "questions" not in st.session_state:
    st.session_state.questions = []

# load the default retriever on app start up
if "retriever" not in st.session_state:
    # check if persistance folder exists
    if os.path.exists('chroma_db_MCQGen'):
        embeddings = downlode_hugging_face_embeddings()
        vectorstore = Chroma(
            persist_directory='chroma_db_MCQGen', 
            embedding_function=embeddings
        )
        st.session_state.retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
    else:
        st.session_state.retriever = None

#side bar for document upload and processing
with st.sidebar:
    st.header("Upload Medical Document")
    uploaded_file = st.file_uploader("Upload a medical book or document to generate the questions", type=["pdf"])

    if uploaded_file:
        if st.button("Process Document"):
            with st.spinner("Analysing your document. This may take a few minutes..."):
                # tempory file path to store the uploded pdf
                with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                    tmp_file.write(uploaded_file.getvalue())
                    tmp_path = tmp_file.name

                # load and split the document from app.py functions
                loader = load_docs(tmp_path)
                chunks=split_docs(loader)
                embeddings=downlode_hugging_face_embeddings()

                vectorstore = create_vectorstore(chunks, embeddings, persist_dir='chroma_db_MCQGen')
                st.session_state.retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

                os.remove(tmp_path) #Clean up
                st.success("Analysis Complete")


# document upload block(Stage-1)
topic = st.text_input("Enter Topic:")
if st.button("Generate MCQs"):
    # Check if a document has been processed first
    if st.session_state.retriever is not None: 
        with st.spinner("Generating..."):
            # CHANGE 'retriever' TO 'st.session_state.retriever'
            st.session_state.questions = generate_mcq(topic, st.session_state.retriever)
        st.rerun()
    else:
        st.error("Please upload and process a document in the sidebar first!")

# quiz block(Stage-2)
if st.session_state.questions:
    with st.form("quiz_form", border=False):
        for i, mcq in enumerate(st.session_state.questions):
            st.write(f"## Question {i+1}")
            st.write(f"Question: {mcq['question']}")

            user_choice = st.radio(label="Select an option", options=mcq['option'], index=None, key=f"radio_{i}")

            if st.session_state.get(f"radio_{i}") and st.session_state.get("FormSubmitter:quiz_form-Submit"):
                if user_choice.strip().lower() == mcq['correct_answer'].strip().lower(): 
                    st.success("Correct Answer!", icon="✅")
                else:
                    st.error("Incorrect!", icon="❌")

            with st.expander(label="Check Answer"):
                st.write(f"Correct Answer: {mcq['correct_answer']}")

        st.write("---")
        submit_btn = st.form_submit_button("Submit", type="primary")

    if submit_btn:
        st.success("Quiz Submitted!")
        
else:
    st.info("The quiz will appear here once you enter a topic and click 'Generate MCQs'.")