content_pipeline.py

"""
Recall — Module A: Content Pipeline.  OWNER: Frank

document -> clean text -> chunks -> list[Card]

Runs in STUB mode out of the box (returns a hardcoded demo deck). Replace the
TODO bodies with real logic; the public signatures must NOT change — app.py and
learning_engine depend on them.
"""
from __future__ import annotations

import os
import re

import llm
from schema import Card, new_card, validate_card

# STUB is owned by llm (single source of truth) and read dynamically as
# `llm.STUB` so every module agrees and runtime/reload changes are honored.


# ---- Public interface ------------------------------------------------------

class ExtractionError(Exception):
    """Raised for bad input (no file, corrupt/empty/image-only PDF) with a
    user-facing message. app.py catches this and shows it — never a crash."""


def extract_text(file) -> str:
    """
    file: a path (str) or a Gradio file object with `.name`. Handles PDF + .txt.
    Returns plain text. Paste-text path in app.py bypasses this.

    Raises ExtractionError (with a friendly message) for bad inputs: no file,
    a missing path, a corrupt/password-protected PDF, an image-only/scanned PDF
    with no selectable text, or an empty file.
    """
    if llm.STUB:
        return ("Photosynthesis is the process by which plants convert light "
                "energy into chemical energy. It occurs in the chloroplasts. "
                "The Calvin cycle fixes carbon dioxide into glucose.")

    if file is None:
        raise ExtractionError("No file provided — upload a PDF/.txt or paste notes.")

    path = getattr(file, "name", file)
    if not os.path.isfile(path):
        raise ExtractionError("Couldn't find that file. Try uploading it again.")

    if str(path).lower().endswith(".pdf"):
        from pypdf import PdfReader
        try:
            reader = PdfReader(path)
            text = "\n\n".join((page.extract_text() or "") for page in reader.pages)
        except Exception:
            raise ExtractionError(
                "Couldn't read that PDF — it may be corrupted or password-protected."
            )
        if not _normalize(text):
            raise ExtractionError(
                "This PDF has no selectable text (looks scanned/image-only). "
                "Try a text-based PDF, or paste the notes instead."
            )
    else:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            text = f.read()
        if not _normalize(text):
            raise ExtractionError("That file is empty — nothing to study from.")

    return _normalize(text)


def generate_deck(text: str, n: int = 12) -> list[Card]:
    """
    Turn source text into ~n good question cards. Quality over quantity.
    Always returns valid Cards (bad model output is skipped, never crashes).
    """
    if llm.STUB:
        return [
            new_card(
                "What does photosynthesis convert light energy into?",
                "Chemical energy (stored as glucose).",
                topic="Photosynthesis",
                source_chunk=text[:120],
                difficulty=1,
            ),
            new_card(
                "Where in the plant cell does photosynthesis occur?",
                "In the chloroplasts.",
                topic="Cell Biology",
                source_chunk=text[:120],
                difficulty=1,
            ),
            new_card(
                "What does the Calvin cycle fix carbon dioxide into?",
                "Glucose.",
                topic="Photosynthesis",
                source_chunk=text[:120],
                difficulty=2,
            ),
        ]

    cards: list[Card] = []
    for chunk in chunk_text(text):
        cards.extend(_cards_from_chunk(chunk))
        if len(cards) >= n:
            break
    return _dedupe(cards)[:n]


# ---- Image-only / scanned PDFs (multimodal model) --------------------------
# When a PDF has no extractable text layer (scanned/photographed), there's
# nothing to chunk — so render the pages to images and let the multimodal model
# (MiniCPM-V) read them directly. Additive to the text path above (unchanged);
# only meaningful with a vision model (RECALL_MODEL=v46, RECALL_STUB=0).

MAX_PDF_IMAGE_PAGES = int(os.getenv("RECALL_MAX_PDF_PAGES", "8"))


def is_image_only_pdf(file) -> bool:
    """True if `file` is a PDF whose pages have no selectable text (scanned/
    image-only) — the case extract_text() rejects and this module renders
    instead. False for non-PDFs / unreadable files (let extract_text surface the
    real error)."""
    path = getattr(file, "name", file)
    if not str(path).lower().endswith(".pdf") or not os.path.isfile(path):
        return False
    try:
        from pypdf import PdfReader
        reader = PdfReader(path)
        text = "".join((page.extract_text() or "") for page in reader.pages)
    except Exception:
        return False
    return not _normalize(text)


def render_pdf_images(file, max_pages: int = MAX_PDF_IMAGE_PAGES) -> list:
    """Render up to `max_pages` PDF pages to RGB PIL.Images for the vision model.
    Returns [] for a non-PDF. Requires PyMuPDF + Pillow (real-model deps)."""
    path = getattr(file, "name", file)
    if not str(path).lower().endswith(".pdf"):
        return []
    import io

    import fitz  # PyMuPDF
    from PIL import Image

    images = []
    with fitz.open(path) as doc:
        for page in list(doc)[:max_pages]:
            pix = page.get_pixmap(dpi=150)
            images.append(Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB"))
    return images


VISION_TARGET_CARDS = 4   # how many distinct questions to aim for from the slides
VISION_MAX_ROUNDS = 6     # cap on model calls so scan-deck latency stays bounded


def _vision_system_prompt() -> str:
    # Labeled lines, NOT JSON: MiniCPM-V writes the right content but mangles JSON
    # (unquoted string values, missing commas) so every object failed to parse.
    # A four-line labeled format has no quotes/commas to get wrong.
    return (
        "You are a quiz generator. You are given the page images of a study "
        "document. Write exactly ONE quiz question testing the material.\n\n"
        "GROUNDING (critical): the question AND its answer must be answerable using "
        "ONLY what is visible in the page images. Do NOT introduce facts, names, or "
        "numbers that are not shown.\n\n"
        "Output EXACTLY these four lines and NOTHING else — no JSON, no quotes, no "
        "extra commentary:\n"
        "QUESTION: <the question, one clear sentence>\n"
        "ANSWER: <concise reference answer, 1-2 sentences>\n"
        "TOPIC: <short concept tag, e.g. Cell Biology>\n"
        "DIFFICULTY: <1, 2, or 3>"
    )


_LABELS = ("QUESTION", "ANSWER", "TOPIC", "DIFFICULTY")


def _labeled_field(label: str, text: str) -> str:
    """Pull one LABEL: value out of the model's labeled-line reply. Captures a
    value that may wrap across lines, up to the next known label or end of text."""
    others = "|".join(L for L in _LABELS if L != label)
    # Labels must start a line (^ in MULTILINE) so stray prose like "Here is your
    # question:" can't be mistaken for the QUESTION field. DOTALL lets a value wrap
    # across lines up to the next known label or end of text.
    m = re.search(
        rf"(?im)^\s*{label}\s*[:\-]\s*(.+?)(?=^\s*(?:{others})\s*[:\-]|\Z)",
        text, re.DOTALL,
    )
    return m.group(1).strip().strip('"').strip() if m else ""


def _parse_labeled_card(text: str, source_chunk: str) -> Card | None:
    """Build a validated Card from a labeled-line vision reply, or None if unusable."""
    question = _labeled_field("QUESTION", text)
    answer = _labeled_field("ANSWER", text)
    topic = _labeled_field("TOPIC", text) or "General"
    dm = re.search(r"[1-3]", _labeled_field("DIFFICULTY", text))
    difficulty = int(dm.group()) if dm else 1
    card = new_card(
        question=question, answer=answer, topic=topic,
        source_chunk=source_chunk, difficulty=difficulty,
    )
    return card if validate_card(card) else None


def _vision_list_facts(images: list, want: int) -> list[str]:
    """Ask the vision model to enumerate the distinct facts on the page(s), one per
    line. Parsing is line-based (strip bullets/numbering), so the model can't break
    it. Drives per-fact question generation below."""
    messages = [
        {"role": "system", "content": (
            "You are given the page images of a study document. List the distinct "
            "key facts or concepts a student should learn from it.\n\n"
            "Use ONLY information visible in the images — do not invent anything.\n"
            f"Output one fact per line, up to {want} lines. Each line is a short "
            "statement (no numbering, no bullets, no extra commentary).")},
        {"role": "user", "content": list(images) + [
            "List the key facts, one per line."]},
    ]
    reply = llm.chat(messages, max_tokens=512)
    facts: list[str] = []
    seen: set[str] = set()
    for line in reply.splitlines():
        # strip leading bullets / numbering ("1.", "-", "*", "•")
        line = re.sub(r"^\s*(?:[-*•]|\d+[.)])\s*", "", line).strip().strip('"').strip()
        if len(line) < 8:
            continue
        key = _norm_key(line)
        if key and key not in seen:
            seen.add(key)
            facts.append(line)
    return facts


def _vision_question_for(images: list, focus: str, level: str) -> Card | None:
    """Generate ONE labeled-format question about a specific fact/focus."""
    ask = (f"Write ONE quiz question from these document page images, in the "
           f"four-line labeled format, specifically about this fact: {focus}\n"
           f"Aim for difficulty {level}.")
    messages = [
        {"role": "system", "content": _vision_system_prompt()},
        {"role": "user", "content": list(images) + [ask]},
    ]
    reply = llm.chat(messages, max_tokens=512)
    return _parse_labeled_card(reply, source_chunk="[scanned PDF page]")


def generate_deck_from_images(images: list, n: int = 12) -> list[Card]:
    """Generate question cards from PDF page images via the multimodal model —
    the image-only/scanned-PDF counterpart to generate_deck(text). Always returns
    valid Cards (bad model output is skipped, never crashes).

    The vision model (MiniCPM-V) answers with a SINGLE question no matter how the
    prompt asks for several, mangles JSON (unquoted values, missing commas), AND
    ignores "ask about a different fact" — so a one-shot or a self-diversifying
    loop both collapse to one card (verified from raw GPU output). Instead we work
    in two phases: enumerate the slide's distinct facts, then ask one labeled-line
    question per fact. Targeting a named fact each call is what makes the questions
    actually differ. Bounded by VISION_MAX_ROUNDS for latency.

    Stub returns the canned demo deck so the flow still runs without a model/GPU.
    """
    if llm.STUB:
        return generate_deck("")  # canned demo deck — no model/GPU
    if not images:
        return []

    target = min(n, VISION_TARGET_CARDS)
    levels = ["1 (direct recall)", "2 (application/explanation)", "3 (synthesis)"]

    facts = _vision_list_facts(images, target * 2)
    # If fact extraction came up thin, fall back to plain "another question" asks so
    # we still try to build a deck rather than give up.
    if len(facts) < target:
        facts = facts + ["a different concept from the images"] * (target - len(facts))

    out: list[Card] = []
    seen: set[str] = set()
    for i, focus in enumerate(facts):
        if len(out) >= target or i >= VISION_MAX_ROUNDS:
            break
        card = _vision_question_for(images, focus, levels[i % 3])
        if card is None:
            continue
        key = _norm_key(card["question"])
        if key and key not in seen:
            seen.add(key)
            out.append(card)

    return _dedupe(out)[:n]


# ---- Internals (Frank implements) ------------------------------------------

def _normalize(text: str) -> str:
    return " ".join(text.split())


def chunk_text(text: str, size: int = 2500) -> list[str]:
    """
    Public chunker (used by generate_deck and by app.py's debug panel).

    Character-based splitter targeting ~500–800 tokens per chunk (1 token ≈ 4 chars,
    so size=2500 ≈ 625 tokens). Breaks at sentence boundaries where possible.
    """
    # Split into sentences on . ! ? followed by whitespace or end-of-string.
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    chunks: list[str] = []
    current: list[str] = []
    current_len = 0

    for sentence in sentences:
        slen = len(sentence)
        if current and current_len + 1 + slen > size:
            chunks.append(" ".join(current))
            current = [sentence]
            current_len = slen
        else:
            current.append(sentence)
            current_len += (1 + slen) if current_len else slen

    if current:
        chunks.append(" ".join(current))

    return chunks or [text]


def _cards_from_chunk(chunk: str) -> list[Card]:
    """Ask the model for JSON cards, parse defensively, validate each."""
    messages = [
        {"role": "system", "content": (
            "You are a quiz generator. Given a study passage, produce 3 to 5 quiz questions "
            "that test different aspects of the material at varying difficulty levels.\n\n"
            "GROUNDING RULES (critical):\n"
            "- Every question AND its answer must be answerable using ONLY the passage below.\n"
            "- Do NOT introduce facts, names, numbers, or examples that are not in the passage.\n"
            "- If the passage is too short or thin for a question, write fewer questions "
            "rather than inventing content.\n"
            "- Quote or paraphrase the passage's own wording in the reference answer.\n\n"
            "Return ONLY a JSON array. Each element must have exactly these keys:\n"
            "  question   — the question text (clear, specific, one sentence)\n"
            "  answer     — a concise reference answer (1-3 sentences), grounded in the passage\n"
            "  topic      — a short tag naming the concept (e.g. 'Cell Biology')\n"
            "  difficulty — integer: 1 (direct recall), 2 (application/explanation), "
            "3 (synthesis/comparison)\n\n"
            "Output format example — return ONE array of OBJECTS exactly like this "
            "(not an array of strings), no other text:\n"
            '[{"question": "What does X do?", "answer": "X does Y.", "topic": "Topic A", '
            '"difficulty": 1}, {"question": "Why does Z occur?", "answer": "Because ...", '
            '"topic": "Topic A", "difficulty": 2}]\n\n'
            "Mix all three difficulty levels across the questions. "
            "No prose, no markdown fences, no explanation outside the JSON array."
        )},
        {"role": "user", "content": f"Passage:\n{chunk}\n\nGenerate the JSON array now."},
    ]
    # Strict-JSON with one repair pass: chat_json feeds a malformed reply back to
    # the model demanding clean JSON before giving up. A chunk that still won't
    # parse is skipped (returns []), never crashing the deck.
    data = llm.chat_json(messages, max_tokens=600, retries=1)
    if isinstance(data, dict):  # tolerate a single object instead of an array
        data = [data]
    if not isinstance(data, list):
        return []
    out: list[Card] = []
    for item in data:
        if not isinstance(item, dict):
            continue
        try:
            difficulty = max(1, min(3, int(item.get("difficulty") or 1)))
        except (ValueError, TypeError):
            difficulty = 1
        card = new_card(
            question=str(item.get("question", "")).strip(),
            answer=str(item.get("answer", "")).strip(),
            topic=str(item.get("topic", "General")).strip() or "General",
            source_chunk=chunk,
            difficulty=difficulty,
        )
        if validate_card(card):
            out.append(card)
    return out


MIN_QUESTION_CHARS = 10   # shorter than this isn't a real question
MIN_ANSWER_CHARS = 2      # an answer needs at least a token of substance


def _norm_key(text: str) -> str:
    """Normalize for near-identical matching: lowercase, strip punctuation,
    collapse whitespace. 'What is X?' and 'what is x' collapse to one key."""
    return re.sub(r"[^a-z0-9 ]", "", text.lower()).strip()


def _dedupe(cards: list[Card]) -> list[Card]:
    """Drop near-identical questions and low-quality cards.

    Near-identical = same question once punctuation/case/whitespace is
    normalized away. Low-quality = a question/answer too short to be useful.
    ~12 great cards beats 40 mediocre ones.
    """
    seen: set[str] = set()
    out: list[Card] = []
    for c in cards:
        question = c["question"].strip()
        answer = c["answer"].strip()
        if len(question) < MIN_QUESTION_CHARS or len(answer) < MIN_ANSWER_CHARS:
            continue  # drop low-quality
        key = _norm_key(question)
        if not key or key in seen:
            continue  # drop empty/near-duplicate
        seen.add(key)
        out.append(c)
    return out


# STRETCH(Frank): difficulty dial — backs Arturo's UI toggle (NAH-32).
def regenerate(card: Card, direction: str) -> Card:
    """Rewrite a card harder or easier on the SAME concept, grounded in the same
    source_chunk. direction: 'harder' | 'easier'. Always returns a valid Card —
    on any failure it falls back to the original so the UI never breaks.
    """
    step = 1 if direction == "harder" else -1
    new_difficulty = max(1, min(3, card["difficulty"] + step))

    if llm.STUB:
        prefix = "[harder] " if direction == "harder" else "[easier] "
        return new_card(
            prefix + card["question"],
            card["answer"],
            topic=card["topic"],
            source_chunk=card["source_chunk"],
            difficulty=new_difficulty,
            parent_id=card.get("parent_id"),
        )

    level = ("more challenging (synthesis/comparison/application)" if direction == "harder"
             else "simpler (direct recall)")
    messages = [
        {"role": "system", "content": (
            "You rewrite a single quiz question to a different difficulty while "
            "keeping the SAME underlying concept. Stay grounded in the passage — "
            "do not introduce facts that are not in it. "
            "Return ONLY a JSON object with keys: question, answer, topic. "
            "Example (return ONE object exactly like this, no other text):\n"
            '{"question": "What does X do?", "answer": "X does Y.", "topic": "Topic A"}'
        )},
        {"role": "user", "content": (
            f"Passage:\n{card['source_chunk']}\n\n"
            f"Current question: {card['question']}\n"
            f"Current answer: {card['answer']}\n"
            f"Topic: {card['topic']}\n\n"
            f"Rewrite it to be {level}, same concept."
        )},
    ]
    data = llm.chat_json(messages, max_tokens=300, retries=1)
    if isinstance(data, dict):
        out = new_card(
            str(data.get("question", "")).strip(),
            str(data.get("answer", "")).strip(),
            topic=str(data.get("topic", card["topic"])).strip() or card["topic"],
            source_chunk=card["source_chunk"],
            difficulty=new_difficulty,
            parent_id=card.get("parent_id"),
        )
        if validate_card(out):
            return out
    return card  # safe fallback — never break the study loop