md_to_gost/word2md/refs.py

"""Second pass: replace «Рисунок 1.1» / «табл. 2» with @Тип:id."""

from __future__ import annotations

import re

from .blocks import (
    Block,
    CaptionBlock,
    EquationBlock,
    ImageBlock,
    InlineSpan,
    ListBlock,
    ListingBlock,
    ParagraphBlock,
    TableBlock,
)
from .captions import caption_unique_id

# Number like 1 / 1.1 / 2.3.4 — do NOT use [\d.]+ (eats trailing sentence period).
_NUM = r"(?P<num>\d+(?:\.\d+)*)"

_REF_PATTERNS = [
    (
        re.compile(
            rf"(?P<full>(?:Рисунк(?:е|а|у|ом)?|Рисунок|Рис\.?)\s+{_NUM})",
            re.IGNORECASE,
        ),
        "figure",
    ),
    (
        re.compile(
            rf"(?P<full>(?:Таблиц[аыеу]|табл\.?)\s+{_NUM})",
            re.IGNORECASE,
        ),
        "table",
    ),
    (
        re.compile(
            rf"(?P<full>(?:Листинг(?:а|у|е|ом)?|лист\.?)\s+{_NUM})",
            re.IGNORECASE,
        ),
        "listing",
    ),
    (
        re.compile(
            rf"(?P<full>(?:Формул[аыеу]|форм\.?)\s+{_NUM})",
            re.IGNORECASE,
        ),
        "equation",
    ),
]

_TYPE_LABEL = {
    "figure": "Рисунок",
    "table": "Таблица",
    "listing": "Листинг",
    "equation": "Формула",
}


def collect_number_maps(blocks: list[Block]) -> dict[str, dict[str, str]]:
    """kind → {number → unique_id}."""
    maps: dict[str, dict[str, str]] = {
        "figure": {},
        "table": {},
        "listing": {},
        "equation": {},
    }
    for b in blocks:
        if isinstance(b, CaptionBlock) and not b.is_continuation:
            maps[b.kind][b.number] = b.unique_id or caption_unique_id(b.kind, b.number)
        elif isinstance(b, ImageBlock) and b.caption_id:
            num = _id_to_number(b.caption_id, "fig")
            if num:
                maps["figure"][num] = b.caption_id
        elif isinstance(b, TableBlock) and b.caption_id:
            num = _id_to_number(b.caption_id, "tbl")
            if num:
                maps["table"][num] = b.caption_id
        elif isinstance(b, ListingBlock) and b.caption_id:
            num = _id_to_number(b.caption_id, "lst")
            if num:
                maps["listing"][num] = b.caption_id
        elif isinstance(b, EquationBlock) and b.number and b.unique_id:
            maps["equation"][b.number] = b.unique_id
    return maps


def _id_to_number(uid: str, prefix: str) -> str | None:
    if not uid.startswith(prefix):
        return None
    rest = uid[len(prefix):]
    if not rest:
        return None
    return rest.replace("_", ".")


def rewrite_text(text: str, maps: dict[str, dict[str, str]]) -> str:
    """Replace object references; leave bibliography [1] alone."""
    if not text:
        return text

    def repl_factory(kind: str):
        label = _TYPE_LABEL[kind]
        num_map = maps.get(kind) or {}

        def repl(m: re.Match) -> str:
            num = m.group("num")
            uid = num_map.get(num)
            if not uid:
                return m.group(0)
            return f"@{label}:{uid}"

        return repl

    out = text
    for pattern, kind in _REF_PATTERNS:
        out = pattern.sub(repl_factory(kind), out)
    return out


def _rewrite_paragraph_spans(para: ParagraphBlock, maps: dict[str, dict[str, str]]) -> None:
    """Rewrite refs even when Word split «Таблице 5.1» across runs."""
    spans = para.spans
    if not spans:
        return
    new_spans: list[InlineSpan] = []
    i = 0
    while i < len(spans):
        s = spans[i]
        if s.math or s.href:
            new_spans.append(s)
            i += 1
            continue
        j = i
        chunk: list[InlineSpan] = []
        while j < len(spans) and not spans[j].math and not spans[j].href:
            chunk.append(spans[j])
            j += 1
        joined = "".join(p.text or "" for p in chunk)
        rewritten = rewrite_text(joined, maps)
        if rewritten == joined:
            new_spans.extend(chunk)
        else:
            first = chunk[0]
            same_fmt = all(
                p.bold == first.bold and p.italic == first.italic and p.strike == first.strike
                for p in chunk
            )
            new_spans.append(
                InlineSpan(
                    text=rewritten,
                    bold=first.bold if same_fmt else False,
                    italic=first.italic if same_fmt else False,
                    strike=first.strike if same_fmt else False,
                )
            )
        i = j
    para.spans = new_spans


def apply_reference_pass(blocks: list[Block]) -> list[Block]:
    maps = collect_number_maps(blocks)
    for b in blocks:
        if isinstance(b, ParagraphBlock):
            _rewrite_paragraph_spans(b, maps)
        elif isinstance(b, ListBlock):
            b.items = [rewrite_text(item, maps) for item in b.items]
    return blocks