MathApp/rag/services/ingest.py

"""
هماهنگ‌کننده‌ی ingestion: متن/PDF درس → chunk/خلاصه/embedding → ذخیره در
DB (LessonChunk، منبعِ حقیقت) + upsert به Qdrant (index).

نامِ collection و کلیدهای payload عمداً همون قراردادِ Rag_backend را دنبال
می‌کنن تا با دادهٔ موجودِ تیمِ دیگه سازگار بمونه.
"""
import logging
import os
import tempfile
from pathlib import Path

from django.conf import settings
from django.db import transaction

from ..models import IngestedPage, LessonChunk
from .chunking import chunk_and_summarize
from .config import build_collection_name, resolve_grade_title
from .embedding import embed_text
from .known_collections import is_math_subject, resolve_math_chapter
from . import pdf as pdf_service
from .qdrant_store import ensure_collection, upsert_points
from .textbooks import Textbook, get_textbook

logger = logging.getLogger(__name__)

# قیمتِ هر یک میلیون توکن (دلار) برای گزارشِ هزینه‌ی ingest. اندازه‌گیریِ
# ۱۱۱ صفحه‌ی واقعیِ ریاضی نهم: ۱٬۴۲۵ توکن ورودی و ~۴۰۰ توکن خروجی در هر صفحه،
# یعنی حدود ۰٫۰۰۷۵ دلار — نه ۰٫۰۱۷ که اول تخمین زده بودیم.
MODEL_PRICE_PER_MTOK = {
    'gpt-4o': (2.50, 10.00),
    'gpt-4o-mini': (0.15, 0.60),
}


def estimate_page_cost(model: str) -> float:
    """هزینه‌ی تقریبیِ یک صفحه بر پایه‌ی مصرفِ اندازه‌گیری‌شده‌ی واقعی."""
    price_in, price_out = MODEL_PRICE_PER_MTOK.get(model, (2.50, 10.00))
    return 1425 / 1e6 * price_in + 400 / 1e6 * price_out


def _payload(chunk: LessonChunk) -> dict:
    """
    payloadِ نقطه‌ی Qdrant.

    ⚠️ کلیدِ `chapter` (رشته‌ی کاملِ فصل) حتماً باید نوشته بشه: جستجوی ریاضی در
    retrieval.py دقیقاً روی همین کلید فیلتر می‌کنه (نه `chapter_number`)، پس
    بدونش هر چیزی که ingest کنیم برای همیشه نامرئی می‌مونه — بدونِ هیچ خطایی.

    دو شکلِ payload داریم و باید هر کدام دقیقاً مثلِ نقاطِ موجودِ همان نوع
    بمانند (نقاطِ «ریاضی هفتم» و «علوم تجربی سوم» را مقایسه کنید):

    - ریاضی: grade / chapter / type / keywords / generalSolution / problem /
      solution — و عمداً **بدونِ `summary`**. اگر summary بنویسیم،
      `retrieval._payload_text` آن را به generalSolution ترجیح می‌دهد و
      به‌جای راه‌حل، متنِ صافِ «Type: ... Keywords: ...» را به‌عنوانِ context
      به مدل می‌دهیم — یعنی نقاطِ ما بی‌سروصدا بدتر از نقاطِ تیمِ دیگر جواب
      می‌دهند. `lesson_title` هم به همین دلیل نمی‌آید: `_payload_title`
      ترجیحش می‌دهد و عنوانِ عمومیِ فصل جایِ نامِ نوعِ مسئله را می‌گیرد.
    - درس‌های متنی: lesson_number + summary، همان‌طور که بود.

    کلیدهای دفتری (chapter_number / chunk_order / source_page) در هر دو حالت
    می‌آیند؛ روی فیلتر و متنِ بازیابی‌شده اثری ندارند و ردیابیِ «این نقطه از
    کدام صفحه‌ی کتاب آمده» را ممکن می‌کنند.
    """
    # منشأ را از `problem_type` تشخیص می‌دهیم، نه از پر بودنِ راه‌حل: صفحه‌ای
    # که vision فقط `type` و `keywords` از آن درآورده هم ردیفِ ساختاریافته است،
    # ولی با شرطِ قبلی «متنی» حساب می‌شد و `summary` می‌گرفت — یعنی همان آلودگیِ
    # «Type: … Keywords: …» که این تفکیک برای جلوگیری از آن ساخته شده بود.
    is_structured = bool(
        chunk.problem_type or chunk.general_solution or chunk.problem or chunk.solution
    )

    payload: dict = {
        'lesson_number': chunk.lesson_number,
        'chapter_number': chunk.chapter,
        'chunk_order': chunk.order,
    }
    if chunk.lesson_id is not None:
        payload['lesson_id'] = chunk.lesson_id
    if chunk.source_page is not None:
        payload['source_page'] = chunk.source_page
    if not is_structured:
        payload['lesson_title'] = chunk.title
        payload['summary'] = chunk.summary

    for key, value in (
        ('chapter', chunk.chapter_title),
        ('grade', chunk.grade_title),
        ('type', chunk.problem_type),
        ('keywords', chunk.keywords),
        ('generalSolution', chunk.general_solution),
        ('problem', chunk.problem),
        ('solution', chunk.solution),
    ):
        if value:
            payload[key] = value
    return payload


def _persist_and_upsert(rows: list[dict], *, lesson_id, grade_id, subject,
                        lesson_number, chapter, title, collection) -> list[LessonChunk]:
    ensure_collection(collection)

    # embedding قبل از باز کردنِ transaction: هر فراخوانی چند صد میلی‌ثانیه
    # شبکه است و نگه داشتنِ قفلِ نوشتنِ SQLite در آن مدت، بقیه‌ی درخواست‌ها را
    # پشتِ «database is locked» می‌خواباند.
    vectors = [embed_text(row['summary']) for row in rows]

    created: list[LessonChunk] = []
    with transaction.atomic():
        for row, vector in zip(rows, vectors):
            created.append(LessonChunk.objects.create(
                lesson_id=lesson_id,
                grade_id=grade_id,
                subject=subject,
                lesson_number=lesson_number,
                chapter=chapter,
                order=row['order'],
                title=title,
                original_text=row['original_text'],
                summary=row['summary'],
                embedding_vector=vector,
                collection=collection,
                source_page=row.get('source_page'),
                # فیلدهای ساختاری فقط از مسیرِ PDF/vision می‌آن؛ ingest_text
                # اون‌ها رو نداره و خالی می‌مونن (payload هم اون‌وقت نمی‌نویسدشون).
                **row.get('structured', {}),
            ))

    # ⚠️ upsert بعد از commit است تا نقطه‌ای در Qdrant بدونِ ردیفِ متناظر در DB
    # نماند. ولی اگر همین upsert شکست بخورد، ردیف‌ها commit شده‌اند و
    # `_done_pages` آن صفحه را «انجام‌شده» می‌بیند — یعنی chunkها برای همیشه در
    # DB می‌مانند و هیچ‌وقت به Qdrant نمی‌رسند، بی‌آنکه کسی بفهمد. پس در صورتِ
    # خطا ردیف‌های همین دسته را پس می‌گیریم تا صفحه واقعاً ناتمام بماند و اجرای
    # بعدی دوباره سراغش برود.
    try:
        upsert_points(collection, [
            {'id': c.qdrant_point_id, 'vector': v, 'payload': _payload(c)}
            for c, v in zip(created, vectors)
        ])
    except Exception:
        LessonChunk.objects.filter(pk__in=[c.pk for c in created]).delete()
        logger.exception(
            'rag: upsert به collection %r شکست خورد؛ %d ردیفِ همین دسته پس گرفته شد',
            collection, len(created),
        )
        raise
    return created


def ingest_text(*, lesson_id: int, grade_id, subject: str, lesson_number: int,
                text: str, chapter: int | None = None, title: str = '') -> dict:
    grade_title = resolve_grade_title(grade_id)
    collection = build_collection_name(subject, grade_title)
    title = title or f'lesson_{lesson_number}'

    rows = chunk_and_summarize(text)
    if not rows:
        return {'status': 'empty', 'chunks_created': 0, 'collection': collection}

    # برای ریاضی، فیلترِ جستجو رشته‌ی کاملِ فصله؛ از شماره‌ی فصل درش می‌آریم تا
    # محتوایی که این‌طوری ingest می‌شه هم قابلِ پیدا شدن باشه. اگه فصل ناشناخته
    # باشه (پایه‌ای که هنوز در known_collections نیست) خالی می‌مونه.
    structured = {'grade_title': grade_title}
    if is_math_subject(subject) and chapter is not None:
        structured['chapter_title'] = resolve_math_chapter(subject, grade_title, chapter) or ''
    for row in rows:
        row.setdefault('structured', structured)

    created = _persist_and_upsert(
        rows, lesson_id=lesson_id, grade_id=grade_id, subject=subject,
        lesson_number=lesson_number, chapter=chapter, title=title, collection=collection,
    )
    return {
        'status': 'success',
        'collection': collection,
        'chunks_created': len(created),
        'chunk_ids': [c.id for c in created],
    }


def ingest_pdf(*, lesson_id: int, grade_id, subject: str, lesson_number: int,
               pdf_bytes: bytes, chapter: int | None = None, title: str = '',
               chapter_name: str = '') -> dict:
    """
    هر صفحه‌ی PDF → عکس → استخراجِ انواعِ مسئله با vision → هر problemType یک
    chunk. برای درس‌های متنی به‌جای این از ingest_text استفاده کن.
    """
    grade_title = resolve_grade_title(grade_id)
    collection = build_collection_name(subject, grade_title)
    title = title or f'lesson_{lesson_number}'
    chapter_name = chapter_name or title

    fd, tmp_path = tempfile.mkstemp(suffix='.pdf')
    try:
        with os.fdopen(fd, 'wb') as fh:
            fh.write(pdf_bytes)

        start_page, detected_grade = pdf_service.detect_grade_and_start_page(tmp_path)
        images = pdf_service.extract_page_images(tmp_path, first_page=start_page)

        rows: list[dict] = []
        order = 0
        for image in images:
            try:
                parsed = pdf_service.extract_problem_types(image, detected_grade or grade_title, chapter_name)
            except Exception:
                logger.exception('rag.ingest_pdf: تحلیلِ یک صفحه شکست خورد؛ رد شد')
                continue
            page_chapter = (parsed.get('chapter') or '').strip() or chapter_name
            page_grade = (parsed.get('grade') or '').strip() or (detected_grade or grade_title)
            for problem_type in parsed.get('problemTypes', []) or []:
                text = pdf_service.problem_type_to_text(problem_type)
                if not text:
                    continue
                order += 1
                example = problem_type.get('exampleFromImage', {}) or {}
                rows.append({
                    'order': order,
                    'original_text': text,
                    'summary': text,
                    # ساختار رو نگه می‌داریم، نه فقط متنِ صاف‌شده‌اش را: هم
                    # `chapter` برای فیلترِ جستجو لازمه، هم problem/solution
                    # بعداً ورودیِ ساختِ سؤال‌های چهارگزینه‌ای می‌شن.
                    'structured': {
                        'chapter_title': page_chapter,
                        'grade_title': page_grade,
                        'problem_type': (problem_type.get('type') or '').strip(),
                        'keywords': problem_type.get('keywords') or [],
                        'general_solution': (problem_type.get('generalSolution') or '').strip(),
                        'problem': (example.get('problem') or '').strip(),
                        'solution': (example.get('solution') or '').strip(),
                    },
                })
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass

    if not rows:
        return {'status': 'empty', 'chunks_created': 0, 'collection': collection}

    created = _persist_and_upsert(
        rows, lesson_id=lesson_id, grade_id=grade_id, subject=subject,
        lesson_number=lesson_number, chapter=chapter, title=title, collection=collection,
    )
    return {
        'status': 'success',
        'collection': collection,
        'chunks_created': len(created),
        'chunk_ids': [c.id for c in created],
    }


# ─── ingestِ کتابِ کامل ────────────────────────────────────────────────────
#
# `ingest_pdf` بالا یک PDF را یک‌تکه و زیرِ یک نامِ فصل می‌گیرد؛ برای کتابِ
# درسیِ کامل (چند فصل، ۱۵۰+ صفحه) سه چیز کم دارد و همین‌ها اینجا اضافه شده‌اند:
#   ۱. فصل‌بندی — بازه‌ی صفحه‌ی هر فصل از rag/services/textbooks.py می‌آید تا
#      هر نقطه زیرِ فصلِ درستِ خودش بنشیند (فیلترِ retrieval روی همین است).
#   ۲. ازسرگیری — صفحه‌ای که قبلاً chunk ساخته دوباره پردازش نمی‌شود، پس قطعیِ
#      شبکه وسطِ کار پول و زمانِ رفته را دور نمی‌ریزد و نقطه‌ی تکراری نمی‌سازد.
#   ۳. صفحه‌به‌صفحه بودن — هم مصرفِ حافظه ثابت می‌ماند، هم هر صفحه بلافاصله
#      ذخیره می‌شود، هم خطای یک صفحه بقیه را از کار نمی‌اندازد.


def _textbook_path(book: Textbook) -> Path:
    root = Path(getattr(settings, 'TEXTBOOK_DIR', Path(settings.BASE_DIR) / 'data' / 'textbooks'))
    return root / book.filename


def _done_pages(collection: str, chapter_title: str) -> set[int]:
    """
    صفحه‌هایی که نباید دوباره پردازش شوند.

    هم نشانه‌های IngestedPage (شاملِ صفحه‌های بی‌مسئله) و هم chunkهای موجود
    را می‌گیرد؛ دومی برای دادهٔ ingest‌شده پیش از افزودنِ جدولِ نشانه لازم است.
    """
    marked = set(
        IngestedPage.objects
        .filter(collection=collection, chapter_title=chapter_title)
        .values_list('page', flat=True)
    )
    chunked = set(
        LessonChunk.objects
        .filter(collection=collection, chapter_title=chapter_title, source_page__isnull=False)
        .values_list('source_page', flat=True)
    )
    return marked | chunked


def ingest_book(*, subject: str, grade: str, chapters=None, dpi: int = 200,
                max_pages: int | None = None, resume: bool = True,
                progress=None) -> dict:
    """
    یک کتابِ درسیِ کامل از data/textbooks/ را فصل‌به‌فصل ingest می‌کند.

    chapters   : فهرستِ شماره‌ی فصل‌ها؛ None یعنی همه‌ی فصل‌های کتاب.
    max_pages  : سقفِ تعدادِ صفحه‌ی پردازش‌شده در این اجرا (برای کنترلِ هزینه).
    resume     : صفحه‌هایی که قبلاً chunk ساخته‌اند دوباره پردازش نشوند.
    progress   : callback(dict) برای گزارشِ زنده به فرمانِ مدیریتی.

    خروجی: خلاصه‌ی اجرا؛ هیچ استثنایی برای خطای یک صفحه بالا نمی‌آید — صفحه‌ی
    خطادار در `failed_pages` گزارش می‌شود و اجرای بعدی دوباره سراغش می‌رود.
    """
    book = get_textbook(subject, grade)
    if book is None:
        raise ValueError(f'کتابی برای «{subject} {grade}» در textbooks.py ثبت نشده است.')

    # کتابِ بدونِ filename (ششم/هفتم — دادهٔ واقعی‌شان روی Qdrantِ تیمِ دیگر
    # است) از این مسیر نمی‌آید. بدونِ این بررسی، مسیر به خودِ پوشه‌ی
    # data/textbooks/ می‌رسید که وجود دارد، پس `exists()` رد می‌شد و کار
    # بی‌سروصدا جلو می‌رفت تا جایی که pdf2image روی یک پوشه شکست بخورد.
    if not book.filename:
        raise FileNotFoundError(
            f'«{book.key}» فایلِ PDF ندارد و از مسیرِ ingest نمی‌آید. '
            f'برای واردکردنِ دادهٔ موجودش: '
            f'python manage.py rag_import_collection --collection "{book.key}"'
        )

    path = _textbook_path(book)
    if not path.exists() or not path.is_file():
        raise FileNotFoundError(f'فایلِ کتاب پیدا نشد: {path}')

    collection = build_collection_name(book.subject, book.grade)
    ensure_collection(collection)

    wanted = list(book.chapters) if chapters is None else [
        c for n in chapters if (c := book.chapter(n)) is not None
    ]

    def emit(**event):
        if progress is not None:
            progress(event)

    totals = {'pages_done': 0, 'pages_skipped': 0, 'chunks_created': 0,
              'prompt_tokens': 0, 'completion_tokens': 0}
    failed_pages: list[int] = []
    budget = max_pages

    for chapter in wanted:
        first_pdf, last_pdf = book.pdf_pages(chapter)
        already = _done_pages(collection, chapter.title) if resume else set()
        emit(kind='chapter', chapter=chapter.number, title=chapter.title,
             pages=(first_pdf, last_pdf), already=len(already))

        # order در هر فصل ادامه‌دار می‌ماند تا اجرای دوم شماره‌ها را از نو نزند.
        order = LessonChunk.objects.filter(
            collection=collection, chapter_title=chapter.title,
        ).count()

        for page in range(first_pdf, last_pdf + 1):
            if page in already:
                totals['pages_skipped'] += 1
                continue
            if budget is not None and budget <= 0:
                emit(kind='budget_reached')
                return _book_result(book, collection, totals, failed_pages, partial=True)

            try:
                images = pdf_service.extract_page_images(
                    str(path), first_page=page, last_page=page, dpi=dpi,
                )
                if not images:
                    raise RuntimeError('صفحه به تصویر تبدیل نشد')
                parsed = pdf_service.extract_problem_types(
                    images[0], book.grade_label, chapter.title,
                )
            except Exception as exc:  # noqa: BLE001
                logger.exception('rag.ingest_book: صفحه %s شکست خورد', page)
                failed_pages.append(page)
                emit(kind='page_failed', page=page, error=str(exc))
                if budget is not None:
                    budget -= 1
                continue

            page_usage = parsed.pop('_usage', None) or {}
            totals['prompt_tokens'] += int(page_usage.get('promptTokens') or 0)
            totals['completion_tokens'] += int(page_usage.get('completionTokens') or 0)

            rows = []
            for problem_type in parsed.get('problemTypes', []) or []:
                text = pdf_service.problem_type_to_text(problem_type)
                if not text:
                    continue
                order += 1
                example = problem_type.get('exampleFromImage', {}) or {}
                rows.append({
                    'order': order,
                    'source_page': page,
                    'original_text': text,
                    'summary': text,
                    'structured': {
                        # عنوانِ فصل از فهرستِ خودمان می‌آید، نه از چیزی که
                        # مدل در تصویر دیده: فیلترِ retrieval روی رشته‌ی دقیق
                        # است و یک حرفِ جابه‌جا کلِ فصل را نامرئی می‌کند.
                        'chapter_title': chapter.title,
                        'grade_title': book.grade_label,
                        'problem_type': (problem_type.get('type') or '').strip(),
                        'keywords': problem_type.get('keywords') or [],
                        'general_solution': (problem_type.get('generalSolution') or '').strip(),
                        'problem': (example.get('problem') or '').strip(),
                        'solution': (example.get('solution') or '').strip(),
                    },
                })

            if rows:
                try:
                    created = _persist_and_upsert(
                        rows, lesson_id=None, grade_id=None, subject=book.subject,
                        lesson_number=chapter.number, chapter=chapter.number,
                        title=chapter.title, collection=collection,
                    )
                except Exception as exc:  # noqa: BLE001
                    logger.exception('rag.ingest_book: ذخیره‌ی صفحه %s شکست خورد', page)
                    failed_pages.append(page)
                    emit(kind='page_failed', page=page, error=str(exc))
                    if budget is not None:
                        budget -= 1
                    continue
                totals['chunks_created'] += len(created)

            # نشانه بعد از ذخیره‌ی موفق زده می‌شود، نه قبلش: اگر ذخیره شکست
            # بخورد صفحه باید در اجرای بعدی دوباره دیده شود.
            IngestedPage.objects.update_or_create(
                collection=collection, page=page,
                defaults={'chapter_title': chapter.title, 'chunk_count': len(rows)},
            )
            totals['pages_done'] += 1
            if budget is not None:
                budget -= 1
            emit(kind='page', page=page, chapter=chapter.number, chunks=len(rows))

    return _book_result(book, collection, totals, failed_pages, partial=False)


def _book_result(book: Textbook, collection: str, totals: dict,
                 failed_pages: list[int], *, partial: bool) -> dict:
    model = settings.OPENAI_VISION_MODEL
    price_in, price_out = MODEL_PRICE_PER_MTOK.get(model, (2.50, 10.00))
    cost = (
        totals['prompt_tokens'] / 1e6 * price_in
        + totals['completion_tokens'] / 1e6 * price_out
    )
    return {
        'status': 'partial' if partial or failed_pages else 'success',
        'book': book.key,
        'collection': collection,
        'pages_processed': totals['pages_done'],
        'pages_skipped': totals['pages_skipped'],
        'chunks_created': totals['chunks_created'],
        'failed_pages': failed_pages,
        'prompt_tokens': totals['prompt_tokens'],
        'completion_tokens': totals['completion_tokens'],
        'cost_usd': round(cost, 4),
        'model': model,
    }