pdf-quiz-generator/backend/app/tasks/quiz_tasks.py

"""Async quiz extraction task with step-by-step progress reporting via Redis."""
import json
import logging
import time

from app.tasks import celery_app
from app.database import SessionLocal

logger = logging.getLogger(__name__)

EXPIRE_SECONDS = 3600  # keep progress for 1 hour


def _redis():
    import redis
    from app.config import settings
    return redis.from_url(settings.REDIS_URL, decode_responses=True)


def _push_step(r, job_id: str, step: str, message: str):
    key = f"extraction:steps:{job_id}"
    entry = json.dumps({"step": step, "message": message, "ts": time.time()})
    r.rpush(key, entry)
    r.expire(key, EXPIRE_SECONDS)


@celery_app.task(name="extract_quiz", bind=True)
def extract_quiz(
    self,
    job_id: str,
    user_id: int,
    section_id: int,
    title: str,
    mode: str,
    time_limit_minutes: int | None,
    model_id: str | None,
    question_category_id: int | None,
):
    """Background quiz extraction task with live progress reporting."""
    r = _redis()
    r.set(f"extraction:status:{job_id}", "running", ex=EXPIRE_SECONDS)
    # Ensure user job index exists (in case set before task ran)
    r.lpush(f"extraction:user_jobs:{user_id}", job_id)
    r.expire(f"extraction:user_jobs:{user_id}", 86400)

    db = SessionLocal()
    try:
        from app.models.section import Section
        from app.models.pdf_document import PDFDocument
        from app.services import ai_service, vector_service, pdf_service, embedding_service
        from app.models.quiz import Quiz
        from app.models.question import Question
        from app.models.quiz_question_link import QuizQuestionLink
        from app.config import settings
        import os

        _push_step(r, job_id, "start", "Starting extraction…")

        section = db.query(Section).filter(Section.id == section_id).first()
        if not section:
            raise ValueError("Section not found")

        document = db.query(PDFDocument).filter(PDFDocument.id == section.document_id).first()
        if not document:
            raise ValueError("Document not found")

        total_pages = section.end_page - section.start_page + 1
        _push_step(r, job_id, "text", f"Loading text from pages {section.start_page}–{section.end_page} ({total_pages} pages)…")

        # Determine model first (used for chunk logging)
        if model_id:
            from app.models.ai_model_config import AIModelConfig
            config = db.query(AIModelConfig).filter(AIModelConfig.model_id == model_id).first()
            api_key = config.api_key if config and config.api_key else None
            model_name = config.name if config else model_id
        else:
            model_id, api_key = ai_service.get_model_for_task(db, "extraction")
            model_name = model_id

        # For large page ranges, process in chunks of 50 pages to avoid truncation
        CHUNK_PAGES = 50
        import json as _json

        if total_pages <= CHUNK_PAGES:
            chunks = [(section.start_page, section.end_page)]
        else:
            chunks = []
            p = section.start_page
            while p <= section.end_page:
                end = min(p + CHUNK_PAGES - 1, section.end_page)
                chunks.append((p, end))
                p = end + 1

        n_chunks = len(chunks)
        if n_chunks > 1:
            _push_step(r, job_id, "text", f"Large section: splitting into {n_chunks} chunks of up to {CHUNK_PAGES} pages each.")

        def _normalize_ocr(text: str) -> str:
            """Normalize common OCR artifacts in PREP PDFs."""
            return (text
                    .replace("Pref erred", "Preferred")
                    .replace("Pre ferred", "Preferred")
                    .replace("Prefer red", "Preferred")
                    .replace("ltem", "Item")
                    .replace("ltcm", "Item"))

        # --- Detect PDF format ---
        # First check if document has INLINE correct answers (standard PREP 2012 format).
        # If yes → always use standard extraction, even if "Preferred Response:" appears
        # later in explanations.
        first_50_content = vector_service.get_pages_text(
            document_id=section.document_id,
            start_page=section.start_page,
            end_page=min(section.start_page + 49, section.end_page),
        )
        has_inline_answers = bool(first_50_content and (
            "Correct Answer:" in first_50_content or
            "correct answer:" in first_50_content.lower()
        ))

        # Only look for end-of-document answer key if there are NO inline answers
        answer_section_start = None
        if not has_inline_answers:
            scan_step = 10
            for scan_p in range(section.start_page, section.end_page, scan_step):
                scan_end = min(scan_p + scan_step - 1, section.end_page)
                raw_chunk = vector_service.get_pages_text(
                    document_id=section.document_id, start_page=scan_p, end_page=scan_end,
                )
                if not raw_chunk:
                    continue
                scan_chunk = _normalize_ocr(raw_chunk)
                if "Preferred Response:" in scan_chunk or "preferred response:" in scan_chunk.lower():
                    answer_section_start = max(section.start_page, scan_p - 5)
                    break

        has_end_answer_key = answer_section_start is not None
        if has_inline_answers:
            _push_step(r, job_id, "ai", "Detected inline answer format (Correct Answer: X). Using standard extraction.")
        elif has_end_answer_key:
            pass  # logged below

        if has_end_answer_key:
            _push_step(r, job_id, "ai", f"Detected separate answer key section starting around page {answer_section_start}. Using two-phase extraction.")
            # Restrict question chunks to pages strictly BEFORE the answer section
            # Cap the end page so no chunk bleeds into the answer section
            q_end = answer_section_start - 1
            chunks = [(s, min(e, q_end)) for s, e in chunks if s <= q_end]
            if not chunks:
                chunks = [(section.start_page, q_end)]
            n_chunks = len(chunks)  # update display count after filtering

        all_valid_questions = []
        all_skipped = []

        if has_end_answer_key:
            # === PHASE 1: Extract questions (allow null correct_answer) ===
            raw_questions = []
            for chunk_idx, (start_p, end_p) in enumerate(chunks, 1):
                _push_step(r, job_id, "ai", f"Phase 1 – Chunk {chunk_idx}/{n_chunks}: pages {start_p}–{end_p} (questions)…")
                chunk_content = vector_service.get_pages_text(
                    document_id=section.document_id, start_page=start_p, end_page=end_p,
                )
                if not chunk_content:
                    continue
                try:
                    chunk_qs = ai_service.extract_questions_no_answers(
                        _normalize_ocr(chunk_content), page_info=f"{start_p}-{end_p}",
                        page_ref=start_p, model_id=model_id, api_key=api_key,
                    )
                    raw_questions.extend(chunk_qs)
                    _push_step(r, job_id, "ai", f"  Pages {start_p}–{end_p}: {len(chunk_qs)} questions found. Total: {len(raw_questions)}.")
                except Exception as e:
                    _push_step(r, job_id, "ai", f"  Pages {start_p}–{end_p} failed: {e}. Continuing…")

            # === PHASE 2: Extract answer key from the answer section ===
            _push_step(r, job_id, "ai", f"Phase 2 – Extracting answer key from pages {answer_section_start}–{section.end_page}…")
            # Process answer section in chunks too (it may be large)
            answer_key: dict = {}
            for ans_start in range(answer_section_start, section.end_page + 1, CHUNK_PAGES):
                ans_end = min(ans_start + CHUNK_PAGES - 1, section.end_page)
                answer_content = vector_service.get_pages_text(
                    document_id=section.document_id,
                    start_page=ans_start,
                    end_page=ans_end,
                )
                if not answer_content:
                    continue
                chunk_key = ai_service.extract_answer_key(
                    _normalize_ocr(answer_content), page_info=f"{ans_start}-{ans_end}",
                    model_id=model_id, api_key=api_key,
                )
                answer_key.update(chunk_key)
                _push_step(r, job_id, "ai", f"  Answer pages {ans_start}–{ans_end}: +{len(chunk_key)} answers. Total: {len(answer_key)}.")
            _push_step(r, job_id, "ai", f"  Answer key complete: {len(answer_key)} items.")

            # === PHASE 3: Match questions to answers ===
            _push_step(r, job_id, "ai", "Phase 3 – Matching questions to answers…")
            for q in raw_questions:
                item_num = q.get("item_number")
                if not item_num:
                    # No item number — skip
                    all_skipped.append(q.get("question_text", "")[:120])
                    continue

                letter = answer_key.get(item_num)
                if not letter:
                    all_skipped.append(q.get("question_text", "")[:120])
                    continue

                # Resolve letter → full option text
                options = q.get("options") or []
                letter_idx = ord(letter.upper()) - ord("A")
                if 0 <= letter_idx < len(options):
                    q["correct_answer"] = options[letter_idx]
                    all_valid_questions.append(q)
                else:
                    all_skipped.append(q.get("question_text", "")[:120])

            _push_step(r, job_id, "ai", f"Matching complete: {len(all_valid_questions)} matched, {len(all_skipped)} unmatched.")

        else:
            # === STANDARD: inline correct answer format ===
            for chunk_idx, (start_p, end_p) in enumerate(chunks, 1):
                if n_chunks > 1:
                    _push_step(r, job_id, "ai", f"Chunk {chunk_idx}/{n_chunks}: pages {start_p}–{end_p} → {model_name}…")
                else:
                    _push_step(r, job_id, "ai", f"Sending pages {start_p}–{end_p} to {model_name}…")

                chunk_content = vector_service.get_pages_text(
                    document_id=section.document_id, start_page=start_p, end_page=end_p,
                )
                if not chunk_content:
                    _push_step(r, job_id, "ai", f"  No text found for pages {start_p}–{end_p}, skipping.")
                    continue
                try:
                    chunk_data = ai_service.extract_questions(
                        chunk_content, page_info=f"{start_p}-{end_p}",
                        page_ref=start_p, model_id=model_id, api_key=api_key,
                    )
                    chunk_skipped = []
                    if chunk_data and chunk_data[0].get("skipped"):
                        chunk_skipped = chunk_data[0].pop("skipped")
                    chunk_valid = [q for q in chunk_data if q.get("correct_answer")]
                    all_valid_questions.extend(chunk_valid)
                    all_skipped.extend(chunk_skipped)
                    _push_step(r, job_id, "ai", f"  Pages {start_p}–{end_p}: {len(chunk_valid)} questions extracted{f', {len(chunk_skipped)} skipped' if chunk_skipped else ''}. Total so far: {len(all_valid_questions)}.")
                except Exception as e:
                    _push_step(r, job_id, "ai", f"  Pages {start_p}–{end_p} failed: {e}. Continuing…")

        valid_questions = all_valid_questions
        skipped = all_skipped

        _push_step(r, job_id, "ai", f"Extraction complete: {len(valid_questions)} total valid questions{f', {len(skipped)} skipped' if skipped else ''}.")

        if not valid_questions:
            raise ValueError("No valid questions extracted. The AI could not find questions with correct answers in this page range.")

        # Extract images
        _push_step(r, job_id, "images", "Extracting question images…")
        file_path = os.path.join(settings.UPLOAD_DIR, document.filename)
        page_images = {}
        if os.path.exists(file_path):
            try:
                page_images = pdf_service.extract_all_images(
                    file_path, document.id, section.start_page, section.end_page
                )
            except Exception as e:
                _push_step(r, job_id, "images", f"Image extraction skipped: {e}")

        # Create quiz
        quiz = Quiz(
            section_id=section_id,
            user_id=user_id,
            title=title,
            questions_count=len(valid_questions),
            mode=mode,
            time_limit_minutes=time_limit_minutes,
            skipped_questions=_json.dumps(skipped) if skipped else None,
        )
        db.add(quiz)
        db.flush()

        _push_step(r, job_id, "save", f"Saving {len(valid_questions)} questions and generating embeddings…")

        for pos, q in enumerate(valid_questions):
            page_ref = q.get("page_reference")
            image_path = None
            if page_ref and page_ref in page_images and page_images[page_ref]:
                image_path = page_images[page_ref].pop(0)
                if not page_images[page_ref]:
                    del page_images[page_ref]

            question = Question(
                source_quiz_id=quiz.id,
                question_category_id=question_category_id,
                question_text=q["question_text"],
                question_type=q["question_type"],
                options=q.get("options"),
                correct_answer=q["correct_answer"],
                explanation=q.get("explanation", ""),
                page_reference=page_ref,
                image_path=image_path,
            )
            db.add(question)
            db.flush()
            db.add(QuizQuestionLink(quiz_id=quiz.id, question_id=question.id, position=pos))
            try:
                embedding_service.embed_question(question)
            except Exception:
                pass

        db.commit()
        db.refresh(quiz)

        _push_step(r, job_id, "done", f"Quiz ready! {len(valid_questions)} questions extracted and saved.")

        r.set(f"extraction:status:{job_id}", "completed", ex=EXPIRE_SECONDS)
        r.set(f"extraction:quiz_id:{job_id}", str(quiz.id), ex=EXPIRE_SECONDS)
        return quiz.id

    except Exception as e:
        logger.exception(f"Quiz extraction failed for job {job_id}")
        _push_step(r, job_id, "error", f"Extraction failed: {e}")
        r.set(f"extraction:status:{job_id}", "failed", ex=EXPIRE_SECONDS)
        r.set(f"extraction:error:{job_id}", str(e)[:500], ex=EXPIRE_SECONDS)
        raise
    finally:
        db.close()