MathApp/rag/management/commands/rag_ingest_book.py

"""
ingestِ یک کتابِ درسیِ کامل از data/textbooks/ به Qdrant.

نمونه:
    python manage.py rag_ingest_book --subject ریاضی --grade نهم --plan
    python manage.py rag_ingest_book --subject ریاضی --grade نهم --chapters 1 --max-pages 3
    python manage.py rag_ingest_book --subject ریاضی --grade نهم

هر اجرا از جایی که مانده ادامه می‌دهد (صفحه‌های ثبت‌شده رد می‌شوند)، پس
قطع کردنش با Ctrl-C بی‌خطر است و اجرای دوباره کارِ رفته را دوباره نمی‌خرد.
"""
from django.core.management.base import BaseCommand, CommandError

from rag.services.clients import RagNotConfigured
from rag.services.config import build_collection_name, rag_is_configured
from rag.services.ingest import (_done_pages, estimate_page_cost, ingest_book,
                                 pages_needing_refresh)
from rag.services.textbooks import TEXTBOOKS, get_textbook
from rag.models import IngestedPage, LessonChunk


class Command(BaseCommand):
    help = 'ingestِ کتابِ درسی (PDF) به LessonChunk + Qdrant، فصل‌به‌فصل و قابلِ ازسرگیری.'

    def add_arguments(self, parser):
        parser.add_argument('--subject', default='ریاضی')
        parser.add_argument('--grade', required=True, help='مثلاً: نهم')
        parser.add_argument('--chapters', nargs='*', type=int, default=None,
                            help='شماره‌ی فصل‌ها؛ خالی یعنی همه.')
        parser.add_argument('--max-pages', type=int, default=None,
                            help='سقفِ صفحه در این اجرا (کنترلِ هزینه).')
        parser.add_argument('--dpi', type=int, default=200)
        parser.add_argument('--no-resume', action='store_true',
                            help='صفحه‌های قبلاً ثبت‌شده را هم دوباره پردازش کن (نقطه‌ی تکراری می‌سازد).')
        parser.add_argument('--plan', action='store_true',
                            help='فقط نشان بده چه چیزی پردازش می‌شود؛ هیچ فراخوانیِ پولی نکن.')
        parser.add_argument('--refresh', choices=['figures', 'unjudgeable', 'both', 'all'],
                            default=None,
                            help='صفحه‌های معیوب را دوباره بخوان: دادهٔ قبلیِ هر صفحه '
                                 'درست پیش از خواندنِ دوباره پاک می‌شود. '
                                 'figures=ارجاع به شکل/جدولِ نوشته‌نشده، '
                                 'unjudgeable=رأیِ «قابلِ قضاوت نیست»، both=هر دو، '
                                 'all=همه‌ی صفحه‌های فصل‌های انتخاب‌شده.')
        parser.add_argument('--refresh-pages', default='',
                            help='به‌جای انتخابِ خودکار، همین صفحه‌های PDF را بازخوانی کن '
                                 '(مثلاً: 12,15,20-24).')
        parser.add_argument('--yes', action='store_true',
                            help='بدونِ پرسشِ تأیید ادامه بده (بازخوانی دادهٔ موجود را پاک می‌کند).')

    def handle(self, *args, **opts):
        subject, grade = opts['subject'], opts['grade']
        book = get_textbook(subject, grade)
        if book is None:
            known = '، '.join(sorted(TEXTBOOKS)) or '—'
            raise CommandError(f'کتابی برای «{subject} {grade}» ثبت نشده. موجود: {known}')
        if not book.filename:
            raise CommandError(
                f'«{book.key}» فایلِ PDF ندارد؛ دادهٔ واقعی‌اش از قبل روی Qdrant است.\n'
                f'برای واردکردنش:  python manage.py rag_import_collection '
                f'--collection "{book.key}"'
            )

        collection = build_collection_name(book.subject, book.grade)
        wanted = list(book.chapters) if not opts['chapters'] else [
            c for n in opts['chapters'] if (c := book.chapter(n)) is not None
        ]
        if not wanted:
            raise CommandError('هیچ فصلی با شماره‌های داده‌شده پیدا نشد.')

        refresh = self._resolve_refresh(opts, collection, book, wanted)

        if opts['plan']:
            return self._plan(book, collection, wanted, refresh)

        if refresh and not opts['yes'] and not self._confirm(collection, refresh):
            self.stdout.write('لغو شد؛ هیچ‌چیزی پاک یا خرج نشد.')
            return

        if not rag_is_configured():
            raise CommandError(
                'RAG پیکربندی نشده (OPENAI_API_KEY / QDRANT_API_KEY خالی است یا RAG_FORCE_MOCK روشن است).'
            )

        self.stdout.write(self.style.MIGRATE_HEADING(
            f'ingest «{book.key}» → collection «{collection}»'
        ))
        try:
            result = ingest_book(
                subject=book.subject, grade=book.grade,
                chapters=[c.number for c in wanted],
                dpi=opts['dpi'], max_pages=opts['max_pages'],
                resume=not opts['no_resume'], refresh_pages=refresh,
                progress=self._progress,
            )
        except (RagNotConfigured, FileNotFoundError, ValueError) as exc:
            raise CommandError(str(exc)) from exc

        self.stdout.write('')
        style = self.style.SUCCESS if result['status'] == 'success' else self.style.WARNING
        self.stdout.write(style(
            f"وضعیت: {result['status']} — {result['pages_processed']} صفحه پردازش شد، "
            f"{result['pages_skipped']} صفحه از قبل موجود بود، "
            f"{result['chunks_created']} chunk ساخته شد."
        ))
        if result.get('chunks_replaced'):
            self.stdout.write(self.style.WARNING(
                f"بازخوانی: {result['chunks_replaced']} chunkِ قدیمی پاک و با نسخه‌ی تازه "
                f"جایگزین شد. رأیِ بازبینیِ این‌ها پاک شده — "
                f"«python manage.py rag_verify_solutions» را دوباره اجرا کنید."
            ))
        self.stdout.write(
            f"مصرف: {result['prompt_tokens']:,} توکن ورودی + "
            f"{result['completion_tokens']:,} خروجی  →  ≈ ${result['cost_usd']:.2f} "
            f"({result['model']})"
        )
        if result['failed_pages']:
            self.stdout.write(self.style.ERROR(
                'صفحه‌های خطادار (اجرای دوباره سراغشان می‌رود): '
                + '، '.join(str(p) for p in result['failed_pages'])
            ))

    # ------------------------------------------------------------------
    @staticmethod
    def _parse_pages(spec: str) -> set[int]:
        """«12,15,20-24» → {12, 15, 20, 21, 22, 23, 24}"""
        pages: set[int] = set()
        for part in (p.strip() for p in spec.split(',') if p.strip()):
            if '-' in part:
                first, _, last = part.partition('-')
                try:
                    pages.update(range(int(first), int(last) + 1))
                except ValueError as exc:
                    raise CommandError(f'بازه‌ی نامعتبر در --refresh-pages: {part!r}') from exc
            else:
                try:
                    pages.add(int(part))
                except ValueError as exc:
                    raise CommandError(f'شماره‌ی نامعتبر در --refresh-pages: {part!r}') from exc
        return pages

    def _resolve_refresh(self, opts, collection, book, wanted) -> set[int]:
        """کدام صفحه‌ها بازخوانی شوند — محدود به فصل‌های انتخاب‌شده."""
        if opts['refresh_pages']:
            pages = self._parse_pages(opts['refresh_pages'])
        elif opts['refresh'] == 'all':
            # همه‌ی صفحه‌های فصل‌های این اجرا. انتخابِ هدفمند به سیگنال‌هایی
            # تکیه می‌کند که اندازه‌گیری نشان داد نویزی‌اند (regexِ «شکل» هم
            # «به شکل مخروطی» را می‌گیرد، و ۷۷٪ از رأی‌های unjudgeable از همان
            # پیش‌فیلترِ گشاد آمده)، و اختلافِ هزینه‌اش ناچیز است.
            pages = set()
            for chapter in wanted:
                first, last = book.pdf_pages(chapter)
                pages.update(range(first, last + 1))
        elif opts['refresh']:
            mode = opts['refresh']
            pages = pages_needing_refresh(
                collection,
                figures=mode in ('figures', 'both'),
                unjudgeable=mode in ('unjudgeable', 'both'),
            )
        else:
            return set()

        # فقط صفحه‌هایی که داخلِ فصل‌های همین اجرا هستند؛ وگرنه --chapters
        # بی‌اثر می‌شد و یک اجرای محدود کلِ کتاب را پاک می‌کرد.
        in_scope: set[int] = set()
        for chapter in wanted:
            first, last = book.pdf_pages(chapter)
            in_scope.update(p for p in pages if first <= p <= last)
        return in_scope

    def _confirm(self, collection, refresh) -> bool:
        chunks = LessonChunk.objects.filter(
            collection=collection, source_page__in=refresh)
        verified = chunks.filter(solution_verdict='agrees').count()
        self.stdout.write(self.style.WARNING(
            f'\n⚠️  بازخوانی {len(refresh)} صفحه را پاک و دوباره می‌سازد: '
            f'{chunks.count()} chunk، که {verified} تایشان رأیِ «agrees» دارند '
            f'و رأیشان از بین می‌رود.'
        ))
        answer = input('ادامه بدهم؟ [yes/no] ').strip().lower()
        return answer in ('y', 'yes', 'بله')

    def _plan(self, book, collection, wanted, refresh=frozenset()):
        self.stdout.write(self.style.MIGRATE_HEADING(
            f'نقشه‌ی «{book.key}» → collection «{collection}»  (فایل: {book.filename})'
        ))
        total_pages = total_left = total_refresh = 0
        for chapter in wanted:
            first, last = book.pdf_pages(chapter)
            pages = last - first + 1
            done = len(_done_pages(collection, chapter.title))
            again = sum(1 for p in refresh if first <= p <= last)
            left = max(pages - done, 0) + again
            total_pages += pages
            total_left += left
            total_refresh += again
            self.stdout.write(
                f'  فصل {chapter.number}: {chapter.title}\n'
                f'      صفحه‌ی کتاب {chapter.first_page}–{chapter.last_page}'
                f'  →  صفحه‌ی PDF {first}–{last}   ({pages} صفحه، {done} انجام‌شده، {left} مانده'
                + (f'، {again} بازخوانی)' if again else ')')
            )
        self.stdout.write('')
        self.stdout.write(f'جمع: {total_pages} صفحه، {total_left} صفحه مانده.')
        if total_refresh:
            chunks = LessonChunk.objects.filter(
                collection=collection, source_page__in=refresh)
            self.stdout.write(self.style.WARNING(
                f'از آن، {total_refresh} صفحه بازخوانی است: {chunks.count()} chunkِ موجود '
                f'پاک و از نو ساخته می‌شود ({chunks.filter(solution_verdict="agrees").count()} '
                f'تایشان رأیِ «agrees» دارند و باید دوباره verify شوند).'
            ))
        from django.conf import settings

        model = settings.OPENAI_VISION_MODEL
        per_page = estimate_page_cost(model)
        self.stdout.write(self.style.WARNING(
            f'هر صفحه یک فراخوانیِ vision دارد ({model}، ~${per_page:.4f} در صفحه، '
            f'از پنلِ صورتحساب اندازه‌گیری شده): '
            f'حدود ${total_left * per_page:.2f} برای صفحه‌های مانده.'
        ))

    def _progress(self, event):
        kind = event.get('kind')
        if kind == 'chapter':
            first, last = event['pages']
            self.stdout.write(self.style.HTTP_INFO(
                f"\n▸ فصل {event['chapter']}: {event['title']}  "
                f"(PDF {first}–{last}، {event['already']} صفحه از قبل موجود)"
            ))
        elif kind == 'page':
            self.stdout.write(f"    صفحه {event['page']}: {event['chunks']} chunk")
        elif kind == 'page_failed':
            self.stdout.write(self.style.ERROR(
                f"    صفحه {event['page']}: خطا — {event['error'][:120]}"
            ))
        elif kind == 'page_cleared':
            self.stdout.write(self.style.WARNING(
                f"    صفحه {event['page']}: {event['chunks_deleted']} chunkِ قدیمی پاک شد"
            ))
        elif kind == 'budget_reached':
            self.stdout.write(self.style.WARNING('    سقفِ --max-pages پر شد؛ توقف.'))