"""
شکلِ پاسخِ دستیار: مدل یک توصیفِ JSONِ کوتاه در <diagram>…</diagram> می‌نویسد و
اپ خودش آن را به SVG تبدیل می‌کند.

چرا مدل خودش SVG/HTML نمی‌نویسد (سنجیده ۲۰۲۶-۰۹-۲۸ روی gpt-4.1-mini):
- **امنیت:** HTMLِ مدل (که متنِ کتاب و سوالِ کاربر را هم خوانده) اگر مستقیم در
  صفحه بنشیند، راهِ تزریقِ اسکریپت است. JSONِ اعتبارسنجی‌شده فقط داده است.
- **درستیِ هندسه:** مدل در حسابِ مختصات (cos 30° و …) خطا می‌کند؛ اینجا فقط
  می‌گوید «نقطه روی دایره با زاویه‌ی ۳۰°» و فرانت مختصات را حساب می‌کند.
- **هزینه و سرعت:** یک JSONِ ۱۰۰–۲۰۰ توکنی به‌جای SVGِ چندصدتوکنی.
- Structured Outputs (json_schema) را این مدل دارد ولی کلِ پاسخ را JSON
  می‌کند و استریمِ متن را خراب می‌کند؛ برای همین فقط بلوکِ آخر JSON است.

این ماژول سه کار می‌کند: قاعده‌ی پرامپت (`DIAGRAM_RULE`)، جدا کردنِ بلوک از متن
(قبل از `mathtext`، که JSON را خراب می‌کرد)، و اعتبارسنجیِ سخت‌گیر. شکلِ
نامعتبر بی‌صدا دور ریخته می‌شود؛ متنِ پاسخ هیچ‌وقت به‌خاطرِ شکل از دست نمی‌رود.
آزموده نشده: رفتارِ مدل‌های دیگر (فقط gpt-4.1-mini سنجیده شد، ۱۸ سوال).
"""
import json
import logging
import math
import re

logger = logging.getLogger(__name__)

OPEN, CLOSE = '<diagram>', '</diagram>'
_BLOCK = re.compile(r'<diagram>(.*?)(?:</diagram>|$)', re.S)

DIAGRAM_RULE = (
    'شکل (اپ این را به تصویر تبدیل می‌کند):\n'
    '- وقتی سوال درباره‌ی یکی از این‌هاست، در پایانِ پاسخ دقیقاً یک بلوکِ '
    '<diagram>…</diagram> بگذار: شکلِ هندسی (مثلث، چندضلعی، دایره، زاویه، خطِ موازی)، '
    'دایره‌ی مثلثاتی، حرکت یا مقایسه روی محورِ اعداد، نقطه یا نمودار در صفحه‌ی مختصات، '
    'کسر یا درصد به‌صورتِ تکه‌های یک شکل. هر سوالی که کلمه‌ی مثلث، دایره، زاویه، محور، '
    'نمودار، کسر یا تکه دارد، شکل می‌خواهد. فقط برای حسابِ عددی یا ساده‌کردنِ عبارتِ '
    'جبری که شکلی ندارد، شکل نگذار.\n'
    '- داخلِ بلوک فقط یک JSON است، نه متنِ دیگر، نه SVG، نه HTML. مختصات را خودت حساب '
    'نکن؛ نقطه‌ی روی دایره را با polar بده ([شعاع، زاویه به درجه]) تا اپ دقیق رسمش کند.\n'
    '- در متنِ پاسخ به شکل اشاره کن («در شکلِ پایین…») و برچسب‌های شکل همان نام‌های متن باشند.\n'
    '- سه نوع شکل:\n'
    '  ۱) صفحه‌ی مختصات یا هندسه:\n'
    '  {"kind":"plane","title":"…","axes":true,"points":{"O":[0,0],"A":[1,0],"P":{"polar":[1,30]}},'
    '"items":[{"type":"circle","center":"O","r":1},{"type":"segment","from":"O","to":"P"},'
    '{"type":"angle","vertex":"O","from":"A","to":"P","label":"30°"},'
    '{"type":"rightAngle","vertex":"H","from":"A","to":"B"},{"type":"polygon","points":["A","B","C"]},'
    '{"type":"line","through":["A","B"]},{"type":"arrow","from":"A","to":"B"},'
    '{"type":"point","point":"V","label":"رأس (1, −1)"},'
    '{"type":"function","expr":"x^2-2*x","from":-1,"to":3,"label":"y = x² − 2x"},'
    '{"type":"text","at":[0.5,1.2],"text":"…"}]}\n'
    '  نقطه‌ها با نامشان برچسب می‌خورند (برچسبِ دیگر با item نوعِ point)؛ نقطه‌ی کمکی که '
    'برچسب نمی‌خواهد با نامِ شروع‌شده با «_» بیاید. ضلع یا اندازه را با label روی segment '
    'بنویس: {"type":"segment","from":"A","to":"B","label":"۳"}. در function فقط x، عدد، '
    '+ − * / ^ ، پرانتز و sin cos tan sqrt abs log بنویس. "axes": false برای شکلِ هندسیِ بی‌محور '
    '(مثلث، چندضلعی)؛ true فقط برای صفحه‌ی مختصات و نمودار. line خطِ بی‌پایان است؛ ضلع، قطر و شعاع '
    'همیشه segment‌اند.\n'
    '  ۲) محورِ اعداد: {"kind":"numberline","title":"…","min":-5,"max":5,'
    '"points":[{"at":2,"label":"شروع"},{"at":-1,"label":"پایان"}],'
    '"jumps":[{"from":2,"to":-1,"label":"۳ واحد به چپ"}]}\n'
    '  ۳) کسر: {"kind":"fraction","title":"…","shape":"circle","parts":4,"shaded":3}\n'
    '- مثال: «وترِ مثلثِ قائم‌الزاویه با ضلع‌های ۳ و ۴» → polygon با سه رأس، rightAngle روی رأسِ قائمه '
    'و label روی هر ضلع (وتر با «?» یا جوابش).\n'
    '- title کوتاه و فارسی. حداکثر ۱۲ نقطه و ۱۲ item.'
)

# کلمه‌هایی که یعنی «این سوال شکل دارد». سنجشِ ۲۰۲۶-۰۹-۲۸: قاعده‌ی تهِ پرامپتِ
# سیستم به‌تنهایی برای مثلث ۱ از ۴ و پنج‌ضلعی ۰ از ۴ شکل می‌گذاشت؛ همان درسِ
# «/» — قاعده کنارِ خودِ سوال اثر دارد، نه تهِ پرامپت.
_SHAPE_WORDS = re.compile(
    r'مثلث|دایره|زاویه|ضلع|چندضلعی|مربع|مستطیل|لوزی|ذوزنقه|متوازی|شعاع|قطر|وتر|محیط|مساحت|'
    r'محور|نمودار|مختصات|سهمی|خط\s*راست|شیب|نقطه|کسر|تکه|پیتزا|کیک|درصد|سینوس|کسینوس|تانژانت|'
    r'sin|cos|tan|°|درجه'
)

DIAGRAM_HINT = (
    'این سوال شکل دارد: در پایانِ پاسخ بلوکِ <diagram>…</diagram> را طبقِ قاعده‌ی «شکل» بگذار.'
)


def hint_for(question: str) -> str:
    """یادآوریِ کنارِ سوال، فقط برای سوال‌هایی که کلمه‌ی شکل‌دار دارند."""
    return DIAGRAM_HINT if _SHAPE_WORDS.search(question or '') else ''


MAX_POINTS = 16
MAX_ITEMS = 16
MAX_LABEL = 40
MAX_COORD = 1000
_EXPR_OK = re.compile(r'^[0-9x+\-*/^().,\s]*(?:(?:sin|cos|tan|sqrt|abs|log|ln|pi)[0-9x+\-*/^().,\s]*)*$')
_NAME_OK = re.compile(r'^[\w′\']{1,6}$')


def _num(value, lo=-MAX_COORD, hi=MAX_COORD):
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError('not a number')
    if not math.isfinite(value) or not lo <= value <= hi:
        raise ValueError('out of range')
    return round(float(value), 4)


def _label(value) -> str:
    if value is None:
        return ''
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        value = str(value)
    if not isinstance(value, str):
        raise ValueError('label')
    return value.strip()[:MAX_LABEL]


def _xy(value):
    if not (isinstance(value, (list, tuple)) and len(value) == 2):
        raise ValueError('xy')
    return [_num(value[0]), _num(value[1])]


def _point_def(value):
    if isinstance(value, dict) and 'polar' in value:
        polar = value['polar']
        if not (isinstance(polar, (list, tuple)) and len(polar) == 2):
            raise ValueError('polar')
        out = {'polar': [_num(polar[0], 0, MAX_COORD), _num(polar[1], -3600, 3600)]}
        if 'from' in value:
            out['from'] = _ref(value['from'])
        return out
    return _xy(value)


def _ref(value):
    """نامِ نقطه یا مختصاتِ مستقیم."""
    if isinstance(value, str):
        if not _NAME_OK.match(value):
            raise ValueError('ref name')
        return value
    return _xy(value)


def _plane_item(item: dict) -> dict:
    kind = item.get('type')
    out = {'type': kind}
    if kind == 'circle':
        out.update(center=_ref(item['center']), r=_num(item['r'], 0, MAX_COORD))
    elif kind in ('segment', 'arrow'):
        out.update(**{'from': _ref(item['from'])}, to=_ref(item['to']))
    elif kind == 'line':
        a, b = item['through']
        out['through'] = [_ref(a), _ref(b)]
    elif kind in ('angle', 'rightAngle'):
        out.update(vertex=_ref(item['vertex']), **{'from': _ref(item['from'])}, to=_ref(item['to']))
    elif kind == 'polygon':
        points = item['points']
        if not isinstance(points, list) or not 3 <= len(points) <= 12:
            raise ValueError('polygon')
        out['points'] = [_ref(p) for p in points]
    elif kind == 'point':
        out['point'] = _ref(item['point'])
    elif kind == 'function':
        expr = item['expr']
        if not isinstance(expr, str) or len(expr) > 80 or not _EXPR_OK.match(expr.replace(' ', '')):
            raise ValueError('expr')
        out.update(expr=expr, **{'from': _num(item.get('from', -5))}, to=_num(item.get('to', 5)))
        if out['from'] >= out['to']:
            raise ValueError('function range')
    elif kind == 'text':
        out.update(at=_xy(item['at']), text=_label(item.get('text')))
        if not out['text']:
            raise ValueError('empty text')
    else:
        raise ValueError(f'unknown item {kind!r}')
    if kind != 'text' and item.get('label') is not None:
        out['label'] = _label(item['label'])
    return out


def validate(spec) -> dict | None:
    """
    نسخه‌ی تمیزِ spec، یا None. فقط کلیدها و نوع‌های شناخته‌شده رد می‌شوند؛
    item یا نقطه‌ی خراب تنها حذف می‌شود، ولی شکلی که چیزی برای کشیدن ندارد None است.
    """
    if not isinstance(spec, dict):
        return None
    kind = spec.get('kind', 'plane')
    try:
        title = _label(spec.get('title'))
    except ValueError:
        title = ''
    try:
        if kind == 'plane':
            points = {}
            raw_points = spec.get('points') or {}
            if not isinstance(raw_points, dict):
                return None
            for name, value in list(raw_points.items())[:MAX_POINTS]:
                if isinstance(name, str) and _NAME_OK.match(name):
                    try:
                        points[name] = _point_def(value)
                    except (ValueError, KeyError, TypeError):
                        continue
            items = []
            for item in (spec.get('items') or [])[:MAX_ITEMS]:
                if isinstance(item, dict):
                    try:
                        clean = _plane_item(item)
                    except (ValueError, KeyError, TypeError):
                        continue
                    items.append(clean)
            if not items and not points:
                return None
            # بی‌محور مگر مدل صریحاً بخواهد یا نمودارِ تابع باشد: مثلثی که محور
            # گرفته بود عددهای محور را روی برچسبِ ضلع‌ها می‌انداخت (۲۰۲۶-۰۹-۲۸).
            axes = spec.get('axes')
            if not isinstance(axes, bool):
                axes = any(i['type'] == 'function' for i in items)
            return {'kind': 'plane', 'title': title, 'axes': axes,
                    'points': points, 'items': items}
        if kind == 'numberline':
            lo, hi = _num(spec.get('min', -5)), _num(spec.get('max', 5))
            if not lo < hi or hi - lo > 200:
                return None
            marks = []
            for p in (spec.get('points') or [])[:MAX_POINTS]:
                try:
                    marks.append({'at': _num(p['at']), 'label': _label(p.get('label'))})
                except (ValueError, KeyError, TypeError):
                    continue
            jumps = []
            for j in (spec.get('jumps') or [])[:MAX_ITEMS]:
                try:
                    jumps.append({'from': _num(j['from']), 'to': _num(j['to']), 'label': _label(j.get('label'))})
                except (ValueError, KeyError, TypeError):
                    continue
            return {'kind': 'numberline', 'title': title, 'min': lo, 'max': hi,
                    'points': marks, 'jumps': jumps}
        if kind == 'fraction':
            parts = int(_num(spec.get('parts'), 1, 24))
            shaded = int(_num(spec.get('shaded', 0), 0, parts))
            shape = spec.get('shape') if spec.get('shape') in ('circle', 'rect') else 'circle'
            return {'kind': 'fraction', 'title': title, 'shape': shape, 'parts': parts, 'shaded': shaded}
    except (ValueError, KeyError, TypeError):
        return None
    return None


def extract(text: str) -> tuple[str, dict | None]:
    """(متنِ بی‌بلوک، spec یا None). بلوکِ بسته‌نشده هم جدا می‌شود تا JSON در متن نماند."""
    if OPEN not in (text or ''):
        return text, None
    match = _BLOCK.search(text)
    spec = parse(match.group(1)) if match else None
    return _BLOCK.sub('', text).rstrip(), spec


def parse(body: str) -> dict | None:
    body = (body or '').strip()
    if body.startswith('```'):
        body = body.strip('`').removeprefix('json').strip()
    try:
        return validate(json.loads(body))
    except (json.JSONDecodeError, TypeError):
        logger.info('diagram: JSONِ نامعتبر از مدل، دور ریخته شد')
        return None


class StreamSplitter:
    """
    متنِ استریم را عبور می‌دهد و بلوکِ <diagram> را کنار می‌گذارد. تهِ هر تکه که
    ممکن است آغازِ «<diagram>» باشد نگه داشته می‌شود، چون برچسب می‌تواند بینِ
    دو تکه‌ی استریم نصف شود.
    """

    def __init__(self):
        self._pending = ''
        self._block: list[str] | None = None
        self.diagram: dict | None = None

    def feed(self, delta: str) -> str:
        if self._block is not None:
            self._block.append(delta)
            joined = ''.join(self._block)
            if CLOSE not in joined:
                return ''
            body, rest = joined.split(CLOSE, 1)
            self.diagram = self.diagram or parse(body)
            self._block = None
            return self.feed(rest) if rest else ''
        text = self._pending + delta
        self._pending = ''
        at = text.find(OPEN)
        if at >= 0:
            self._block = []
            return text[:at] + self.feed(text[at + len(OPEN):])
        for keep in range(min(len(OPEN) - 1, len(text)), 0, -1):
            if OPEN.startswith(text[-keep:]):
                self._pending = text[-keep:]
                return text[:-keep]
        return text

    def flush(self) -> str:
        if self._block is not None:
            self.diagram = self.diagram or parse(''.join(self._block))
            self._block = None
            return ''
        tail, self._pending = self._pending, ''
        return tail