# ai_quiz_from_pdf_gradio.py
# Gradio app: AI Quiz Generation System from PDF Documents using LLM (Groq: meta-llama/llama-4-maverick-17b-128e-instruct)
#
# Prerequisites:
#   - pip install -r requirements.txt
#   - Set environment variable: GROQ_API_KEY=<your_api_key>
#
# Run:
#   - python ai_quiz_from_pdf_gradio.py
#
# Notes:
#   - The app uses PyMuPDF to extract text from PDFs.
#   - The LLM is used to generate quizzes (MCQ/TrueFalse/ShortAnswer/Essay) and to grade Essay answers.
#   - Tabs are used for a clean interface: "Upload & Generate", "Take Quiz", "Grade Essays", "Export".

import os
import io
import json
import time
from typing import List, Dict, Any, Tuple, Optional
import pandas as pd

import fitz  # PyMuPDF
import gradio as gr

# --- LLM (Groq) ---
from groq import Groq

GROQ_MODEL = "meta-llama/llama-4-maverick-17b-128e-instruct"

def call_groq(system_prompt: str, user_prompt: str, temperature: float = 0.7, max_tokens: int = 2048) -> str:
    """
    Calls the Groq chat completion API and returns the full combined response text.
    Streams tokens and accumulates them into a final string.
    """
    client = Groq(api_key = "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")  # Uses GROQ_API_KEY env var
    completion = client.chat.completions.create (
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=temperature,
        max_completion_tokens=max_tokens,
        top_p=1,
        stream=True,
        stop=None
    )
    chunks = []
    for chunk in completion:
        piece = chunk.choices[0].delta.content or ""
        chunks.append(piece)
    return "".join(chunks)

# -------------------- PDF Utilities --------------------

def extract_text_from_pdf(pdf_bytes: bytes, max_pages: int = 50) -> str:
    """
    Extracts text from a PDF byte stream using PyMuPDF.
    Limits to max_pages for performance.
    """
    text_parts = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        page_count = min(len(doc), max_pages)
        for i in range(page_count):
            page = doc.load_page(i)
            text_parts.append(page.get_text("text"))
    full_text = "\n".join(text_parts)
    # Basic cleanup
    lines = [ln.strip() for ln in full_text.splitlines() if ln.strip()]
    return "\n".join(lines)

def chunk_text(text: str, max_chars: int = 4000, overlap: int = 200) -> List[str]:
    """
    Splits text into overlapping chunks to keep prompts within context limits.
    """
    chunks = []
    start = 0
    n = len(text)
    while start < n:
        end = min(start + max_chars, n)
        chunk = text[start:end]
        chunks.append(chunk)
        if end == n:
            break
        start = end - overlap
        if start < 0:
            start = 0
    return chunks

# -------------------- Quiz Generation --------------------

QUIZ_SCHEMA_EXAMPLE = {
    "metadata": {
        "source": "string (e.g., PDF title or filename)",
        "difficulty": "Easy|Medium|Hard"
    },
    "questions": [
        # MCQ
        {"id": 1, "type": "mcq", "question": "Q text", "options": ["A", "B", "C", "D"], "answer": "B", "explanation": "why"},
        # True/False
        {"id": 2, "type": "true_false", "question": "Statement", "answer": True, "explanation": "why"},
        # Short Answer (open, brief)
        {"id": 3, "type": "short_answer", "question": "Q text", "expected_answer": "reference answer", "rubric": "key points"},
        # Essay (theory)
        {"id": 4, "type": "essay", "question": "Essay prompt", "rubric": "grading rubric/bullets", "max_score": 10}
    ]
}

SYSTEM_PROMPT_GENERATE = """You are an expert educational content creator.
You will receive textbook or lecture content and must produce a structured quiz in strict JSON.
Output ONLY valid JSON that follows the given schema. Do not add commentary.
JSON schema keys: metadata{source, difficulty}, questions[list].
Question item types: "mcq", "true_false", "short_answer", "essay".
For "mcq": include 4 distinct options, exactly one correct "answer", and a brief "explanation".
For "true_false": "answer" must be true or false (boolean), include "explanation".
For "short_answer": include "expected_answer" and a concise "rubric" with key points.
For "essay": include a clear "rubric" and "max_score" as an integer.
Make questions faithful to the source text (no hallucinations). Prefer concise, unambiguous wording.
"""

def build_generation_prompt(chunk: str, source_name: str, difficulty: str,
                             n_mcq: int, n_tf: int, n_short: int, n_essay: int) -> str:
    return f"""
Source name: {source_name}
Desired difficulty: {difficulty}
Required counts -> MCQ: {n_mcq}, True/False: {n_tf}, Short Answer: {n_short}, Essay: {n_essay}

Source text:
\"\"\"
{chunk}
\"\"\"

Produce a single JSON object matching this schema (example):
{json.dumps(QUIZ_SCHEMA_EXAMPLE, indent=2)}

Ensure total number of questions equals the requested counts combined.
Use simple language appropriate for undergraduates.
"""

def merge_quizzes(quizzes: List[Dict[str, Any]], source_name: str, difficulty: str) -> Dict[str, Any]:
    """
    Merge multiple chunk-level quiz JSONs into one. Re-index question IDs.
    """
    merged = {"metadata": {"source": source_name, "difficulty": difficulty}, "questions": []}
    qid = 1
    for q in quizzes:
        for item in q.get("questions", []):
            item["id"] = qid
            merged["questions"].append(item)
            qid += 1
    return merged

def generate_quiz_from_text(text: str, source_name: str, difficulty: str,
                            n_mcq: int, n_tf: int, n_short: int, n_essay: int,
                            temperature: float = 0.7) -> Tuple[Dict[str, Any], str]:
    """
    Generates a quiz by calling the LLM on one or more chunks, distributing question counts across chunks.
    Returns (quiz_json, raw_model_output_debug).
    """
    chunks = chunk_text(text, max_chars=3500, overlap=200)
    total_required = n_mcq + n_tf + n_short + n_essay
    if total_required == 0:
        return {"metadata": {"source": source_name, "difficulty": difficulty}, "questions": []}, ""

    # Divide counts across chunks (simple even split)
    c = max(1, len(chunks))
    split = lambda total: [total // c + (1 if i < (total % c) else 0) for i in range(c)]
    mcq_split = split(n_mcq)
    tf_split = split(n_tf)
    short_split = split(n_short)
    essay_split = split(n_essay)

    partial_quizzes = []
    debug_texts = []
    for i, ch in enumerate(chunks):
        if mcq_split[i] + tf_split[i] + short_split[i] + essay_split[i] == 0:
            continue
        user_prompt = build_generation_prompt(
            chunk=ch,
            source_name=source_name,
            difficulty=difficulty,
            n_mcq=mcq_split[i],
            n_tf=tf_split[i],
            n_short=short_split[i],
            n_essay=essay_split[i]
        )
        raw = call_groq(SYSTEM_PROMPT_GENERATE, user_prompt, temperature=temperature, max_tokens=2048)
        debug_texts.append(raw)
        # attempt to parse JSON (robustly strip fences, trailing text)
        cleaned = raw.strip()
        if cleaned.startswith("```"):
            cleaned = cleaned.strip("`")
            # if it included language tag like ```json
            if cleaned.startswith("json"):
                cleaned = cleaned[len("json"):].strip()
        # Find first and last braces to extract JSON
        start = cleaned.find("{")
        end = cleaned.rfind("}")
        if start != -1 and end != -1 and end > start:
            cleaned = cleaned[start:end+1]
        try:
            obj = json.loads(cleaned)
            partial_quizzes.append(obj)
        except Exception as e:
            partial_quizzes.append({"metadata":{"source":source_name, "difficulty":difficulty},"questions":[]})

    merged = merge_quizzes(partial_quizzes, source_name, difficulty)
    return merged, "\n\n---\n\n".join(debug_texts)

# -------------------- Grading --------------------

SYSTEM_PROMPT_GRADE_ESSAY = """You are a strict but fair examiner.
Grade the student's essay against the rubric and max_score provided.
Return ONLY a JSON object with keys: score (integer), feedback (string).
Be concise and specific in feedback. Do not add commentary outside JSON.
"""

def build_grading_prompt(question_text: str, rubric: str, max_score: int, student_answer: str) -> str:
    return f"""
Question:
{question_text}

Rubric (key points to award marks):
{rubric}

Max Score: {max_score}

Student Answer:
\"\"\"
{student_answer}
\"\"\"

Return JSON only: {{"score": <int>, "feedback": "<text>"}}
"""

def grade_essay_answer(question_text: str, rubric: str, max_score: int, student_answer: str, temperature: float = 0.2) -> Dict[str, Any]:
    raw = call_groq(SYSTEM_PROMPT_GRADE_ESSAY, build_grading_prompt(question_text, rubric, max_score, student_answer), temperature=temperature, max_tokens=512)
    cleaned = raw.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.startswith("json"):
            cleaned = cleaned[len("json"):].strip()
    start = cleaned.find("{"); end = cleaned.rfind("}")
    if start != -1 and end != -1 and end > start:
        cleaned = cleaned[start:end+1]
    try:
        obj = json.loads(cleaned)
    except Exception:
        obj = {"score": 0, "feedback": "Grading failed to parse. Please retry."}
    # clamp score
    try:
        obj["score"] = int(obj.get("score", 0))
        if obj["score"] < 0: obj["score"] = 0
        if obj["score"] > max_score: obj["score"] = max_score
    except Exception:
        obj["score"] = 0
    return obj

# -------------------- Gradio App --------------------

with gr.Blocks(title="AI Quiz Generation from PDF (LLM)") as demo:
    gr.Markdown("# AI Quiz Generation System from PDF Documents using LLM")
    gr.Markdown("Upload a PDF, generate MCQ/True/Short/Essay questions, take the quiz, and grade essay answers via LLM.\n")

    state_pdf_text = gr.State("")
    state_quiz = gr.State({"metadata": {}, "questions": []})
    state_source_name = gr.State("")
    state_debug = gr.State("")
        # Extra states for step-by-step quiz flow
    state_non_essay = gr.State([])                 # holds MCQ/TF/Short questions
    state_take_idx = gr.State(0)                   # index for Take Quiz flow
    state_take_results = gr.State({"attempted": 0, "correct": 0, "details": []})

    state_essay_list = gr.State([])                # holds Essay questions
    state_essay_idx = gr.State(0)                  # index for Essay flow
    state_essay_results = gr.State([])             # list of {"id","score","max","feedback"}


    with gr.Tabs():
        with gr.Tab("1) Upload & Generate"):
            with gr.Row():
                pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
                source_name_tb = gr.Textbox(label="Source Name (optional)", placeholder="e.g., Intro_to_AI_Notes.pdf")
            with gr.Accordion("Extraction & Generation Settings", open=True):
                with gr.Row():
                    difficulty_dd = gr.Dropdown(choices=["Easy","Medium","Hard"], value="Medium", label="Difficulty")
                    temp_slider = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Creativity (temperature)")
                with gr.Row():
                    n_mcq = gr.Number(value=5, precision=0, label="Number of MCQ")
                    n_tf = gr.Number(value=3, precision=0, label="Number of True/False")
                    n_short = gr.Number(value=2, precision=0, label="Number of Short Answer")
                    n_essay = gr.Number(value=1, precision=0, label="Number of Essay")
            extract_btn = gr.Button("Extract Text")
            extracted_preview = gr.Textbox(label="Extracted Text Preview (first 2000 chars)", lines=12)
            generate_btn = gr.Button("Generate Quiz")
            quiz_json_out = gr.JSON(label="Generated Quiz (JSON)")
            debug_out = gr.Textbox(label="Raw Model Output (debug)", visible=False)

            def do_extract(pdf_file, source_name):
                if not pdf_file:
                    return gr.update(value=""), "", "No file provided."
                with open(pdf_file.name, "rb") as f:
                    pdf_bytes = f.read()
                text = extract_text_from_pdf(pdf_bytes, max_pages=50)
                preview = text[:2000]
                src = source_name or os.path.basename(pdf_file.name)
                return preview, text, src

            extract_btn.click(
                do_extract,
                inputs=[pdf_input, source_name_tb],
                outputs=[extracted_preview, state_pdf_text, state_source_name]
            )

            def do_generate(text, source_name, difficulty, temp, n1, n2, n3, n4):
                if not text or not text.strip():
                    return {"metadata": {}, "questions": []}, "No text to generate from.", {"metadata": {}, "questions": []}
                quiz, dbg = generate_quiz_from_text(
                    text=text,
                    source_name=source_name or "Uploaded PDF",
                    difficulty=difficulty,
                    n_mcq=int(n1 or 0),
                    n_tf=int(n2 or 0),
                    n_short=int(n3 or 0),
                    n_essay=int(n4 or 0),
                    temperature=float(temp)
                )
                return quiz, dbg, quiz

            generate_btn.click(
                do_generate,
                inputs=[state_pdf_text, state_source_name, difficulty_dd, temp_slider, n_mcq, n_tf, n_short, n_essay],
                outputs=[quiz_json_out, debug_out, state_quiz]
            )

        with gr.Tab("2) Take Quiz"):
            gr.Markdown("Answer one question at a time. MCQ and True/False use radio buttons; Short Answer uses a textbox. A summary appears at the end.")

            take_progress = gr.Markdown("")
            take_q_text = gr.Textbox(label="Question", lines=4, interactive=False)

            take_mcq_radio = gr.Radio(choices=[], label="Choose one (MCQ)", visible=False)
            take_tf_radio = gr.Radio(choices=["True", "False"], label="True/False", visible=False)
            take_short_tb = gr.Textbox(label="Short Answer", visible=False, placeholder="Type your short answer here...")

            take_submit_btn = gr.Button("Submit & Next")
            take_feedback = gr.JSON(label="Feedback (this question)", visible=False)
            take_summary = gr.JSON(label="Final Summary", visible=False)

            def _split_non_essay(quiz: Dict[str, Any]):
                return [q for q in quiz.get("questions", []) if q.get("type") in ["mcq", "true_false", "short_answer"]]

            def _progress_str(idx: int, total: int) -> str:
                return f"**Question {min(idx+1, total)} of {total}**" if total else "**No questions available.**"

            def _show_non_essay_question(non_essay: List[Dict[str, Any]], idx: int):
                total = len(non_essay)
                if total == 0:
                    # nothing to show
                    return (
                        _progress_str(0, 0),
                        gr.update(value="No non-essay questions in this quiz.", interactive=False),
                        gr.update(visible=False),  # mcq
                        gr.update(visible=False),  # tf
                        gr.update(visible=False),  # short
                        gr.update(visible=False),  # feedback
                        gr.update(value={"message": "Nothing to attempt."}, visible=True)  # summary
                    )
                if idx >= total:
                    # finished
                    return (
                        _progress_str(total, total),
                        gr.update(value="You have completed all non-essay questions.", interactive=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=True)
                    )

                q = non_essay[idx]
                qtype = q.get("type")
                qtext = q.get("question", "")

                if qtype == "mcq":
                    opts = q.get("options", []) or []
                    return (
                        _progress_str(idx, total),
                        gr.update(value=qtext, interactive=False),
                        gr.update(choices=opts, value=None, visible=True),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                    )
                elif qtype == "true_false":
                    return (
                        _progress_str(idx, total),
                        gr.update(value=qtext, interactive=False),
                        gr.update(visible=False),
                        gr.update(value=None, visible=True),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                    )
                else:  # short_answer
                    return (
                        _progress_str(idx, total),
                        gr.update(value=qtext, interactive=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(value="", visible=True),
                        gr.update(visible=False),
                        gr.update(visible=False),
                    )

            def init_take_quiz(quiz: Dict[str, Any]):
                non_essay = _split_non_essay(quiz)
                idx0 = 0
                results0 = {"attempted": 0, "correct": 0, "details": []}
                ui = _show_non_essay_question(non_essay, idx0)
                # return UI + states
                return (*ui, non_essay, idx0, results0)

            def _compare_short_answer(gold: str, ans: str) -> bool:
                gold = (gold or "").strip().lower()
                ans = (ans or "").strip().lower()
                if not gold or not ans:
                    return False
                gold_tokens = set(gold.split())
                ans_tokens = set(ans.split())
                # require at least 25% token overlap (min 1)
                common = gold_tokens & ans_tokens
                return len(common) >= max(1, len(gold_tokens) // 4)

            def submit_next(mcq_choice, tf_choice, short_ans,
                            non_essay: List[Dict[str, Any]],
                            idx: int,
                            results_state: Dict[str, Any]):
                total = len(non_essay)
                if total == 0 or idx >= total:
                    # already done
                    return (
                        _progress_str(total, total),
                        gr.update(value="You have completed all non-essay questions.", interactive=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(
                            value={"summary": results_state},
                            visible=True
                        ),
                        idx,
                        results_state,
                    )

                q = non_essay[idx]
                qid = q.get("id")
                qtype = q.get("type")
                explanation = q.get("explanation", "")
                feedback = {}
                is_correct = False
                your_answer = None

                if qtype == "mcq":
                    your_answer = mcq_choice
                    is_correct = (str(your_answer).strip() == str(q.get("answer", "")).strip())
                elif qtype == "true_false":
                    your_answer = tf_choice
                    gold = str(q.get("answer", "")).strip().lower()
                    val = str(your_answer or "").strip().lower()
                    # normalize radio labels
                    if val in ["t", "true", "yes", "1"]:
                        val = "true"
                    elif val in ["f", "false", "no", "0"]:
                        val = "false"
                    is_correct = (val == gold)
                else:  # short_answer
                    your_answer = short_ans
                    gold = str(q.get("expected_answer", "")).strip()
                    is_correct = _compare_short_answer(gold, your_answer)

                # update results
                results_state = dict(results_state)  # copy
                results_state["attempted"] += 1
                if is_correct:
                    results_state["correct"] += 1
                results_state["details"] = list(results_state.get("details", [])) + [{
                    "id": qid,
                    "type": qtype,
                    "your_answer": your_answer,
                    "correct_answer": q.get("answer", q.get("expected_answer", "")),
                    "correct": is_correct
                }]

                feedback = {
                    "id": qid,
                    "type": qtype,
                    "your_answer": your_answer,
                    "correct_answer": q.get("answer", q.get("expected_answer", "")),
                    "correct": is_correct,
                }
                if explanation and qtype in ["mcq", "true_false"]:
                    feedback["explanation"] = explanation

                # advance
                idx_next = idx + 1

                # if finished, show summary
                if idx_next >= total:
                    return (
                        _progress_str(total, total),
                        gr.update(value="Completed non-essay questions.", interactive=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(value=feedback, visible=True),
                        gr.update(value={
                            "summary": {
                                "attempted": results_state["attempted"],
                                "correct": results_state["correct"],
                                "accuracy": round(100.0 * results_state["correct"] / max(1, results_state["attempted"]), 2)
                            },
                            "details": results_state["details"]
                        }, visible=True),
                        idx_next,
                        results_state,
                    )

                # otherwise show next question
                ui_next = _show_non_essay_question(non_essay, idx_next)
                return (
                    *ui_next[:5],
                    gr.update(value=feedback, visible=True),  # feedback visible
                    gr.update(visible=False),                # summary hidden until end
                    idx_next,
                    results_state,
                )

            # reinitialize when a new quiz is generated
            state_quiz.change(
                init_take_quiz,
                inputs=[state_quiz],
                outputs=[
                    take_progress, take_q_text,
                    take_mcq_radio, take_tf_radio, take_short_tb,
                    take_feedback, take_summary,
                    state_non_essay, state_take_idx, state_take_results
                ]
            )

            # submit current answer and move next
            take_submit_btn.click(
                submit_next,
                inputs=[
                    take_mcq_radio, take_tf_radio, take_short_tb,
                    state_non_essay, state_take_idx, state_take_results
                ],
                outputs=[
                    take_progress, take_q_text,
                    take_mcq_radio, take_tf_radio, take_short_tb,
                    take_feedback, take_summary,
                    state_take_idx, state_take_results
                ]
            )

        with gr.Tab("3) Essay/Theory Grading (LLM)"):
            gr.Markdown("Essay questions are presented one at a time. Submit your answer to get LLM grading, then proceed to the next.")

            essay_progress = gr.Markdown("")
            essay_q_view = gr.Textbox(label="Essay Question", lines=4, interactive=False)

            with gr.Accordion("Rubric & Max Score", open=False):
                essay_rubric_view = gr.Textbox(label="Rubric", lines=6, interactive=False)
                essay_max_score = gr.Number(label="Max Score", value=10, precision=0, interactive=False)

            essay_answer_tb = gr.Textbox(label="Your Essay/Theory Answer", lines=10, placeholder="Write your answer here...")
            essay_submit_btn = gr.Button("Submit & Grade with LLM")

            essay_grade_json = gr.JSON(label="Grade (this essay)", visible=False)
            essay_final_summary = gr.JSON(label="Essay Summary", visible=False)

            def _essays_only(quiz: Dict[str, Any]):
                return [q for q in quiz.get("questions", []) if q.get("type") == "essay"]

            def _essay_progress_str(idx: int, total: int) -> str:
                return f"**Essay {min(idx+1, total)} of {total}**" if total else "**No essay questions available.**"

            def _show_essay(essays: List[Dict[str, Any]], idx: int):
                total = len(essays)
                if total == 0:
                    return (
                        _essay_progress_str(0, 0),
                        gr.update(value="No essay questions in this quiz.", interactive=False),
                        gr.update(value="", interactive=False),
                        gr.update(value=10, interactive=False),
                        gr.update(value="", visible=False),
                        gr.update(visible=False),
                        gr.update(value={"message": "Nothing to grade."}, visible=True)
                    )
                if idx >= total:
                    return (
                        _essay_progress_str(total, total),
                        gr.update(value="All essay questions completed.", interactive=False),
                        gr.update(value="", interactive=False),
                        gr.update(value=10, interactive=False),
                        gr.update(value="", visible=False),
                        gr.update(visible=False),
                        gr.update(visible=True)
                    )

                q = essays[idx]
                return (
                    _essay_progress_str(idx, total),
                    gr.update(value=q.get("question", ""), interactive=False),
                    gr.update(value=q.get("rubric", ""), interactive=False),
                    gr.update(value=int(q.get("max_score", 10)), interactive=False),
                    gr.update(value="", visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                )

            def init_essay_flow(quiz: Dict[str, Any]):
                essays = _essays_only(quiz)
                idx0 = 0
                results0 = []
                ui = _show_essay(essays, idx0)
                return (*ui, essays, idx0, results0)

            def grade_and_next_essay(answer: str,
                                     essays: List[Dict[str, Any]],
                                     idx: int,
                                     results: List[Dict[str, Any]]):
                total = len(essays)
                if total == 0 or idx >= total:
                    # already done
                    return (
                        _essay_progress_str(total, total),
                        gr.update(value="All essay questions completed.", interactive=False),
                        gr.update(value="", interactive=False),
                        gr.update(value=10, interactive=False),
                        gr.update(value="", visible=False),
                        gr.update(visible=False),
                        gr.update(value={"message": "Completed."}, visible=True),
                        essays, idx, results
                    )

                q = essays[idx]
                qtext = q.get("question", "")
                rubric = q.get("rubric", "")
                max_score = int(q.get("max_score", 10))

                if not str(answer or "").strip():
                    # ask to provide an answer
                    return (
                        _essay_progress_str(idx, total),
                        gr.update(value=qtext, interactive=False),
                        gr.update(value=rubric, interactive=False),
                        gr.update(value=max_score, interactive=False),
                        gr.update(value=answer or "", visible=True),
                        gr.update(value={"error": "Provide an answer before submitting."}, visible=True),
                        gr.update(visible=False),
                        essays, idx, results
                    )

                # grade with LLM
                grade_obj = grade_essay_answer(qtext, rubric, max_score, answer)

                # store result
                results = list(results) + [{
                    "id": q.get("id"),
                    "score": grade_obj.get("score", 0),
                    "max": max_score,
                    "feedback": grade_obj.get("feedback", "")
                }]

                idx_next = idx + 1

                # if finished, compute summary
                if idx_next >= total:
                    total_score = sum(r["score"] for r in results)
                    max_total = sum(r["max"] for r in results)
                    summary = {
                        "completed": len(results),
                        "total_score": total_score,
                        "max_total": max_total,
                        "percentage": round(100.0 * total_score / max(1, max_total), 2),
                        "details": results
                    }
                    return (
                        _essay_progress_str(total, total),
                        gr.update(value="All essay questions completed.", interactive=False),
                        gr.update(value=rubric, interactive=False),
                        gr.update(value=max_score, interactive=False),
                        gr.update(value="", visible=False),
                        gr.update(value=grade_obj, visible=True),
                        gr.update(value=summary, visible=True),
                        essays, idx_next, results
                    )

                # otherwise move to next essay
                ui_next = _show_essay(essays, idx_next)
                return (
                    *ui_next[:5],
                    gr.update(value=grade_obj, visible=True),  # show grade for current
                    gr.update(visible=False),                  # final summary hidden
                    essays, idx_next, results
                )

            # initialize the essay flow on new quiz
            state_quiz.change(
                init_essay_flow,
                inputs=[state_quiz],
                outputs=[
                    essay_progress, essay_q_view,
                    essay_rubric_view, essay_max_score,
                    essay_answer_tb, essay_grade_json, essay_final_summary,
                    state_essay_list, state_essay_idx, state_essay_results
                ]
            )

            # submit current essay answer and move to next
            essay_submit_btn.click(
                grade_and_next_essay,
                inputs=[essay_answer_tb, state_essay_list, state_essay_idx, state_essay_results],
                outputs=[
                    essay_progress, essay_q_view,
                    essay_rubric_view, essay_max_score,
                    essay_answer_tb, essay_grade_json, essay_final_summary,
                    state_essay_list, state_essay_idx, state_essay_results
                ]
            )

        with gr.Tab("4) Export"):
            gr.Markdown("Download your quiz as JSON or CSV.")
            export_json_btn = gr.Button("Download Quiz (JSON)")
            export_csv_btn = gr.Button("Download Quiz (CSV)")
            file_json = gr.File(label="Quiz JSON")
            file_csv = gr.File(label="Quiz CSV")

            def export_json(quiz: Dict[str,Any]):
                path = "quiz_export.json"
                with open(path, "w", encoding="utf-8") as f:
                    json.dump(quiz, f, ensure_ascii=False, indent=2)
                return path

            def export_csv(quiz: Dict[str,Any]):
                import csv
                path = "quiz_export.csv"
                fields = ["id","type","question","options","answer","expected_answer","rubric","max_score","explanation"]
                with open(path, "w", newline="", encoding="utf-8") as f:
                    writer = csv.DictWriter(f, fieldnames=fields)
                    writer.writeheader()
                    for q in quiz.get("questions", []):
                        row = {
                            "id": q.get("id"),
                            "type": q.get("type"),
                            "question": q.get("question",""),
                            "options": "|".join(q.get("options", [])) if isinstance(q.get("options"), list) else "",
                            "answer": q.get("answer",""),
                            "expected_answer": q.get("expected_answer",""),
                            "rubric": q.get("rubric",""),
                            "max_score": q.get("max_score",""),
                            "explanation": q.get("explanation",""),
                        }
                        writer.writerow(row)
                return path

            export_json_btn.click(export_json, inputs=[state_quiz], outputs=[file_json])
            export_csv_btn.click(export_csv, inputs=[state_quiz], outputs=[file_csv])

    # Footer/help
    with gr.Accordion("Help & Notes", open=False):
        gr.Markdown("""
**Setup**
- Install dependencies from `requirements.txt`.
- Launch: `python ai_quiz_from_pdf_gradio.py` and open the local URL.

**Usage**
1. Go to **Upload & Generate**: upload your PDF, tweak counts, click *Extract Text* then *Generate Quiz*.
2. In **Take Quiz**, answer MCQ/True-False/Short-Answer and click *Submit* to auto-grade.
3. In **Essay/Theory Grading**, choose an Essay question, paste your answer, and click *Grade*.
4. In **Export**, download the quiz as JSON/CSV for sharing or record-keeping.

**Notes**
- Essay grading uses the LLM with a rubric and max_score. Results are heuristic.
- For scanned PDFs (images), add OCR (e.g., `pytesseract`) in future work.
""")

if __name__ == "__main__":
    # If the API key isn't set, warn in console.
    if not os.getenv("GROQ_API_KEY"):
        print("WARNING: GROQ_API_KEY is not set. Set it before generating or grading with the LLM.")
    demo.launch(debug=True)