import streamlit as st
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
import fitz
import json

# Title
st.title("📄 PDF-based MCQ Generator")

# Sidebar
st.sidebar.title("Upload & Settings")

# Upload PDF
pdf_file = st.sidebar.file_uploader("Upload a PDF file", type=["pdf"])

# Number of questions
number_of_questions = st.sidebar.slider("Number of questions", min_value=1, max_value=20, value=5)

# Session states
if "mcqs" not in st.session_state:
    st.session_state.mcqs = []
if "current_q" not in st.session_state:
    st.session_state.current_q = 0
if "user_answers" not in st.session_state:
    st.session_state.user_answers = {}
if "quiz_finished" not in st.session_state:
    st.session_state.quiz_finished = False

# Gemini setup
GOOGLE_API_KEY = "AIzaSyCB5NLx39vOAlfRQBDmnEG3uLBgLraGvH4"
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=GOOGLE_API_KEY,
    temperature=0.7
)

template = """
You are an expert MCQ generator. Generate {number} unique multiple-choice questions from the given text.
Each question must have exactly 1 correct answer and 3 incorrect options.
Strictly return output in the following JSON format (no explanations, no markdown):

[
  {{
    "question": "What is ...?",
    "options": ["Option A", "Option B", "Option C", "Option D"],
    "answer": "Option D"
  }},
  ...
]

TEXT:
{text}
"""

prompt = PromptTemplate(
    input_variables=["text", "number"],
    template=template
)

mcq_chain = LLMChain(llm=llm, prompt=prompt)

# PDF text extractor
def extract_text_from_pdf(pdf):
    doc = fitz.open(stream=pdf.read(), filetype="pdf")
    full_text = ""
    for page in doc:
        full_text += page.get_text()
    doc.close()
    return full_text

# Generate MCQs
if st.sidebar.button("Generate MCQs"):
    if pdf_file is None:
        st.error("Please upload a PDF file.")
    else:
        with st.spinner("Extracting text and generating MCQs..."):
            text = extract_text_from_pdf(pdf_file)
            try:
                response = mcq_chain.run(text=text, number=str(number_of_questions))
                # st.subheader("🔍 Raw Output (Debugging)")
                # st.code(response)
                mcqs_json = json.loads(response.strip())
                st.session_state.mcqs = mcqs_json
                st.session_state.current_q = 0
                st.session_state.user_answers = {}
                st.session_state.quiz_finished = False
                st.success("✅ MCQs generated successfully!")
            except Exception as e:
                st.error(f"Error generating MCQs: {e}")

# Display question
if st.session_state.mcqs and not st.session_state.quiz_finished:
    idx = st.session_state.current_q
    q_data = st.session_state.mcqs[idx]

    st.subheader(f"Question {idx + 1}: {q_data['question']}")
    selected_option = st.radio(
        "Choose an answer:",
        q_data["options"],
        key=f"radio_{idx}"
    )

    if st.button("Next"):
        st.session_state.user_answers[idx] = selected_option

        if st.session_state.current_q < len(st.session_state.mcqs) - 1:
            st.session_state.current_q += 1
        else:
            st.session_state.quiz_finished = True
            st.success("🎉 Quiz completed!")

# Show result
if st.session_state.quiz_finished:
    st.header("📊 Quiz Results")
    score = 0
    total = len(st.session_state.mcqs)

    for i, q in enumerate(st.session_state.mcqs):
        user_ans = st.session_state.user_answers.get(i)
        correct_ans = q["answer"]
        if user_ans == correct_ans:
            score += 1
        st.markdown(f"**Q{i+1}: {q['question']}**")
        st.markdown(f"- Your answer: {user_ans}")
        st.markdown(f"- Correct answer: {correct_ans}")
        st.markdown("---")

    st.success(f"✅ You scored {score} out of {total}")