MathApp/rag/services/pdf.py

"""
استخراج از PDF کتابِ درسی + تحلیلِ عکسِ صفحه با مدلِ vision.

برداشتِ پاک‌سازی‌شده از mathapp/services/read_pdf.py و info_question.py مخزنِ
Rag_backend. به‌جای نوشتنِ نتیجه در یک فایلِ JSON روی دیسک، دیکشنری
برمی‌گردونه تا لایه‌ی ingest تصمیم بگیره.

pymupdf (fitz) و pdf2image اختیاری‌ان؛ نبودشون فقط این ماژول رو از کار
می‌ندازه (RagNotConfigured)، نه بقیه‌ی پروژه رو.
"""
import base64
import json
from io import BytesIO

from django.conf import settings

from . import clients
from .clients import RagNotConfigured


def extract_page_images(pdf_path: str, first_page: int = 1, dpi: int = 200,
                        last_page: int | None = None):
    """
    لیستی از تصاویرِ PIL برای صفحاتِ PDF در بازه‌ی [first_page, last_page].

    `last_page` برای کتابِ کامل ضروری است: بدونِ آن pdf2image کلِ کتاب را
    یک‌جا به تصویر تبدیل می‌کند (۱۵۰+ صفحه در حافظه) و امکانِ ingest فصل‌به‌فصل
    هم از بین می‌رود.
    """
    try:
        from pdf2image import convert_from_path
    except ModuleNotFoundError as exc:
        raise RagNotConfigured('pdf2image نصب نیست یا poppler موجود نیست.') from exc
    return convert_from_path(pdf_path, dpi=dpi, first_page=first_page, last_page=last_page)


def detect_grade_and_start_page(pdf_path: str) -> tuple[int, str | None]:
    """
    از متنِ صفحه‌ی اولِ PDF، پایه‌ی تحصیلی و صفحه‌ای که محتوای واقعی از اونجا
    شروع می‌شه رو حدس می‌زنه (مثلِ نسخه‌ی Rag_backend).
    """
    try:
        import fitz  # PyMuPDF
    except ModuleNotFoundError as exc:
        raise RagNotConfigured('pymupdf (fitz) نصب نیست.') from exc

    # صفحه‌ی شروعِ محتوای واقعی؛ هرچه اینجا نباشد پیش‌فرضِ ۷ می‌گیرد.
    start_pages = {'پنجم دبستان': 11, 'ششم دبستان': 9, 'نهم': 9, 'دهم': 7}
    # ترتیب مهم است: «دهم» زیررشته‌ی «نهم» نیست ولی «هفتم»/«هشتم» در متنِ
    # شناسنامه‌ی کتابِ نهم هم می‌آیند («در ادامهٔ پایه‌های هفتم و هشتم»)، پس
    # پایه‌های بالاتر باید زودتر بررسی شوند.
    grades = ['دهم', 'نهم', 'هشتم', 'هفتم', 'اول دبستان', 'دوم دبستان',
              'سوم دبستان', 'چهارم دبستان', 'پنجم دبستان', 'ششم دبستان']

    doc = fitz.open(pdf_path)
    try:
        text = doc[0].get_text()
    finally:
        doc.close()

    grade = next((g for g in grades if g in text), None)
    return start_pages.get(grade, 7), grade


_VISION_JSON_TEMPLATE = """
{
  "grade": "",
  "chapter": "",
  "problemTypes": [
    {"type": "", "keywords": [], "generalSolution": "",
     "exampleFromImage": {"problem": "", "solution": ""}}
  ]
}
"""


# پایه → توصیفِ معلمِ متناسب. پرامپت قبلاً «معلمِ ابتدایی، بچه‌های ۸ تا ۱۲
# ساله» را هاردکد کرده بود؛ برای پایه‌های نهم/دهم این هم لحنِ توضیح را بچگانه
# می‌کرد و هم مدل را به ساده‌سازیِ بیش از حدِ مفاهیم (اتحاد، مثلثات، تابع)
# سوق می‌داد.
_TEACHER_BY_GRADE = {
    'اول دبستان': 'a specialist elementary math teacher for children aged 6 to 7',
    'دوم دبستان': 'a specialist elementary math teacher for children aged 7 to 8',
    'سوم دبستان': 'a specialist elementary math teacher for children aged 8 to 9',
    'چهارم دبستان': 'a specialist elementary math teacher for children aged 9 to 10',
    'پنجم دبستان': 'a specialist elementary math teacher for children aged 10 to 11',
    'ششم دبستان': 'a specialist elementary math teacher for children aged 11 to 12',
    'هفتم': 'a specialist middle-school math teacher for students aged 12 to 13',
    'هشتم': 'a specialist middle-school math teacher for students aged 13 to 14',
    'نهم': 'a specialist middle-school math teacher for students aged 14 to 15',
    'دهم': 'a specialist high-school math teacher for students aged 15 to 16',
}


def teacher_persona(grade: str) -> str:
    return _TEACHER_BY_GRADE.get(
        (grade or '').strip(),
        'a specialist school math teacher',
    )


def extract_problem_types(image, grade: str, chapter_name: str) -> dict:
    """
    یک عکسِ صفحه‌ی کتاب → ساختارِ JSONِ انواعِ مسئله (فارسی)، با مدلِ vision.
    خروجی همون شکلی‌ست که uploadِ ریاضیِ تیمِ دیگه هم تولید می‌کرد.

    ⚠️ قاعده‌ی «self-contained» در پرامپت تازه اضافه شده و کتاب‌هایی که تا امروز
    وارد شده‌اند بدونِ آن استخراج شده‌اند. اثرش با عدد: از ۲٬۱۱۹ مسئله‌ی موجود،
    ۸۸۱ تا (۴۲٪) به «شکلِ زیر» یا «جدولِ روبه‌رو» ارجاع می‌دهند و بدونِ آن تصویر
    قابلِ حل نیستند — نه برای دانش‌آموز، نه برای دستیار، نه برای ساختِ سؤال.

    نکته این است که مدل آن شکل را **دیده بود**؛ محدودیتِ vision نبود، فقط از او
    خواسته نشده بود آنچه می‌بیند را بنویسد. برای همین با ingestِ دوباره‌ی همان
    صفحه‌ها قابلِ جبران است، ولی چون هزینه دارد فعلاً انجام نشده. هر کتابِ
    تازه‌ای که از این پس وارد شود این مشکل را ندارد.
    """
    prompt = f"""
You are {teacher_persona(grade)}.
Analyze the uploaded image and extract every unique math problem type for the given chapter.

For each type produce: type, keywords (3-7, Persian), generalSolution (step-by-step),
exampleFromImage (one real problem from the image + full solution).

Output must follow EXACTLY this JSON structure:
{_VISION_JSON_TEMPLATE}

Rules:
- Everything in Persian (Farsi).
- grade = "{grade}"
- chapter = "{chapter_name}"
- exampleFromImage must be a real problem from the image (not invented).
- Keep the mathematical notation of the textbook; do not simplify the concept
  to a lower grade level than "{grade}".
- SELF-CONTAINED PROBLEMS. You can see the figures, tables and diagrams on this
  page, but the reader of your output cannot. Whenever a problem refers to one
  ("the shape below", "the table opposite", "the given diagram"), write what it
  shows INTO the problem text itself: the measurements, the vertex labels, the
  angles, the table's values, what is shaded. Someone with only your text must
  be able to solve it. If a figure genuinely cannot be put into words (a freehand
  drawing task, a colouring exercise), skip that problem instead of writing a
  reference to something the reader will never see.
- The same applies to the solution: no step may depend on looking at the page.
- If the page has no math problem at all (cover, foreword, table of contents,
  a page of pure text), return an empty "problemTypes" list instead of inventing one.
- Output ONLY the JSON object. No markdown, no backticks.
"""

    buffered = BytesIO()
    image.save(buffered, format='JPEG')
    img_b64 = base64.b64encode(buffered.getvalue()).decode()

    client = clients.get_openai_client()
    response = client.chat.completions.create(
        model=settings.OPENAI_VISION_MODEL,
        messages=[
            {'role': 'system', 'content': 'You process educational visual content and respond ONLY with valid JSON.'},
            {'role': 'user', 'content': [
                {'type': 'text', 'text': prompt},
                {'type': 'image_url', 'image_url': {'url': f'data:image/jpeg;base64,{img_b64}', 'detail': 'high'}},
            ]},
        ],
        max_tokens=3000,
        temperature=0.1,
    )
    raw = (response.choices[0].message.content or '').strip()
    if raw.startswith('```json'):
        raw = raw[len('```json'):].strip()
    if raw.startswith('```'):
        raw = raw[3:].strip()
    if raw.endswith('```'):
        raw = raw[:-3].strip()

    parsed = json.loads(raw)
    # مصرفِ واقعیِ همین فراخوانی را همراهِ نتیجه برمی‌گردانیم تا ingest بتواند
    # هزینه‌ی واقعی را گزارش کند. تا پیش از این فقط تخمینِ سرانگشتی داشتیم و
    # آن تخمین ۲.۳ برابرِ واقعیت بود.
    usage = getattr(response, 'usage', None)
    parsed['_usage'] = {
        'promptTokens': getattr(usage, 'prompt_tokens', 0) or 0,
        'completionTokens': getattr(usage, 'completion_tokens', 0) or 0,
    }
    return parsed


def problem_type_to_text(problem_type: dict) -> str:
    """یک problemType → متنِ قابلِ embedding (مثلِ doc_text در Rag_backend)."""
    example = problem_type.get('exampleFromImage', {}) or {}
    keywords = ' '.join(problem_type.get('keywords', []) or [])
    return (
        f"Type: {problem_type.get('type', '')}\n"
        f"Keywords: {keywords}\n"
        f"General solution: {problem_type.get('generalSolution', '')}\n"
        f"Example problem: {example.get('problem', '')}\n"
        f"Example solution: {example.get('solution', '')}"
    ).strip()