chatkit/backend/app/question_generator.py

"""Practice question retrieval from the shared question bank, with LLM answer verification."""
from __future__ import annotations

import asyncio
import json
import logging
import random
import re
from dataclasses import dataclass, field

from openai import AsyncOpenAI

from .config import get_settings
from .database import query_question_bank, save_question_bank_passage
from .question_rewriter import (
    generate_passage_question,
    needs_missing_image,
    needs_passage,
    rewrite_question,
)

logger = logging.getLogger(__name__)

# Matches "(A) ...(B) ...(C) ...(D) ..."
_OPTIONS_RE = re.compile(
    r"\(A\)\s*(.*?)\s*\(B\)\s*(.*?)\s*\(C\)\s*(.*?)\s*\(D\)\s*(.*?)$",
    re.DOTALL,
)

_ANSWER_SYSTEM = (
    "You are an English teacher for Taiwanese junior high students (CEFR A2-B1). "
    "For each multiple-choice question, identify which option (A, B, C, or D) is correct. "
    "Return ONLY a JSON array of uppercase answer letters, one per question, in order. "
    'Example for 3 questions: ["A","C","B"]'
)


def _parse_question(raw: str) -> tuple[str, dict[str, str]] | None:
    """Split `stem (A) x (B) y (C) z (D) w` into (stem, {A,B,C,D}). Returns None if no options."""
    m = _OPTIONS_RE.search(raw)
    if not m:
        return None
    stem = raw[: m.start()].strip()
    opts = {
        "A": m.group(1).strip(),
        "B": m.group(2).strip(),
        "C": m.group(3).strip(),
        "D": m.group(4).strip(),
    }
    return stem, opts


async def _check_answers(questions: list[dict], model: str) -> list[str]:
    """Ask the LLM to identify the correct A/B/C/D for each question in one batch call."""
    lines = []
    for i, q in enumerate(questions, 1):
        opts = q["options"]
        lines.append(
            f"{i}. {q['question']}\n"
            f"   (A) {opts['A']}  (B) {opts['B']}  (C) {opts['C']}  (D) {opts['D']}"
        )

    settings = get_settings()
    client = AsyncOpenAI(api_key=settings.openai_api_key)
    response = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": _ANSWER_SYSTEM},
            {"role": "user", "content": "\n".join(lines)},
        ],
        temperature=0,
        max_tokens=max(64, len(questions) * 6),
    )

    content = (response.choices[0].message.content or "").strip()
    content = re.sub(r"^```json?\s*|\s*```$", "", content)
    try:
        answers = json.loads(content)
        if isinstance(answers, list) and len(answers) == len(questions):
            return [str(a).strip().upper() for a in answers]
    except Exception:
        pass

    # Fallback: pull bare letters from the text
    letters = re.findall(r"\b([ABCD])\b", content)
    return (letters + [""] * len(questions))[: len(questions)]


@dataclass
class QuestionFilter:
    category: str
    tags: list[str] = field(default_factory=list)


def _split_count(total: int, n: int) -> list[int]:
    """Split `total` into `n` near-equal parts (earlier parts get the remainder)."""
    if n <= 0:
        return []
    base, rem = divmod(total, n)
    return [base + 1 if i < rem else base for i in range(n)]


async def _rewrite_candidate(q: dict, model: str) -> dict:
    """Rephrase one ordinary candidate; falls back to the original wording on failure."""
    try:
        rewritten = await rewrite_question(
            q["question"], q["options"], q["answer"], q["category"], passage=q.get("passage"), model=model,
        )
        return {**q, **rewritten}
    except Exception as exc:
        logger.warning("Question rewrite failed, keeping original wording: %s", exc)
        return q


async def _build_passage_candidate(pc: dict, model: str) -> dict | None:
    """Author a fresh passage+question for a row that had none.

    Drops the candidate entirely on failure — never returns a 篇章/文意推論
    question without its passage. Unlike an ordinary rewrite (which can
    safely fall back to the original bank wording), the original stem here
    can't be verified without the passage it was originally paired with, so
    a passage-less fallback would just reproduce the broken, unanswerable
    question this whole path exists to avoid.
    """
    try:
        result = await generate_passage_question(pc["category"], pc.get("tags"), model=model)
        if not result["question"] or not result["passage"]:
            raise ValueError("empty passage/question from LLM")
    except Exception as exc:
        logger.warning("Passage generation failed for category %s, dropping question: %s", pc["category"], exc)
        return None

    try:
        await save_question_bank_passage(pc["question_bank_id"], result["passage"])
    except Exception as exc:
        logger.warning("Failed to persist generated passage for row %d: %s", pc["question_bank_id"], exc)

    return result


async def generate_questions(
    filters: list[QuestionFilter],
    count: int = 10,
    model: str = "gpt-4o-mini",
) -> list[dict]:
    """Return `count` questions sampled from the bank, each rewritten by the LLM.

    Questions are returned as structured multiple-choice: {question, options, answer, category, passage}.

    When a filter selects specific tags, the questions are split evenly across
    those tags (remainder distributed to the first tags) — no other tags in the
    category are pulled in. A tag with zero questions in the bank falls back to
    the broader category for its share (handled by `query_question_bank`); a tag
    with *some* questions never gets padded from other tags.

    Two kinds of bank rows are handled specially instead of being rewritten
    verbatim:
    - Questions that reference an image with no captured description are
      excluded — the correct answer can't be verified without seeing the
      picture (see `question_rewriter.needs_missing_image`).
    - Questions that depend on a reading passage that was never captured
      (篇章/文意推論 categories) get a brand-new passage + question authored
      from scratch instead, since the original stem/answer can't be verified
      without the passage it was originally paired with. The generated
      passage is persisted back onto the bank row for reuse next time.
    """
    if not filters:
        return []

    # One unit per (category, tag): a filter with tags becomes one unit per
    # tag so each tag gets its own even quota; a filter with no tags is a
    # single whole-category unit.
    units: list[tuple[str, list[str]]] = []
    for f in filters:
        if f.tags:
            units.extend((f.category, [tag]) for tag in f.tags)
        else:
            units.append((f.category, []))

    quotas = _split_count(count, len(units))
    candidates: list[dict] = []
    passage_candidates: list[dict] = []

    for (category, tags), quota in zip(units, quotas):
        if quota == 0:
            continue
        rows = await query_question_bank(category, tags, limit=max(quota * 3, 8))
        random.shuffle(rows)
        picked = 0
        for row in rows:
            if picked >= quota:
                break
            raw = (row.get("question_text") or "").strip()
            if needs_missing_image(raw):
                continue
            if needs_passage(row.get("main_category") or category, raw, row.get("passage")):
                passage_candidates.append({
                    "question_bank_id": row["id"],
                    "category": category,
                    "tags": tags,
                })
                picked += 1
                continue
            parsed = _parse_question(raw)
            if parsed is None:
                continue
            stem, opts = parsed
            candidates.append({
                "question": stem,
                "options": opts,
                "answer": (row.get("answer") or "").strip(),
                "category": category,
                "passage": row.get("passage"),
            })
            picked += 1

    if not candidates and not passage_candidates:
        return []

    # Fill in answers via LLM for any ordinary question missing one, before rewriting
    need_check = [q for q in candidates if not q["answer"]]
    if need_check:
        try:
            answers = await _check_answers(need_check, model)
            for q, ans in zip(need_check, answers):
                q["answer"] = ans
        except Exception as exc:
            logger.warning("Answer checking failed: %s", exc)

    rewritten, built_passages = await asyncio.gather(
        asyncio.gather(*(_rewrite_candidate(q, model) for q in candidates)),
        asyncio.gather(*(_build_passage_candidate(pc, model) for pc in passage_candidates)),
    )

    results = list(rewritten) + [p for p in built_passages if p is not None]

    # Old-data passage fallbacks (see _build_passage_candidate) may still be
    # missing an answer — give answer-checking one more shot for those.
    still_need_check = [q for q in results if not q["answer"]]
    if still_need_check:
        try:
            answers = await _check_answers(still_need_check, model)
            for q, ans in zip(still_need_check, answers):
                q["answer"] = ans
        except Exception as exc:
            logger.warning("Fallback answer checking failed: %s", exc)

    random.shuffle(results)
    return results