SkillBridge AI — Test Prep Academy Platform

"""
AI-powered quiz generator. Given a standard and a topic, asks the LLM to
produce a JSON-formatted quiz with multiple-choice and short-answer questions.

The instructor reviews and edits before publishing to students. We deliberately
return raw JSON-shaped dicts here rather than DB rows — the calling service
persists them after the instructor approves.

The important pattern below: AI output is treated as UNTRUSTED. Code fences are
stripped, JSON is parsed with explicit error handling, and every question is
validated (valid type, non-empty prompt, exactly one correct option) before it
reaches the instructor. Malformed output raises rather than showing garbage.
"""
import json
from typing import Literal

from app.services.ai_provider import get_provider
from app.services.standards_catalog import STANDARDS, get_standard_by_id


# Tunable: per-quiz max output tokens. 10 questions ≈ 4000 output tokens.
_MAX_TOKENS_PER_QUIZ = 8000


def _build_generation_prompt(
    standard_label: str,
    standard_description: str,
    topic: str,
    num_questions: int,
    difficulty: str,
    mix: str,
) -> str:
    """The system prompt that drives the AI to produce a JSON quiz."""
    mix_instructions = {
        "all_mc": "All questions must be multiple_choice with 4 options each.",
        "mostly_mc": (
            "Mostly multiple_choice (4 options). Include 1-2 short_answer "
            "questions for higher-order thinking."
        ),
        "mixed": (
            "Roughly half multiple_choice (4 options) and half short_answer. "
            "Use short_answer for questions that need worked solutions."
        ),
    }.get(mix, "Mostly multiple_choice with 4 options each.")

    return f"""You are an expert curriculum designer creating practice questions for a test prep academy.

STANDARD: {standard_label}
ABOUT THIS STANDARD: {standard_description}
TOPIC FOCUS: {topic}
NUMBER OF QUESTIONS: {num_questions}
DIFFICULTY: {difficulty}
QUESTION MIX: {mix_instructions}

YOUR TASK:
Generate {num_questions} original practice questions aligned to this standard and topic.
Return ONLY valid JSON — no markdown fences, no commentary before or after.

REQUIRED JSON SHAPE:
{{
  "questions": [
    {{
      "type": "multiple_choice",
      "prompt": "The question text. Be specific and clear.",
      "options": [
        {{"text": "Option A", "is_correct": false}},
        {{"text": "Option B", "is_correct": true}},
        {{"text": "Option C", "is_correct": false}},
        {{"text": "Option D", "is_correct": false}}
      ],
      "explanation": "Brief explanation shown after grading. 1-2 sentences.",
      "points": 10
    }},
    {{
      "type": "short_answer",
      "prompt": "The question text.",
      "reference_answer": "The expected answer or key elements a correct response must include.",
      "explanation": "Brief explanation of the correct approach. 1-2 sentences.",
      "points": 10
    }}
  ]
}}

CRITICAL RULES:
- Do NOT reproduce real exam questions verbatim. Generate ORIGINAL questions in the style of the standard.
- For multiple_choice: exactly ONE option must have is_correct=true.
- For multiple_choice: distractors should be plausible — represent common student mistakes, not silly answers.
- All questions must be answerable without external context (no "see figure 3").
- Mathematical notation: use plain text (e.g. "x^2", "sqrt(5)", "3/4") — no LaTeX, no special characters.
- Difficulty calibration: {difficulty} means the typical student at this level should find these challenging-but-fair.

Return only the JSON object. Begin your response with {{ and end with }}."""


def generate_quiz(
    prompt: str,
    standard_id: str | None = None,
    num_questions: int = 10,
    difficulty: Literal["easy", "medium", "hard"] = "medium",
    mix: Literal["all_mc", "mostly_mc", "mixed"] = "mostly_mc",
) -> dict:
    """
    Generate a quiz draft. Returns the cleaned questions plus token usage and
    (when no standard was chosen) an AI-suggested standard or topic label for
    auto-tagging.

    Raises ValueError if the AI's output isn't valid JSON or doesn't match the
    expected shape — caller should retry or surface to the instructor.
    """
    if num_questions < 1 or num_questions > 25:
        raise ValueError("num_questions must be between 1 and 25")

    standard = get_standard_by_id(standard_id) if standard_id else None
    standard_label = standard["label"] if standard else "Custom"
    standard_description = (
        standard["description"]
        if standard
        else "No specific standard chosen — use the user prompt as the source of truth."
    )
    system_prompt = _build_generation_prompt(
        standard_label=standard_label,
        standard_description=standard_description,
        topic=prompt,
        num_questions=num_questions,
        difficulty=difficulty,
        mix=mix,
    )

    # If no standard was picked, ask the AI to classify the content — either
    # match the catalog or propose a free-form topic label. This drives the
    # auto-tagging of generated quizzes.
    if standard is None:
        catalog_lines = [f'  - "{s["id"]}": {s["label"]}' for s in STANDARDS]
        system_prompt += (
            "\n\nADDITIONALLY, classify this content with a topic tag. Add TWO top-level fields to your JSON:\n"
            "1. `suggested_standard_id`: if the content clearly matches one of the standards below, "
            "set this to the exact id (otherwise null):\n"
            + "\n".join(catalog_lines)
            + "\n\n2. `suggested_topic_label`: ALWAYS set this. A short 2-5 word topic label "
            'describing what the content covers (e.g. "Calculus Derivatives", "Linear Equations"). '
            "Title Case. Use the most specific accurate label."
        )

    provider = get_provider()
    response = provider.send_message(
        system_prompt=system_prompt,
        messages=[
            {
                "role": "user",
                "content": (
                    f"Generate a {num_questions}-question {difficulty} quiz. "
                    f"User's request: {prompt}. Return JSON only."
                ),
            }
        ],
        max_tokens=_MAX_TOKENS_PER_QUIZ,
    )

    raw = response.content.strip()
    # Defensive: strip code fences if the model returned them despite instructions
    if raw.startswith("```"):
        raw = raw.lstrip("`")
        if raw.lower().startswith("json"):
            raw = raw[4:]
        raw = raw.strip()
        if raw.endswith("```"):
            raw = raw[:-3].strip()

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as e:
        raise ValueError(
            f"AI returned invalid JSON: {e}. First 200 chars: {raw[:200]!r}"
        )

    questions = parsed.get("questions")
    if not isinstance(questions, list) or not questions:
        raise ValueError("AI returned no questions in expected shape")

    # Light validation — surface issues early so the instructor isn't shown garbage
    cleaned: list[dict] = []
    for i, q in enumerate(questions):
        q_type = q.get("type")
        if q_type not in ("multiple_choice", "short_answer"):
            raise ValueError(f"Question {i + 1} has invalid type: {q_type}")
        if not q.get("prompt"):
            raise ValueError(f"Question {i + 1} missing prompt")

        item: dict = {
            "type": q_type,
            "prompt": q["prompt"].strip(),
            "explanation": (q.get("explanation") or "").strip(),
            "points": int(q.get("points") or 10),
        }

        if q_type == "multiple_choice":
            options = q.get("options") or []
            if len(options) < 2:
                raise ValueError(f"Question {i + 1} needs at least 2 options")
            correct_count = sum(1 for o in options if o.get("is_correct"))
            if correct_count != 1:
                raise ValueError(
                    f"Question {i + 1} must have exactly one correct option, found {correct_count}"
                )
            item["options"] = [
                {
                    "text": (o.get("text") or "").strip(),
                    "is_correct": bool(o.get("is_correct")),
                }
                for o in options
            ]
        else:  # short_answer
            item["reference_answer"] = (q.get("reference_answer") or "").strip()

        cleaned.append(item)

    return {
        "questions": cleaned,
        "input_tokens": response.input_tokens,
        "output_tokens": response.output_tokens,
        "tokens_used": response.input_tokens + response.output_tokens,
    }