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,
                   'یازدهم': 7, 'دوازدهم': 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',
    'یازدهم': 'a specialist high-school math teacher for students aged 16 to 17',
    'دوازدهم': 'a specialist high-school math teacher for students aged 17 to 18, '
               'preparing them for the university entrance exam',
}


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


# بک‌اسلش‌هایی که همیشه escapeِ واقعی‌اند و نباید دست بخورند.
_JSON_SAFE_ESCAPES = set('"\\/u')
# این‌ها هم escapeِ معتبرِ JSON‌اند و هم آغازِ دستورهای پرکاربردِ LaTeX:
# \frac \times \theta \beta \rightarrow \nabla. تفکیک از روی نویسه‌ی بعدی.
_JSON_AMBIGUOUS_ESCAPES = set('bfnrt')


def repair_json_escapes(raw: str) -> str:
    """
    بک‌اسلشِ نامعتبر در رشته‌های JSON را دوبرابر می‌کند تا قابلِ تجزیه شود.

    مدل در متنِ ریاضی LaTeX تولید می‌کند و `\sqrt{2}` از نظرِ JSON خطاست، پس
    json.loads کلِ صفحه را دور می‌ریخت — دقیقاً صفحه‌های پرفرمول، یعنی
    باارزش‌ترین‌ها.

    ⚠️ نکته‌ای که نسخه‌ی اولِ همین تابع را غلط کرد: `\f`، `\t`، `\b`، `\r` و
    `\n` از نظرِ JSON معتبرند، پس «سالم» رد می‌شدند و `\frac` به formfeed+«rac»
    و `\theta` به tab+«heta» تبدیل می‌شد — خرابیِ بی‌صدا به‌جای خطای بلند.
    قاعده‌ی تفکیک: دستورِ LaTeX همیشه با حرفِ ASCII ادامه پیدا می‌کند
    (`\frac`, `\times`)، ولی escapeِ واقعی با رقم، فاصله، فارسی یا پایانِ
    رشته می‌آید (`\n۲.`، `\nحالت`). پس فقط وقتی control character حساب
    می‌شود که نویسه‌ی بعدی حرفِ ASCII نباشد.
    """
    out: list[str] = []
    i, n = 0, len(raw)
    while i < n:
        ch = raw[i]
        if ch != '\\':
            out.append(ch)
            i += 1
            continue
        nxt = raw[i + 1] if i + 1 < n else ''
        after = raw[i + 2] if i + 2 < n else ''
        looks_like_latex = after.isascii() and after.isalpha()
        keep = nxt in _JSON_SAFE_ESCAPES or (
            nxt in _JSON_AMBIGUOUS_ESCAPES and not looks_like_latex
        )
        if keep:
            out.append(ch)
            out.append(nxt)
            i += 2
        else:
            out.append('\\\\')
            i += 1
    return ''.join(out)


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'}},
            ]},
        ],
        # ۳۰۰۰ برای صفحه‌ی متراکم کم است: خروجی وسطِ JSON قطع می‌شود،
        # json.loads خطا می‌دهد و کلِ صفحه در failed_pages می‌افتد. قاعده‌ی
        # self-contained هم خروجی را بلندتر کرده (توصیفِ شکل داخلِ متن)، پس
        # سقف بالا رفت. هزینه‌ی واقعی به توکنِ تولیدشده است نه به این سقف.
        max_tokens=6000,
        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()

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        # مدل در متنِ ریاضی LaTeX می‌نویسد (\frac، \sqrt، \vec) و بک‌اسلشِ تنها
        # در JSON نامعتبر است، پس json.loads کلِ صفحه را دور می‌ریخت. این دقیقاً
        # صفحه‌های پرفرمول — باارزش‌ترین‌ها — را بی‌سروصدا حذف می‌کرد؛ روی حسابان
        # یازدهم/دوازدهم چهار صفحه هرچه تلاش شد «خطادار» ماندند تا همین‌جا پیدا شد.
        parsed = json.loads(repair_json_escapes(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()