import streamlit as st
import random
import time
from transformers import pipeline

# Load question generation pipeline
qg_pipeline = pipeline("text2text-generation", model="iarfmoose/t5-base-question-generator")

st.set_page_config(page_title="AI MCQ Quiz", layout="centered")
st.title("🧠 Auto-Generated AI Quiz")

# Step 1: User input
input_text = st.text_area("Enter a topic or paragraph:")

if st.button("Generate Quiz"):
    if not input_text.strip():
        st.warning("Please enter some text first.")
    else:
        with st.spinner("Generating quiz questions..."):
            result = qg_pipeline(f"generate questions: {input_text}", max_length=512, do_sample=True)
            raw_questions = result[0]['generated_text'].split("\n")
            quiz_data = []

            for q in raw_questions:
                if q.strip() and "?" in q:
                    question = q.strip()
                    correct = "Not provided"  # Placeholder (you can use a separate model for answer extraction)
                    options = [correct]
                    # Generate 3 fake options (distractors)
                    fake_answers = [
                        "Option A", "Option B", "Option C", "Option D", "None of these",
                        "Random guess", "All of the above", "Incorrect answer"
                    ]
                    while len(options) < 4:
                        option = random.choice(fake_answers)
                        if option not in options:
                            options.append(option)
                    random.shuffle(options)

                    quiz_data.append({
                        "question": question,
                        "options": options,
                        "answer": correct  # For now this is unknown
                    })

        if quiz_data:
            st.session_state.quiz = quiz_data
            st.session_state.current_q = 0
            st.session_state.score = 0
            st.session_state.start_time = time.time()
            st.experimental_rerun()

# Proceed to quiz if it's generated
if "quiz" in st.session_state and st.session_state.current_q < len(st.session_state.quiz):
    quiz = st.session_state.quiz
    q_index = st.session_state.current_q
    q_data = quiz[q_index]

    st.markdown(f"### ❓ Question {q_index + 1}:")
    st.markdown(f"**{q_data['question']}**")

    time_left = max(0, 15 - int(time.time() - st.session_state.start_time))
    st.info(f"⏱️ Time left: {time_left} seconds")

    if time_left == 0:
        st.warning("⏰ Time's up! Moving to next question...")
        st.session_state.current_q += 1
        st.session_state.start_time = time.time()
        st.experimental_rerun()

    selected = st.radio("Your Answer:", q_data["options"], index=None, key=q_index)

    if st.button("Submit Answer"):
        if selected:
            if selected == q_data["answer"]:
                st.success("✅ Correct!")
                st.session_state.score += 1
            else:
                st.error(f"❌ Incorrect! Correct: {q_data['answer']}")
            time.sleep(1.5)
            st.session_state.current_q += 1
            st.session_state.start_time = time.time()
            st.experimental_rerun()
        else:
            st.warning("Please choose an option.")

elif "quiz" in st.session_state and st.session_state.current_q >= len(st.session_state.quiz):
    st.success(f"🎉 Quiz completed! Score: {st.session_state.score}/{len(st.session_state.quiz)}")
    if st.button("Restart"):
        for key in ["quiz", "current_q", "score", "start_time"]:
            st.session_state.pop(key, None)
        st.experimental_rerun()