md_to_gost/md2gost/label_pass.py

"""Pre-assign section-scoped numbers and resolve @Type:label references."""

from __future__ import annotations

import re

from .numberer import Numberer, APPENDIX_LETTERS
from .renderable import Renderable
from .renderable.heading import Heading
from .renderable.requires_numbering import RequiresNumbering
from .renderable.equation import Equation
from .renderable.paragraph import Paragraph
from .renderable.caption import CaptionInfo


APPENDIX_RE = re.compile(r"^ПРИЛОЖЕНИЕ\s+([А-ЯA-Z])\b", re.IGNORECASE)


class LabelRegistry:
    def __init__(self):
        self._map: dict[tuple[str, str], str] = {}

    def put(self, category: str, name: str, number: str):
        self._map[(category.lower(), name)] = number
        # also store without category for loose lookup
        self._map[("*", name)] = f"{category} {number}"

    def get(self, category: str, name: str) -> str | None:
        return self._map.get((category.lower(), name))

    def get_display(self, category: str, name: str) -> str | None:
        n = self.get(category, name)
        if n:
            return f"{category} {n}"
        # try Russian category aliases
        aliases = {
            "рисунок": "Рисунок",
            "figure": "Рисунок",
            "таблица": "Таблица",
            "table": "Таблица",
            "листинг": "Листинг",
            "listing": "Листинг",
            "формула": "Формула",
            "formula": "Формула",
            "equation": "Формула",
        }
        cat = aliases.get(category.lower(), category)
        n = self.get(cat, name)
        if n:
            return f"{cat} {n}"
        return None


def assign_numbers(renderables: list[Renderable],
                   numbered_equations: set[str] | None = None,
                   numbering_scope: str = "section") -> LabelRegistry:
    """Walk renderables once, set_number on objects, fill LabelRegistry."""
    numbered_equations = numbered_equations or set()
    numberer = Numberer(mode=numbering_scope)
    registry = LabelRegistry()
    section_count = 0
    appendix_index = 0

    for r in renderables:
        if isinstance(r, Heading):
            text = (r.text or "").strip()
            upper = text.upper()
            if r.level == 1 and r.is_numbered:
                section_count += 1
                numberer.enter_section(section_count)
            m = APPENDIX_RE.match(upper)
            if m:
                numberer.enter_appendix(m.group(1).upper())
            elif upper.startswith("ПРИЛОЖЕНИЕ"):
                parts = upper.split()
                if len(parts) >= 2 and parts[1] in APPENDIX_LETTERS:
                    numberer.enter_appendix(parts[1])
                elif appendix_index < len(APPENDIX_LETTERS):
                    numberer.enter_appendix(APPENDIX_LETTERS[appendix_index])
                    appendix_index += 1

        if isinstance(r, Equation):
            label = r.unique_name
            if r.needs_numbering or (label and label in numbered_equations):
                r.enable_numbering()
                num = numberer.next_number("Формула", label)
                r.set_number(num)
                if label:
                    registry.put("Формула", label, num)
            continue

        if isinstance(r, RequiresNumbering):
            label = getattr(r, "unique_name", None)
            num = numberer.next_number(r.numbering_category, label)
            r.set_number(num)
            if label:
                registry.put(r.numbering_category, label, num)
            # Diagram +listing → also number Листинг with same label
            from .renderable.diagram import DiagramFigure
            if isinstance(r, DiagramFigure) and r.listing is not None:
                lnum = numberer.next_number("Листинг", label)
                r.set_listing_number(lnum)
                if label:
                    registry.put("Листинг", label, lnum)
            # Images nested in paragraphs
            if isinstance(r, Paragraph):
                pass

        # Images attached to paragraphs
        if isinstance(r, Paragraph):
            for img in getattr(r, "_images", []):
                label = getattr(img, "unique_name", None)
                num = numberer.next_number(img.numbering_category, label)
                img.set_number(num)
                if label:
                    registry.put(img.numbering_category, label, num)

    return registry


def resolve_pending_in_renderables(renderables: list[Renderable], resolve_fn) -> None:
    """Resolve @Type:label placeholders in paragraphs, lists, headings and table cells."""
    from .renderable.list import List as RList
    from .renderable.table import Table
    from .renderable.heading import Heading

    def resolve_para(p):
        if hasattr(p, "resolve_pending_refs"):
            p.resolve_pending_refs(resolve_fn)

    for r in renderables:
        if isinstance(r, Paragraph):
            resolve_para(r)
        elif isinstance(r, Heading):
            resolve_para(r)
        elif isinstance(r, RList):
            for p in getattr(r, "_paragraphs", []):
                resolve_para(p)
        elif isinstance(r, Table):
            for row in getattr(r, "_rows", []):
                for cell in row:
                    for p in cell:
                        resolve_para(p)


# Patch Paragraph to resolve references at run-creation time via module-level registry
_ACTIVE_REGISTRY: LabelRegistry | None = None


def set_active_registry(registry: LabelRegistry | None):
    global _ACTIVE_REGISTRY
    _ACTIVE_REGISTRY = registry


def resolve_reference(type_: str, name: str) -> str:
    if _ACTIVE_REGISTRY is None:
        return f"{type_}?"
    display = _ACTIVE_REGISTRY.get_display(type_, name)
    return display or f"{type_}?"