MathApp/rag/services/mathtext.py

"""
تبدیلِ LaTeX و markdown در پاسخِ دستیار به متنِ ساده با نمادهای یونیکد.

چرا لازم است: صفحه‌ی «از پی بپرس» پاسخ را همان‌طور که هست، به‌صورتِ متنِ خام
نشان می‌دهد. مدل‌ها برای ریاضی عادت دارند `\\( \\frac{1}{2} \\)` و `### مرحله ۱`
بنویسند؛ روی گوشی این‌ها عیناً دیده می‌شوند و فرمولِ حسابان به رشته‌ای از
بک‌اسلش و آکولاد تبدیل می‌شود. پرامپت از مدل می‌خواهد متنِ ساده بنویسد، ولی
مدل همیشه گوش نمی‌دهد — این ماژول تورِ ایمنیِ قطعی پشتِ آن است.

قاعده‌ی مهم: آکولادی که جزوِ LaTeX نیست دست نمی‌خورد. مجموعه‌ها در ریاضیِ نهم
با `{1, 2, 3}` نوشته می‌شوند و نباید پاک شوند.
"""
import re
from contextvars import ContextVar

# برای دبستان «/» ممیزِ اعشار است (مثلِ کتابشان)، پس کسر «۳ بر ۴» نوشته می‌شود.
_FRACTION_WORD: ContextVar[bool] = ContextVar('fraction_word', default=False)

_SUPERSCRIPT = str.maketrans('0123456789+-=()nixyk', '⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻⁼⁽⁾ⁿⁱˣʸᵏ')
_SUBSCRIPT = str.maketrans('0123456789+-=()', '₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎')
_SUP_CHARS = set('0123456789+-=()nixyk')
_SUB_CHARS = set('0123456789+-=()')

_SYMBOLS = {
    'times': '×', 'cdot': '·', 'div': '÷', 'pm': '±', 'mp': '∓',
    'leq': '≤', 'le': '≤', 'geq': '≥', 'ge': '≥', 'neq': '≠', 'ne': '≠',
    'approx': '≈', 'simeq': '≃', 'equiv': '≡', 'infty': '∞',
    'pi': 'π', 'alpha': 'α', 'beta': 'β', 'gamma': 'γ', 'theta': 'θ',
    'Delta': 'Δ', 'delta': 'δ', 'lambda': 'λ', 'mu': 'μ', 'sigma': 'σ',
    'Sigma': 'Σ', 'phi': 'φ', 'varphi': 'φ', 'omega': 'ω', 'epsilon': 'ε',
    'varepsilon': 'ε', 'rho': 'ρ', 'tau': 'τ',
    'rightarrow': '→', 'to': '→', 'leftarrow': '←', 'Rightarrow': '⇒',
    'implies': '⇒', 'Leftrightarrow': '⇔', 'iff': '⇔', 'mapsto': '↦',
    'in': '∈', 'notin': '∉', 'subset': '⊂', 'subseteq': '⊆', 'supset': '⊃',
    'cup': '∪', 'cap': '∩', 'emptyset': '∅', 'varnothing': '∅',
    'forall': '∀', 'exists': '∃', 'angle': '∠', 'perp': '⊥', 'parallel': '∥',
    'triangle': '△', 'circ': '∘', 'degree': '°', 'prime': '′',
    'ldots': '…', 'dots': '…', 'cdots': '⋯', 'int': '∫', 'sum': 'Σ',
    'prod': '∏', 'therefore': '∴', 'propto': '∝', 'neg': '¬',
    'quad': '  ', 'qquad': '    ', 'lbrace': '{', 'rbrace': '}',
    'setminus': '−', 'mid': '|', 'vert': '|',
}
_PREFIX_SIGNS = {'pm', 'mp', 'neg', 'angle', 'triangle', 'forall', 'exists'}
_BLACKBOARD = {'R': 'ℝ', 'N': 'ℕ', 'Z': 'ℤ', 'Q': 'ℚ', 'C': 'ℂ'}
# \sin و \log و ... فقط بک‌اسلششان حذف می‌شود.
_FUNCTIONS = {
    'sin', 'cos', 'tan', 'cot', 'sec', 'csc', 'log', 'ln', 'exp', 'lim',
    'max', 'min', 'arcsin', 'arccos', 'arctan', 'sinh', 'cosh', 'det', 'gcd',
}
# دستورهایی که فقط محتوای آرگومانشان می‌ماند.
_UNWRAP = {'text', 'textbf', 'textit', 'mathrm', 'mathbf', 'mathit', 'operatorname',
           'boxed', 'displaystyle', 'mbox', 'underline', 'bm', 'textrm'}
_DROP = {'left', 'right', 'big', 'Big', 'bigg', 'Bigg', 'displaystyle', 'limits',
         'nonumber', 'hline', 'noindent'}


def _read_group(s: str, i: int) -> tuple[str, int]:
    """از s[i] یک آرگومانِ LaTeX می‌خواند: {…} با آکولادِ تودرتو، یا یک نویسه."""
    while i < len(s) and s[i] == ' ':
        i += 1
    if i >= len(s):
        return '', i
    if s[i] != '{':
        if s[i] == '\\':
            m = re.match(r'\\[A-Za-z]+', s[i:])
            if m:
                return m.group(0), i + len(m.group(0))
        return s[i], i + 1
    depth, j = 0, i
    while j < len(s):
        if s[j] == '\\' and j + 1 < len(s):
            j += 2
            continue
        if s[j] == '{':
            depth += 1
        elif s[j] == '}':
            depth -= 1
            if depth == 0:
                return s[i + 1:j], j + 1
        j += 1
    return s[i + 1:], len(s)


_ATOM = re.compile(r'√?[\d.]+|√?\w|\(.*\)')


def _needs_parens(expr: str) -> bool:
    """«2h» در مخرج باید (2h) شود، وگرنه 1/2h یعنی h/2."""
    return not _ATOM.fullmatch(expr.strip())


def _script(body: str, table, allowed: set, marker: str) -> str:
    if body and all(ch in allowed for ch in body):
        return body.translate(table)
    return f'{marker}({body})' if len(body) > 1 else f'{marker}{body}'


def _convert(s: str) -> str:
    """یک عبارتِ LaTeX را (بازگشتی) به یونیکد تبدیل می‌کند."""
    out: list[str] = []
    i = 0
    while i < len(s):
        ch = s[i]
        if ch == '\\':
            if i + 1 < len(s) and not s[i + 1].isalpha():
                nxt = s[i + 1]
                if nxt == '\\':
                    out.append('\n')
                elif nxt in '{}%$&#_':
                    out.append(nxt)
                elif nxt in ',;:! ':
                    out.append(' ' if nxt != '!' else '')
                else:
                    out.append(nxt)
                i += 2
                continue
            m = re.match(r'\\([A-Za-z]+)', s[i:])
            name = m.group(1)
            i += len(m.group(0))
            if name == 'frac' or name in ('dfrac', 'tfrac'):
                num, i = _read_group(s, i)
                den, i = _read_group(s, i)
                num, den = _convert(num), _convert(den)
                num = f'({num})' if _needs_parens(num) else num
                den = f'({den})' if _needs_parens(den) else den
                out.append(f'{num} بر {den}' if _FRACTION_WORD.get() else f'{num}/{den}')
            elif name == 'sqrt':
                index = ''
                if i < len(s) and s[i] == '[':
                    close = s.find(']', i)
                    if close != -1:
                        index, i = s[i + 1:close], close + 1
                body, i = _read_group(s, i)
                body = _convert(body)
                root = (index.translate(_SUPERSCRIPT) if index else '') + '√'
                out.append(root + (f'({body})' if len(body) > 1 else body))
            elif name == 'mathbb':
                body, i = _read_group(s, i)
                out.append(_BLACKBOARD.get(body, body))
            elif name in ('vec', 'overrightarrow'):
                body, i = _read_group(s, i)
                out.append(_convert(body) + '⃗')
            elif name in ('overline', 'bar', 'widehat', 'hat'):
                body, i = _read_group(s, i)
                out.append(_convert(body))
            elif name in ('begin', 'end'):
                _env, i = _read_group(s, i)
            elif name in _UNWRAP:
                if i < len(s) and s[i] == '{':
                    body, i = _read_group(s, i)
                    out.append(_convert(body))
            elif name in _DROP:
                # \left( → (  — خودِ پرانتز در دورِ بعدِ حلقه نوشته می‌شود.
                if name in ('left', 'right') and i < len(s) and s[i] == '.':
                    i += 1
            elif name in _SYMBOLS:
                out.append(_SYMBOLS[name])
                # «\pm 2» → «±2»: نمادِ یک‌طرفه به عددش می‌چسبد.
                if name in _PREFIX_SIGNS:
                    while i < len(s) and s[i] == ' ':
                        i += 1
            elif name in _FUNCTIONS:
                out.append(name)
            else:
                out.append(name)
        elif ch in '^_':
            # مدل «x^10» را به معنیِ ده می‌نویسد، نه «x¹0» — رقم‌ها با هم خوانده می‌شوند.
            digits = re.match(r'\d+', s[i + 1:])
            if digits:
                body, i = digits.group(0), i + 1 + len(digits.group(0))
            else:
                body, i = _read_group(s, i + 1)
            if ch == '^' and body in ('\\circ', 'circ', '°'):
                out.append('°')
                continue
            if ch == '^' and body in ('\\prime', "'"):
                out.append('′')
                continue
            body = _convert(body)
            if ch == '_' and ''.join(out).endswith('lim'):
                out.append(f'({body})')
                continue
            if ch == '^':
                out.append(_script(body, _SUPERSCRIPT, _SUP_CHARS, '^'))
            else:
                out.append(_script(body, _SUBSCRIPT, _SUB_CHARS, '_'))
        elif ch == '{':
            # آکولادِ گروه‌بندیِ LaTeX (بعد از پاک شدنِ دستورِ صاحبش)
            body, i = _read_group(s, i)
            out.append(_convert(body))
        elif ch == '&':
            i += 1
        else:
            out.append(ch)
            i += 1
    return ''.join(out)


# \( … \) و \[ … \] و $$ … $$ و $ … $ — محتوا تبدیل و جداکننده حذف می‌شود.
_MATH_SPANS = re.compile(
    r'\\\[(.+?)\\\]|\\\((.+?)\\\)|\$\$(.+?)\$\$|(?<![\w$])\$([^$\n]+?)\$(?!\w)',
    re.DOTALL,
)
# LaTeXِ بی‌جداکننده که مدل گاهی وسطِ متن می‌نویسد: \frac{..}{..}، x^{2}، \sqrt{..}
_BARE_LATEX = re.compile(r'\\[A-Za-z{}]|[\^_]\{|(?<=[\w)])\^\d')


def _math_span(match: re.Match) -> str:
    body = next(g for g in match.groups() if g is not None)
    # بلوکِ نمایشی را مدل خودش در خطِ جدا می‌نویسد؛ خطِ خالیِ اضافه نمی‌سازیم.
    if '\\\\' in body:  # چندخطی (\\ یا cases/aligned)
        return '\n'.join(' '.join(line.split()) for line in _convert(body).strip().split('\n'))
    return ' '.join(_convert(body).split())


def _strip_markdown(text: str) -> str:
    lines = []
    for line in text.split('\n'):
        stripped = line.strip()
        if re.fullmatch(r'\|?\s*:?-{2,}:?\s*(\|\s*:?-{2,}:?\s*)*\|?', stripped):
            continue  # خطِ جداکننده‌ی جدولِ markdown
        if stripped.startswith('|') and stripped.endswith('|') and stripped.count('|') >= 2:
            cells = [c.strip() for c in stripped.strip('|').split('|')]
            line = '  |  '.join(cells)
        line = re.sub(r'^\s{0,3}#{1,6}\s+', '', line)
        line = re.sub(r'^(\s*)[-*]\s+', r'\1• ', line)
        line = re.sub(r'^\s*(-{3,}|\*{3,}|_{3,})\s*$', '', line)
        line = re.sub(r'\*\*(.+?)\*\*', r'\1', line)
        line = re.sub(r'__(.+?)__', r'\1', line)
        line = line.replace('`', '')
        lines.append(line.strip())  # تورفتگی در حبابِ راست‌به‌چپ معنایی ندارد
    return '\n'.join(lines)


def to_plain(text: str, *, fraction_word: bool = False) -> str:
    """
    پاسخِ مدل → متنِ ساده‌ی قابلِ خواندن روی گوشی.
    fraction_word: کسرِ LaTeX به‌صورتِ «a بر b» (برای دبستان) به‌جای a/b.
    """
    if not text:
        return text
    token = _FRACTION_WORD.set(fraction_word)
    try:
        return _to_plain(text)
    finally:
        _FRACTION_WORD.reset(token)


def _to_plain(text: str) -> str:
    text = _MATH_SPANS.sub(_math_span, text)
    if _BARE_LATEX.search(text):
        text = '\n'.join(_convert(line) if _BARE_LATEX.search(line) else line
                         for line in text.split('\n'))
    text = _strip_markdown(text)
    text = re.sub(r'[ \t]+\n', '\n', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


class StreamSanitizer:
    """
    همان to_plain برای پاسخِ استریمی. یک دستورِ LaTeX ممکن است بینِ دو تکه‌ی
    استریم نصف شود (`\\fr` + `ac{1}{2}`)، پس متن تا پایانِ هر خط نگه داشته و
    خط‌به‌خط تبدیل می‌شود؛ یک بلوکِ `\\[ … \\]` چندخطی هم تا بسته شدنش نگه
    داشته می‌شود.
    """

    def __init__(self, *, fraction_word: bool = False):
        self._fraction_word = fraction_word
        self._buffer = ''
        self._emitted_any = False
        self._newlines = 0  # خط‌های خالیِ معلق؛ مثلِ to_plain حداکثر دوتا

    def _ready_cut(self) -> int:
        cut = self._buffer.rfind('\n')
        if cut == -1:
            return -1
        head = self._buffer[:cut + 1]
        # بلوکِ نمایشیِ باز را نصفه نفرست.
        if head.count('\\[') > head.count('\\]') or head.count('$$') % 2:
            return -1
        return cut + 1

    def feed(self, delta: str) -> str:
        self._buffer += delta
        cut = self._ready_cut()
        if cut == -1:
            return ''
        head, self._buffer = self._buffer[:cut], self._buffer[cut:]
        return self._emit(head, final=False)

    def flush(self) -> str:
        head, self._buffer = self._buffer, ''
        return self._emit(head, final=True)

    def _emit(self, chunk: str, *, final: bool) -> str:
        plain = to_plain(chunk, fraction_word=self._fraction_word)
        if not plain:
            # خطِ واقعاً خالی خطِ خالی است؛ خطی که محتوایش حذف شد (جداکننده‌ی
            # جدول) مثلِ to_plain هیچ ردی نمی‌گذارد.
            if not chunk.strip():
                self._newlines += chunk.count('\n')
            return ''
        leading = len(chunk) - len(chunk.lstrip('\n'))
        prefix = '\n' * min(self._newlines + leading, 2) if self._emitted_any else ''
        self._emitted_any = True
        self._newlines = 0 if final else len(chunk) - len(chunk.rstrip('\n'))
        return prefix + plain