src/tutor/services/reading.py

"""Reading exercise pipeline (M1).

Generator proposes, classifier disposes: texts are produced by the LLM at a
requested CEFR level, then *verified* by our own classifier — a mismatch is
sent back once with corrective feedback, and an honest label is served if the
rewrite still disagrees. Comprehension questions are requested as strict JSON,
validated with pydantic (one retry on malformed output), then gated by a
minimal LLM judge before being served (the judge fails open: a broken judge
must not block the product). Every LLM product is cached content-addressed
(prompt version + model + inputs), so repeated demos cost zero quota.
"""

import json
import random
from typing import Literal

from pydantic import BaseModel, Field, ValidationError, model_validator

from tutor.services.cache import FileCache
from tutor.services.llm.base import ChatMessage, LLMClient

PROMPT_VERSION = "1"
CEFR_LEVELS = ["A1", "A2", "B1", "B2", "C1", "C2"]
WORD_TARGETS = {"A1": 80, "A2": 120, "B1": 180, "B2": 220, "C1": 260, "C2": 280}
LEVEL_DESCRIPTORS = {
    "A1": "very short simple sentences, present tense, the ~500 most frequent words",
    "A2": "short sentences, simple past and future, everyday topics and vocabulary",
    "B1": "connected paragraphs, some subordinate clauses, common idioms, familiar topics",
    "B2": "varied sentence structures, abstract topics, opinion and argument, wider vocabulary",
    "C1": "complex structures, low-frequency vocabulary, implicit meaning, nuanced argument",
    "C2": "sophisticated prose, rare vocabulary, dense argumentation, subtle register shifts",
}
DEFAULT_TOPICS = [
    "a city changing through the seasons",
    "an unexpected friendship",
    "how a small invention changed daily life",
    "a journey that did not go as planned",
    "food and memory",
    "the ocean",
    "learning a new skill as an adult",
    "a place that no longer exists",
]


class ReadingError(RuntimeError):
    """A reading exercise could not be built (surfaced to the UI)."""


class Question(BaseModel):
    question: str
    options: list[str] = Field(min_length=3, max_length=4)
    answer_index: int = Field(ge=0)
    explanation: str = ""

    @model_validator(mode="after")
    def _index_in_range(self) -> "Question":
        if self.answer_index >= len(self.options):
            msg = f"answer_index {self.answer_index} out of range for {len(self.options)} options"
            raise ValueError(msg)
        return self


class ClozeBlank(BaseModel):
    """One validated gap: the answer is a verbatim token of the text."""

    answer: str
    start: int  # character offset in the text (computed Python-side, never trusted from the LLM)
    distractors: list[str] = Field(default_factory=list)
    hint: str = ""

    @property
    def end(self) -> int:
        return self.start + len(self.answer)

    def options(self) -> list[str]:
        """Answer + distractors, de-duplicated, for the optional multiple-choice hint."""
        seen, out = set(), []
        for candidate in [self.answer, *self.distractors]:
            key = candidate.casefold()
            if key not in seen:
                seen.add(key)
                out.append(candidate)
        return out


class ClozeExercise(BaseModel):
    blanks: list[ClozeBlank]


class ReadingExercise(BaseModel):
    text: str
    source: Literal["generated", "user"]
    activity: Literal["questions", "cloze"] = "questions"
    requested_level: str | None = None
    classified_level: str | None = None
    classifier_score: float | None = None
    attempts: int = 1
    questions: list[Question] = Field(default_factory=list)
    cloze: ClozeExercise | None = None


TEXT_PROMPT = """You are writing graded reading material for English learners.
Write one self-contained text in English at CEFR level {level}, about: {topic}.
Target length: about {words} words.
Level guidance for {level}: {descriptor}.
Rules: return ONLY the text itself — no title, no preamble, no notes.{feedback}"""

QUESTIONS_PROMPT = """You write comprehension checks for English learners.
Text:
\"\"\"{text}\"\"\"
Create exactly {n} multiple-choice comprehension questions strictly answerable
from the text alone{level_clause}.
Return ONLY a JSON object — no markdown fences, no commentary — exactly this shape:
{{"questions": [{{"question": "...", "options": ["...", "...", "...", "..."],
"answer_index": 0, "explanation": "..."}}]}}
Rules: 4 options per question; exactly one correct option; answer_index is the
0-based index of the correct option; the explanation points to the relevant
part of the text.{retry_clause}"""

JUDGE_PROMPT = """You are a strict reviewer of reading-comprehension questions.
Text:
\"\"\"{text}\"\"\"
Questions (JSON): {questions_json}
For each question, check that: (a) it is answerable from the text alone,
(b) the option at answer_index is correct, (c) no other option is also correct.
Return ONLY JSON: {{"verdicts": [{{"index": 0, "valid": true, "reason": ""}}]}}"""

CLOZE_PROMPT = """You design fill-in-the-blank (cloze) exercises for English learners.
Text:
\"\"\"{text}\"\"\"
Choose exactly {n} words to blank out that test real language mastery at CEFR
level {level}{level_clause} — prefer verbs (tense/form), prepositions,
connectors and collocations over content nouns or names. Avoid blanking the
same word twice, and avoid trivial words (the, a, is).
Each "answer" MUST be a single word copied EXACTLY (same case) from the text.
Provide 3 plausible-but-wrong distractors of the same kind, and a short hint.
Return ONLY a JSON object — no markdown fences, no commentary — exactly:
{{"blanks": [{{"answer": "...", "distractors": ["...", "...", "..."],
"hint": "..."}}]}}{retry_clause}"""


def locate_blanks(text: str, raw_blanks: list[dict]) -> list[ClozeBlank]:
    """Validate LLM-proposed answers against the text and assign real offsets.

    The LLM is not trusted for indexing: each answer must occur verbatim as a
    whole word, occurrences are consumed left to right so repeats land on
    distinct positions, and unfound answers are dropped. Result is sorted by
    position. Raises ValueError if nothing usable remains.
    """
    blanks: list[ClozeBlank] = []
    search_from: dict[str, int] = {}
    for item in raw_blanks:
        answer = str(item.get("answer", "")).strip()
        if not answer or " " in answer:
            continue
        cursor = search_from.get(answer, 0)
        while True:
            start = text.find(answer, cursor)
            if start == -1:
                break
            before = text[start - 1] if start > 0 else " "
            after = text[start + len(answer)] if start + len(answer) < len(text) else " "
            if not before.isalnum() and not after.isalnum():  # whole-word match
                break
            cursor = start + 1
        if start == -1:
            continue  # the model invented a word that is not in the text
        search_from[answer] = start + len(answer)
        blanks.append(
            ClozeBlank(
                answer=answer,
                start=start,
                distractors=[str(d) for d in item.get("distractors", [])][:3],
                hint=str(item.get("hint", "")),
            )
        )
    if not blanks:
        msg = "none of the proposed blanks were found verbatim in the text"
        raise ValueError(msg)
    return sorted(blanks, key=lambda blank: blank.start)


def extract_json(raw: str) -> dict:
    """Parse a JSON object out of an LLM reply, tolerating fences and chatter."""
    start, end = raw.find("{"), raw.rfind("}")
    if start == -1 or end <= start:
        msg = "no JSON object found in the model reply"
        raise ValueError(msg)
    return json.loads(raw[start : end + 1])


async def _ask(llm: LLMClient, prompt: str, *, temperature: float, max_tokens: int) -> str:
    response = await llm.complete(
        [ChatMessage(role="user", content=prompt)],
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.text.strip()


def _parse_questions(raw: str) -> list[Question]:
    payload = extract_json(raw)
    items = payload.get("questions")
    if not isinstance(items, list) or not items:
        msg = "JSON has no non-empty 'questions' list"
        raise ValueError(msg)
    return [Question.model_validate(item) for item in items]


async def _judge_questions(
    llm: LLMClient, *, text: str, questions: list[Question]
) -> list[Question]:
    """Keep only judge-approved questions; fail open if the judge misbehaves."""
    questions_json = json.dumps([q.model_dump() for q in questions], ensure_ascii=False)
    try:
        raw = await _ask(
            llm,
            JUDGE_PROMPT.format(text=text, questions_json=questions_json),
            temperature=0.0,
            max_tokens=400,
        )
        verdicts = {
            int(v["index"]): bool(v["valid"]) for v in extract_json(raw).get("verdicts", [])
        }
    except Exception:  # judge is a gate, not a wall
        return questions
    kept = [q for i, q in enumerate(questions) if verdicts.get(i, True)]
    return kept or questions  # a judge that rejects everything is suspect: fail open


async def generate_questions(
    llm: LLMClient,
    cache: FileCache,
    *,
    text: str,
    level: str | None,
    model_name: str,
    n: int = 3,
) -> list[Question]:
    key = FileCache.key("questions", PROMPT_VERSION, model_name, level or "", str(n), text)
    if (cached := cache.get(key)) is not None:
        return [Question.model_validate(item) for item in cached["questions"]]

    level_clause = f", with question language at or below CEFR {level}" if level else ""
    prompt = QUESTIONS_PROMPT.format(text=text, n=n, level_clause=level_clause, retry_clause="")
    raw = await _ask(llm, prompt, temperature=0.3, max_tokens=900)
    try:
        questions = _parse_questions(raw)
    except (ValueError, ValidationError) as exc:
        retry_clause = (
            f"\nYour previous reply was rejected ({exc}). Output valid JSON of the exact "
            "shape above and nothing else."
        )
        prompt = QUESTIONS_PROMPT.format(
            text=text, n=n, level_clause=level_clause, retry_clause=retry_clause
        )
        raw = await _ask(llm, prompt, temperature=0.0, max_tokens=900)
        try:
            questions = _parse_questions(raw)
        except (ValueError, ValidationError) as exc2:
            msg = f"the model could not produce valid questions ({exc2})"
            raise ReadingError(msg) from exc2

    questions = await _judge_questions(llm, text=text, questions=questions)
    cache.set(key, {"questions": [q.model_dump() for q in questions]})
    return questions


async def generate_cloze(
    llm: LLMClient,
    cache: FileCache,
    *,
    text: str,
    level: str | None,
    model_name: str,
    n: int = 5,
) -> ClozeExercise:
    key = FileCache.key("cloze", PROMPT_VERSION, model_name, level or "", str(n), text)
    if (cached := cache.get(key)) is not None:
        return ClozeExercise.model_validate(cached)

    level_clause = f" (keep distractors at or below CEFR {level})" if level else ""

    def parse(raw: str) -> ClozeExercise:
        payload = extract_json(raw)
        items = payload.get("blanks")
        if not isinstance(items, list) or not items:
            msg = "JSON has no non-empty 'blanks' list"
            raise ValueError(msg)
        return ClozeExercise(blanks=locate_blanks(text, items))

    prompt = CLOZE_PROMPT.format(
        text=text, n=n, level=level or "B1", level_clause=level_clause, retry_clause=""
    )
    raw = await _ask(llm, prompt, temperature=0.4, max_tokens=700)
    try:
        cloze = parse(raw)
    except (ValueError, ValidationError) as exc:
        retry_clause = (
            f"\nYour previous reply was rejected ({exc}). Every answer must be a single "
            "word copied verbatim from the text. Output valid JSON of the exact shape "
            "above and nothing else."
        )
        prompt = CLOZE_PROMPT.format(
            text=text,
            n=n,
            level=level or "B1",
            level_clause=level_clause,
            retry_clause=retry_clause,
        )
        raw = await _ask(llm, prompt, temperature=0.0, max_tokens=700)
        try:
            cloze = parse(raw)
        except (ValueError, ValidationError) as exc2:
            msg = f"the model could not produce a valid cloze exercise ({exc2})"
            raise ReadingError(msg) from exc2

    cache.set(key, cloze.model_dump())
    return cloze


async def generate_verified_text(
    llm: LLMClient,
    classifier,
    *,
    level: str,
    topic: str,
    max_attempts: int = 2,
) -> tuple[str, str | None, float | None, int]:
    """Generate at `level`, verify with the classifier, retry once with feedback."""
    feedback = ""
    text, classified, score = "", None, None
    for attempt in range(1, max_attempts + 1):
        text = await _ask(
            llm,
            TEXT_PROMPT.format(
                level=level,
                topic=topic,
                words=WORD_TARGETS[level],
                descriptor=LEVEL_DESCRIPTORS[level],
                feedback=feedback,
            ),
            temperature=0.8,
            max_tokens=700,
        )
        if not text:
            msg = "the model returned an empty text"
            raise ReadingError(msg)
        if classifier is None:
            return text, None, None, attempt
        prediction = classifier.classify_text(text)
        classified, score = prediction.level, prediction.score
        if classified == level:
            return text, classified, score, attempt
        direction = (
            "simplify it: shorter sentences, more frequent vocabulary"
            if CEFR_LEVELS.index(classified) > CEFR_LEVELS.index(level)
            else "make it harder: longer sentences, less frequent vocabulary, "
            "more advanced structures"
        )
        feedback = (
            f"\nYour previous attempt was classified as {classified} by an automatic "
            f"CEFR grader, but {level} was requested — {direction}."
        )
    return text, classified, score, max_attempts  # served with an honest mismatch label


async def build_reading_exercise(
    llm: LLMClient,
    classifier,
    cache: FileCache,
    *,
    level: str,
    model_name: str,
    topic: str | None = None,
    activity: Literal["questions", "cloze"] = "questions",
) -> ReadingExercise:
    if level not in CEFR_LEVELS:
        msg = f"unknown CEFR level {level!r}"
        raise ReadingError(msg)
    topic = (topic or "").strip() or random.choice(DEFAULT_TOPICS)

    text_key = FileCache.key("text", PROMPT_VERSION, model_name, level, topic)
    if (cached := cache.get(text_key)) is not None:
        text = cached["text"]
        classified, score, attempts = (
            cached.get("classified_level"),
            cached.get("classifier_score"),
            cached.get("attempts", 1),
        )
    else:
        text, classified, score, attempts = await generate_verified_text(
            llm, classifier, level=level, topic=topic
        )
        cache.set(
            text_key,
            {
                "text": text,
                "classified_level": classified,
                "classifier_score": score,
                "attempts": attempts,
            },
        )

    questions: list[Question] = []
    cloze: ClozeExercise | None = None
    if activity == "cloze":
        cloze = await generate_cloze(llm, cache, text=text, level=level, model_name=model_name)
    else:
        questions = await generate_questions(
            llm, cache, text=text, level=level, model_name=model_name
        )
    return ReadingExercise(
        text=text,
        source="generated",
        activity=activity,
        requested_level=level,
        classified_level=classified,
        classifier_score=score,
        attempts=attempts,
        questions=questions,
        cloze=cloze,
    )


async def exercise_from_user_text(
    llm: LLMClient,
    classifier,
    cache: FileCache,
    *,
    text: str,
    model_name: str,
    activity: Literal["questions", "cloze"] = "questions",
) -> ReadingExercise:
    text = text.strip()
    if not text:
        msg = "paste a text first"
        raise ReadingError(msg)
    classified_level: str | None = None
    classifier_score: float | None = None
    if classifier is not None:
        prediction = classifier.classify_text(text)
        classified_level, classifier_score = prediction.level, prediction.score

    questions: list[Question] = []
    cloze: ClozeExercise | None = None
    if activity == "cloze":
        cloze = await generate_cloze(
            llm, cache, text=text, level=classified_level, model_name=model_name
        )
    else:
        questions = await generate_questions(
            llm, cache, text=text, level=classified_level, model_name=model_name
        )
    return ReadingExercise(
        text=text,
        source="user",
        activity=activity,
        classified_level=classified_level,
        classifier_score=classifier_score,
        questions=questions,
        cloze=cloze,
    )