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ِ ریاضیِ تیمِ دیگه هم تولید می‌کرد.
    """
    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}".
- 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()