import json
from typing import List, Any, Dict, Tuple, Union
import gradio as gr
from huggingface_hub import InferenceClient

MAX_QUESTIONS = 20

def parse_quiz_json(content: str) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
    """Parse the quiz JSON content robustly."""
    try:
        if content.startswith("{") or content.startswith("["):
            quiz = json.loads(content)
        else:
            quiz = json.loads(content.lstrip().removeprefix("```json").removesuffix("```"))
        if isinstance(quiz, dict) and "questions" in quiz:
            return quiz["questions"]
        return quiz
    except Exception:
        try:
            quiz = json.loads(json.loads(content))
            if isinstance(quiz, dict) and "questions" in quiz:
                return quiz["questions"]
            return quiz
        except Exception as e:
            return {"error": f"Could not parse JSON: {e}", "raw_output": content}

def generate_quiz(
    topic: str, num_questions: int, q_type: str
) -> Tuple[List[dict[str, Any]], List[str]]:
    """Generate quiz questions and correct answers."""
    if not topic:
        updates = [
            gr.update(visible=True, value="Please enter a topic."),
            *[gr.update(visible=False, value="", interactive=False) for _ in range((MAX_QUESTIONS * 2) - 1)]
        ]
        return updates, []

    system_msg = "You are an expert quiz creator."
    user_msg = (
        f"Create a quiz containing {num_questions} {q_type} questions based on the following topic:\n"
        f"{topic}\n"
        "Respond with a JSON array of objects, each containing a question, options, and the answer.\n"
    )
    if q_type == "Multiple Choice":
        user_msg += "Each question should have 4 answer choices, unlabeled, and one correct answer."

    client = InferenceClient()
    content = ""
    # Stream the response from the model
    for chunk in client.chat.completions.create(
        model="openai/gpt-oss-20b",
        messages=[
            {"role": "system", "content": system_msg},
            {"role": "user", "content": user_msg}
        ],
        max_tokens=1500,
        temperature=0.7,
        stream=True,
        # response_format={
        #     "type": "json_schema",
        #     "json_schema": {
        #         "name": "questions",
        #         "description": "A list of questions, answer choices, and the correct answer",
        #         "schema": {
        #             "type": "array",
        #             "items": {
        #                 "type": "object",
        #                 "properties": {
        #                     "question": {"type": "string"},
        #                     "options": {
        #                         "type": "array",
        #                         "items": {"type": "string"}
        #                     },
        #                     "answer": {"type": "string"}
        #                 },
        #             }
        #         },
        #         "strict": True
        #     }
        # },
    ):
        if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
            content += chunk.choices[0].delta.content

    quiz = parse_quiz_json(content)
    updates = []
    correct_answers = []

    for i in range(MAX_QUESTIONS):
        if isinstance(quiz, dict) and "error" in quiz:
            if i == 0:
                updates.append(gr.update(visible=True, value="Oops! An error occured. Please try again."))
            else:
                updates.append(gr.update(visible=False, value="", interactive=False))
            updates.append(gr.update(visible=False, choices=[], value=None, interactive=False))
        elif isinstance(quiz, list) and i < len(quiz):
            question = quiz[i].get("question", "No question provided")
            options = ["True", "False"] if q_type == "True/False" else quiz[i].get("options", [])
            answer = quiz[i].get("answer", "No answer provided")
            correct_answers.append(answer)
            updates.append(gr.update(visible=True, value=f"{question}"))
            updates.append(gr.update(visible=True, choices=options, interactive=True, value=None))
        else:
            updates.append(gr.update(visible=False, value="", interactive=False))
            updates.append(gr.update(visible=False, choices=[], value=None, interactive=False))
    return updates, correct_answers

def score_quiz(*args) -> str:
    """Score the quiz based on user answers and correct answers."""
    *user_answers, correct_answers = args
    if not correct_answers or not isinstance(correct_answers, list):
        return "No quiz generated yet. Please generate a quiz first."
    score = 0
    total = 0
    for i, user_ans in enumerate(user_answers):
        if i < len(correct_answers) and user_ans is not None:
            if user_ans == correct_answers[i]:
                score += 1
            total += 1
    return f"Your score: {score} / {total}"

with gr.Blocks() as demo:
    gr.Markdown("# AI Quiz Generator")
    gr.Markdown("### Instructions")
    gr.Markdown(
        """1. Enter a topic for the quiz in the textbox above.
2. Select the number of questions you want.
3. Choose the type of questions (Multiple Choice or True/False).
4. Click 'Generate Quiz' to create your quiz.
5. The generated quiz will appear in the output box.
6. After answering, click 'Submit Answers' to see your score."""
    )
    topic_input = gr.Textbox(label="Quiz Topic / Prompt", lines=4, placeholder="Enter the topic for the quiz")
    num_questions_input = gr.Slider(minimum=1, maximum=MAX_QUESTIONS, step=1, label="Number of Questions", value=5)
    q_type_input = gr.Radio(choices=["Multiple Choice", "True/False"], label="Question Type", value="Multiple Choice")
    generate_button = gr.Button("Generate Quiz")
    quiz_components = []
    for i in range(MAX_QUESTIONS):
        quiz_components.append(gr.Textbox(label=f"Q{i+1}", visible=False, interactive=False))
        quiz_components.append(gr.Radio(choices=[], label=f"Options", visible=False))
    score_output = gr.Textbox(label="Score", visible=True, interactive=False)
    submit_button = gr.Button("Submit Answers")
    correct_answers_state = gr.State([])

    def generate_and_store(*args):
        updates, correct_answers = generate_quiz(*args)
        return (*updates, correct_answers, 0)

    generate_button.click(
        fn=generate_and_store,
        inputs=[topic_input, num_questions_input, q_type_input],
        outputs=[*quiz_components, correct_answers_state, score_output],
    )

    radio_outputs = [quiz_components[i] for i in range(1, len(quiz_components), 2)]

    def score_with_state(*args):
        *user_answers, correct_answers = args
        return score_quiz(*user_answers, correct_answers)

    submit_button.click(
        fn=score_with_state,
        inputs=[*radio_outputs, correct_answers_state],
        outputs=[score_output],
    )

    with gr.Row():
        generate_button
        submit_button
    score_output

if __name__ == "__main__":
    demo.launch()