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.
"""
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>"],
    "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.
- 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>"}}"""


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"