import os
import json
import re
import gradio as gr
from groq import Groq

# Initialize Groq Client
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# -------------------------------
# Helper: Clean JSON Response
# -------------------------------
def extract_json(text):
    try:
        # Remove markdown formatting if present
        text = re.sub(r"```json|```", "", text).strip()
        return json.loads(text)
    except:
        return None

# -------------------------------
# Generate Quiz from Groq
# -------------------------------
def generate_quiz(topic):
    prompt = f"""
    Generate 5 MCQs for undergraduate students on topic: {topic}.
    Return ONLY JSON in this format:

    [
        {{
            "question": "Question text",
            "options": ["Option1", "Option2", "Option3", "Option4"],
            "answer": "Correct option exactly as written above"
        }}
    ]
    """

    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.3-70b-versatile",
    )

    raw_output = chat_completion.choices[0].message.content
    questions = extract_json(raw_output)

    return questions

# -------------------------------
# Hint Generator
# -------------------------------
def get_hint(question):
    hint_prompt = f"Give a short hint for this MCQ: {question}"

    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": hint_prompt}],
        model="llama-3.3-70b-versatile",
    )

    return chat_completion.choices[0].message.content

# -------------------------------
# Explanation Generator
# -------------------------------
def get_explanation(question, answer):
    explanation_prompt = f"Explain briefly why '{answer}' is correct for: {question}"

    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": explanation_prompt}],
        model="llama-3.3-70b-versatile",
    )

    return chat_completion.choices[0].message.content

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as app:

    gr.Markdown("# 🎯 Mini AI Quiz App")
    gr.Markdown("### Enter topic → Generate MCQs → Get instant feedback 💡🔥")

    topic_input = gr.Textbox(label="📚 Enter Topic")
    generate_btn = gr.Button("🔥 Generate Quiz")

    question_display = gr.Markdown()
    options_radio = gr.Radio(choices=[], label="Choose your answer 👇", interactive=True)

    submit_btn = gr.Button("Submit Answer ✅")
    feedback_output = gr.Markdown()
    next_btn = gr.Button("Next Question ➡️")

    question_state = gr.State([])
    index_state = gr.State(0)
    score_state = gr.State(0)

    # Start Quiz
    def start_quiz(topic):
        questions = generate_quiz(topic)

        if not questions:
            return "⚠️ Failed to generate quiz. Try again.", [], [], 0, 0

        first_q = questions[0]

        return (
            f"### Q1: {first_q['question']}",
            gr.update(choices=first_q["options"], value=None),
            questions,
            0,
            0
        )

    generate_btn.click(
        start_quiz,
        inputs=topic_input,
        outputs=[question_display, options_radio, question_state, index_state, score_state]
    )

    # Submit Answer
    def check_answer(selected_option, questions, index, score):

        if selected_option is None:
            return "⚠️ Please select an option first!", score

        correct_answer = questions[index]["answer"]
        question_text = questions[index]["question"]

        if selected_option == correct_answer:
            score += 1
            explanation = get_explanation(question_text, correct_answer)
            feedback = f"✅ Correct! 🔥\n\n📘 {explanation}"
        else:
            hint = get_hint(question_text)
            explanation = get_explanation(question_text, correct_answer)
            feedback = f"❌ Incorrect 😭\n\n💡 Hint: {hint}\n\n📘 Explanation: {explanation}"

        return feedback, score

    submit_btn.click(
        check_answer,
        inputs=[options_radio, question_state, index_state, score_state],
        outputs=[feedback_output, score_state]
    )

    # Next Question
    def next_question(questions, index, score):

        index += 1

        if index >= len(questions):
            return (
                f"🎉 Quiz Completed!\n\n🏆 Final Score: {score}/{len(questions)}",
                gr.update(choices=[], value=None),
                index
            )

        next_q = questions[index]

        return (
            f"### Q{index+1}: {next_q['question']}",
            gr.update(choices=next_q["options"], value=None),
            index
        )

    next_btn.click(
        next_question,
        inputs=[question_state, index_state, score_state],
        outputs=[question_display, options_radio, index_state]
    )

app.launch()