pdf-quiz-generator/backend/app/services/extraction_modes.py

"""Additional extraction modes that run alongside (not replacing) the standard extractor.

Modes
-----
questions_only   Extract Q+options with no answers. User fills answers later via QuizEditPage.
two_step         Separate answer key section (PREP 2013): Phase 1 = questions, Phase 2 = key, Phase 3 = match.
regex            AI generates a regex pattern for the document's answer format, then we apply it.
ai_decide        AI samples the document and picks standard / two_step / questions_only.
generate         AI reads plain text/study material and creates MCQ questions from scratch.
"""
from __future__ import annotations
import json
import logging
import re

from app.services import ai_service, vector_service

logger = logging.getLogger(__name__)


def _normalize(text: str) -> str:
    return (text.replace("Pref erred", "Preferred").replace("Pre ferred", "Preferred")
            .replace("Prefer red", "Preferred").replace("ltem", "Item").replace("ltcm", "Item"))


# ─── QUESTIONS ONLY ──────────────────────────────────────────────────────────

QUESTIONS_ONLY_PROMPT = """Extract every question from this PREP exam content.
Do NOT look for correct answers — we only need the question text and answer options.

Return ONLY JSON:
{{"questions": [
  {{
    "item_number": "<digits only, or null>",
    "question_text": "<full vignette + stem>",
    "question_type": "mcq",
    "options": ["<A>", "<B>", "<C>", "<D>", "<E>"],
    "has_figure": false,
    "page_reference": {page_ref}
  }}
]}}

Rules:
- A new question starts with "Item NNN" or "ltem NNN".
- Extract ALL questions even if no answer is visible.
- Do NOT include answer explanations or "Preferred Response:" content as questions.
- has_figure: true ONLY if the question references an image, figure, radiograph, photo, ECG, or chart essential to answering. false for decorative/branding images.
- Return ONLY JSON — no markdown, no preamble.

Content (pages {page_info}):
{content}"""


def extract_questions_only(content: str, page_info: str, page_ref: int | None,
                            model_id: str | None, api_key: str | None) -> list[dict]:
    """Extract questions + options with correct_answer = PENDING (to be filled by admin)."""
    content = ai_service._truncate_content(content)
    prompt = QUESTIONS_ONLY_PROMPT.format(
        content=content, page_info=page_info,
        page_ref=page_ref if page_ref else "null",
    )
    for attempt in range(3):
        try:
            text = ai_service._call_model(prompt, model_id, api_key).strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            data = json.loads(text)
            qs = data.get("questions", data) if isinstance(data, dict) else data
            result = []
            for q in qs:
                if not q.get("question_text"):
                    continue
                result.append({
                    "item_number": str(q.get("item_number") or "").strip().lstrip("0") or None,
                    "question_text": q["question_text"],
                    "question_type": q.get("question_type", "mcq"),
                    "options": q.get("options"),
                    "correct_answer": "PENDING",  # placeholder — admin fills via QuizEditPage
                    "explanation": "",
                    "page_reference": q.get("page_reference"),
                })
            if result:
                return result
            raise ValueError("No questions found")
        except Exception as e:
            logger.warning(f"questions_only attempt {attempt + 1}: {e}")
    raise RuntimeError("questions_only extraction failed after 3 attempts")


# ─── TWO-STEP (SEPARATE ANSWER KEY) ─────────────────────────────────────────

def extract_two_step(
    document_id: int,
    section_start: int,
    section_end: int,
    model_id: str | None,
    api_key: str | None,
    push_step,          # callable(step, message) to report progress
    chunk_pages: int = 50,
) -> tuple[list[dict], list[str]]:
    """
    Two-phase extraction for PDFs with questions in the first half
    and a separate answer key section (e.g. PREP 2013 "Preferred Response:").

    Returns (valid_questions, skipped_list).
    Raises ValueError if answer section not found or no questions matched.
    """
    total = section_end - section_start + 1
    min_answer_start = section_start + max(20, int(total * 0.20))

    # Scan for answer key boundary
    push_step("ai", "Two-step mode: scanning for answer key section…")
    answer_section_start = None
    for scan_p in range(min_answer_start, section_end, 10):
        raw = vector_service.get_pages_text(document_id=document_id, start_page=scan_p,
                                             end_page=min(scan_p + 9, section_end))
        if raw and ("Preferred Response:" in _normalize(raw) or
                    "preferred response:" in raw.lower()):
            answer_section_start = max(section_start, scan_p - 5)
            break

    if not answer_section_start:
        raise ValueError(
            "Two-step mode: could not find a separate answer key section "
            "(no 'Preferred Response:' found after the first 20% of pages). "
            "Try 'Standard' mode instead."
        )

    push_step("ai", f"Answer key section starts around page {answer_section_start}.")

    q_end = answer_section_start - 1
    q_chunks = []
    p = section_start
    while p <= q_end:
        end = min(p + chunk_pages - 1, q_end)
        q_chunks.append((p, end))
        p = end + 1
    n = len(q_chunks)

    # Phase 1 — questions
    raw_questions: list[dict] = []
    for i, (sp, ep) in enumerate(q_chunks, 1):
        push_step("ai", f"Phase 1 – Chunk {i}/{n}: pages {sp}–{ep} (questions)…")
        chunk = vector_service.get_pages_text(document_id=document_id, start_page=sp, end_page=ep)
        if not chunk:
            continue
        try:
            qs = ai_service.extract_questions_no_answers(
                _normalize(chunk), page_info=f"{sp}-{ep}", page_ref=sp,
                model_id=model_id, api_key=api_key,
            )
            raw_questions.extend(qs)
            push_step("ai", f"  Pages {sp}–{ep}: {len(qs)} questions. Total: {len(raw_questions)}.")
        except Exception as e:
            push_step("ai", f"  Pages {sp}–{ep} failed: {e}. Continuing…")

    # Phase 2 — answer key
    push_step("ai", f"Phase 2 – Answer key from pages {answer_section_start}–{section_end}…")
    answer_key: dict = {}
    for ans_start in range(answer_section_start, section_end + 1, chunk_pages):
        ans_end = min(ans_start + chunk_pages - 1, section_end)
        ans_content = vector_service.get_pages_text(document_id=document_id,
                                                     start_page=ans_start, end_page=ans_end)
        if not ans_content:
            continue
        chunk_key = ai_service.extract_answer_key(
            _normalize(ans_content), page_info=f"{ans_start}-{ans_end}",
            model_id=model_id, api_key=api_key,
        )
        answer_key.update(chunk_key)
        push_step("ai", f"  Answer pages {ans_start}–{ans_end}: +{len(chunk_key)}. Total: {len(answer_key)}.")

    push_step("ai", f"Answer key complete: {len(answer_key)} items.")

    # Phase 3 — match
    push_step("ai", "Phase 3 – Matching questions to answers…")
    valid: list[dict] = []
    skipped: list[str] = []
    for q in raw_questions:
        item = q.get("item_number")
        letter = answer_key.get(item) if item else None
        if not letter:
            skipped.append(q.get("question_text", "")[:120])
            continue
        options = q.get("options") or []
        idx = ord(letter.upper()) - ord("A")
        if 0 <= idx < len(options):
            q["correct_answer"] = options[idx]
            valid.append(q)
        else:
            skipped.append(q.get("question_text", "")[:120])

    push_step("ai", f"Matching: {len(valid)} matched, {len(skipped)} unmatched.")
    return valid, skipped


# ─── REGEX MODE ──────────────────────────────────────────────────────────────

REGEX_ANALYSIS_PROMPT = """Look at this PREP exam PDF content and identify the pattern used to mark correct answers.

Describe:
1. The exact text pattern before the correct answer letter (e.g. "Correct Answer:" or "Preferred Response:")
2. Whether answers appear right after each question (inline) or in a separate section at the back
3. A Python regex pattern that would match: the answer indicator + whitespace + the letter (A-E)

Return ONLY JSON:
{{"indicator": "<text before letter>",
 "placement": "inline" | "end_of_doc",
 "regex": "<python regex with one capture group for the letter>",
 "notes": "<any relevant observation>"}}

Sample content (first 30 pages):
{content}"""


def extract_with_regex(
    document_id: int,
    section_start: int,
    section_end: int,
    model_id: str | None,
    api_key: str | None,
    push_step,
    chunk_pages: int = 50,
) -> tuple[list[dict], list[str]]:
    """AI analyses format, generates regex for answer extraction, then applies it."""
    # Sample first 30 pages for analysis
    push_step("ai", "Regex mode: analysing document format…")
    sample = vector_service.get_pages_text(document_id=document_id,
                                            start_page=section_start,
                                            end_page=min(section_start + 29, section_end))
    if not sample:
        raise ValueError("No content found for format analysis")

    prompt = REGEX_ANALYSIS_PROMPT.format(content=_normalize(sample)[:60000])
    analysis_text = ai_service._call_model(prompt, model_id, api_key).strip()
    if analysis_text.startswith("```"):
        analysis_text = analysis_text.split("\n", 1)[1] if "\n" in analysis_text else analysis_text[3:]
        if analysis_text.endswith("```"): analysis_text = analysis_text[:-3]

    analysis = json.loads(analysis_text.strip())
    regex_pattern = analysis.get("regex", r"(?:Correct Answer|Preferred Response)\s*[:\.]?\s*([A-E])")
    placement = analysis.get("placement", "inline")
    push_step("ai", f"Format detected: {placement} answers. Indicator: {analysis.get('indicator','?')!r}. Regex: {regex_pattern!r}")

    # Extract questions using standard no-answer prompt
    push_step("ai", "Extracting questions…")
    q_end = section_end
    if placement == "end_of_doc":
        # Try to find boundary
        for scan_p in range(section_start + 20, section_end, 10):
            raw = vector_service.get_pages_text(document_id=document_id, start_page=scan_p,
                                                 end_page=min(scan_p + 9, section_end))
            if raw:
                norm = _normalize(raw)
                try:
                    if re.search(regex_pattern, norm, re.IGNORECASE):
                        q_end = max(section_start, scan_p - 5)
                        break
                except re.error:
                    break

    raw_questions: list[dict] = []
    q_chunks = []
    p = section_start
    while p <= q_end:
        q_chunks.append((p, min(p + chunk_pages - 1, q_end)))
        p += chunk_pages

    for i, (sp, ep) in enumerate(q_chunks, 1):
        push_step("ai", f"Questions chunk {i}/{len(q_chunks)}: pages {sp}–{ep}…")
        chunk = vector_service.get_pages_text(document_id=document_id, start_page=sp, end_page=ep)
        if not chunk:
            continue
        try:
            qs = ai_service.extract_questions_no_answers(
                _normalize(chunk), page_info=f"{sp}-{ep}", page_ref=sp,
                model_id=model_id, api_key=api_key,
            )
            raw_questions.extend(qs)
        except Exception as e:
            push_step("ai", f"  Chunk {i} failed: {e}")

    push_step("ai", f"{len(raw_questions)} questions extracted. Applying regex to full document for answers…")

    # Apply regex to full document to find item→letter mapping
    answer_map: dict = {}
    for ans_start in range(section_start, section_end + 1, chunk_pages):
        ans_end = min(ans_start + chunk_pages - 1, section_end)
        ans_content = vector_service.get_pages_text(document_id=document_id,
                                                     start_page=ans_start, end_page=ans_end)
        if not ans_content:
            continue
        norm = _normalize(ans_content)
        # Find Item NNN + answer pattern
        combined = r"Item\s+(\d+)[\s\S]{0,300}?" + regex_pattern
        try:
            for m in re.finditer(combined, norm, re.IGNORECASE):
                item_num = m.group(1).lstrip("0") or m.group(1)
                letter = m.group(2).upper() if len(m.groups()) > 1 else ""
                if letter:
                    answer_map[item_num] = letter
        except re.error:
            pass

    push_step("ai", f"Regex found {len(answer_map)} item→answer mappings.")

    # Match
    valid: list[dict] = []
    skipped: list[str] = []
    for q in raw_questions:
        item = q.get("item_number")
        letter = answer_map.get(item) if item else None
        if not letter:
            skipped.append(q.get("question_text", "")[:120])
            continue
        options = q.get("options") or []
        idx = ord(letter.upper()) - ord("A")
        if 0 <= idx < len(options):
            q["correct_answer"] = options[idx]
            valid.append(q)
        else:
            skipped.append(q.get("question_text", "")[:120])

    push_step("ai", f"Matched {len(valid)}, skipped {len(skipped)}.")
    return valid, skipped


# ─── AI DECIDES ──────────────────────────────────────────────────────────────

AI_DECIDE_PROMPT = """You are analysing a PREP medical exam PDF to determine the best extraction strategy.

Sample from first 30 pages:
{sample_start}

---
Sample from last 20 pages:
{sample_end}

Based on these samples, which extraction strategy should be used?

1. standard      — "Correct Answer: X" or "Preferred Response: X" appears right after each question
2. two_step      — questions come first (no answers), then a separate answer key section at the back
3. questions_only — no answer indicators at all (answers unknown)

Return ONLY JSON:
{{"strategy": "standard" | "two_step" | "questions_only",
 "reasoning": "<one sentence>"}}"""


# ─── GENERATE FROM TEXT ──────────────────────────────────────────────────────

GENERATE_PROMPT = """You are a pediatric medical education expert. Read the text below and generate {n} high-quality multiple-choice questions to test comprehension.

Rules:
- The CORRECT answer must be directly supported by the text — do not invent facts
- Create 3 plausible but incorrect distractors using your medical knowledge
- Questions should test understanding of concepts, not just exact word recall
- Prefer clinical application questions over pure recall when the content allows
- Include a 1-2 sentence explanation that cites the key concept from the text
- Spread questions across different parts of the text, not just the first section
- Each option should be a complete, standalone phrase (not "A", "B" labels)

Return ONLY valid JSON (no markdown, no preamble):
{{"questions": [
  {{
    "question_text": "<clear clinical or conceptual question>",
    "question_type": "mcq",
    "options": ["<correct answer text>", "<distractor B>", "<distractor C>", "<distractor D>"],
    "correct_answer": "<exact text of the correct option>",
    "explanation": "<1-2 sentence explanation>",
    "page_reference": {page_ref}
  }}
]}}

Text (pages {page_info}):
{content}"""

QUESTIONS_PER_CHUNK = 8  # target questions per ~50-page chunk


def generate_from_text(content: str, page_info: str, page_ref: int | None,
                       model_id: str | None, api_key: str | None,
                       n: int = QUESTIONS_PER_CHUNK) -> list[dict]:
    """Generate MCQ questions from plain text — correct answers from the document,
    distractors from AI medical knowledge."""
    content = ai_service._truncate_content(content)
    prompt = GENERATE_PROMPT.format(
        content=content,
        page_info=page_info,
        page_ref=page_ref if page_ref is not None else "null",
        n=n,
    )
    for attempt in range(3):
        try:
            text = ai_service._call_model(prompt, model_id, api_key).strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            data = json.loads(text)
            qs = data.get("questions", data) if isinstance(data, dict) else data
            result = []
            for q in qs:
                if not q.get("question_text") or not q.get("correct_answer"):
                    continue
                options = q.get("options") or []
                correct = q["correct_answer"]
                # Validate correct_answer is among options
                if options and correct not in options:
                    # Try to find closest match
                    matches = [o for o in options if correct.lower() in o.lower() or o.lower() in correct.lower()]
                    if matches:
                        correct = matches[0]
                    else:
                        continue  # skip malformed question
                result.append({
                    "question_text": q["question_text"],
                    "question_type": q.get("question_type", "mcq"),
                    "options": options,
                    "correct_answer": correct,
                    "explanation": q.get("explanation", ""),
                    "page_reference": q.get("page_reference"),
                    "item_number": None,
                })
            if result:
                return result
            raise ValueError("No valid questions generated")
        except Exception as e:
            logger.warning(f"generate_from_text attempt {attempt + 1}: {e}")
    raise RuntimeError("generate_from_text failed after 3 attempts")


# ─── AI DECIDES ──────────────────────────────────────────────────────────────

def ai_decide_strategy(
    document_id: int,
    section_start: int,
    section_end: int,
    model_id: str | None,
    api_key: str | None,
) -> str:
    """AI reads samples from start and end of document and decides extraction strategy."""
    sample_start = vector_service.get_pages_text(document_id=document_id,
                                                   start_page=section_start,
                                                   end_page=min(section_start + 29, section_end))
    sample_end = vector_service.get_pages_text(document_id=document_id,
                                                start_page=max(section_start, section_end - 19),
                                                end_page=section_end)
    prompt = AI_DECIDE_PROMPT.format(
        sample_start=_normalize(sample_start or "")[:30000],
        sample_end=_normalize(sample_end or "")[:20000],
    )
    try:
        text = ai_service._call_model(prompt, model_id, api_key).strip()
        if text.startswith("```"):
            text = text.split("\n", 1)[1] if "\n" in text else text[3:]
            if text.endswith("```"): text = text[:-3]
        result = json.loads(text.strip())
        strategy = result.get("strategy", "standard")
        reasoning = result.get("reasoning", "")
        logger.info(f"AI decided: {strategy} — {reasoning}")
        return strategy, reasoning
    except Exception as e:
        logger.warning(f"ai_decide failed: {e}, falling back to standard")
        return "standard", "Fallback to standard due to analysis error"


# ─── FLASHCARD GENERATION ────────────────────────────────────────────────────

FLASHCARD_PROMPT = """You are a pediatric medical education expert. Read the text below and create {n} high-quality flashcards for studying.

Each flashcard has a FRONT (question, term, or concept prompt) and a BACK (answer, definition, or explanation).

Rules:
- Mix question-style fronts ("What is the most common cause of...") and term-style fronts ("Hyperbilirubinemia")
- FRONT should be concise — one sentence or a few words
- BACK should be complete but not verbose — 1-3 sentences with the key facts
- Focus on high-yield facts: diagnostic criteria, treatment protocols, age-specific norms, pathophysiology
- Do NOT repeat the same concept in multiple cards
- Spread cards across different parts of the text
- Each card must be directly supported by the text — do not invent facts

Return ONLY valid JSON (no markdown, no preamble):
{{"cards": [
  {{
    "front": "<question or term>",
    "back": "<answer or definition>",
    "page_reference": {page_ref}
  }}
]}}

Text (pages {page_info}):
{content}"""

FLASHCARDS_PER_CHUNK = 15


def generate_flashcards(content: str, page_info: str, page_ref: int | None,
                        model_id: str | None, api_key: str | None,
                        n: int = FLASHCARDS_PER_CHUNK) -> list[dict]:
    """Generate flashcards from text content using AI."""
    content = ai_service._truncate_content(content)
    prompt = FLASHCARD_PROMPT.format(
        content=content, page_info=page_info,
        page_ref=page_ref if page_ref else "null",
        n=n,
    )
    for attempt in range(3):
        try:
            text = ai_service._call_model(prompt, model_id, api_key).strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            data = json.loads(text)
            cards = data.get("cards", data) if isinstance(data, dict) else data
            result = []
            for c in cards:
                if not c.get("front") or not c.get("back"):
                    continue
                result.append({
                    "front": c["front"].strip(),
                    "back": c["back"].strip(),
                    "page_reference": c.get("page_reference"),
                })
            return result
        except (json.JSONDecodeError, KeyError) as e:
            logger.warning(f"Flashcard generation attempt {attempt + 1} failed: {e}")
            continue
    return []