md_to_gost/md2gost/page_geometry.py

"""A4 page geometry helpers for portrait / landscape sections."""

from __future__ import annotations

import os
import tempfile
import zipfile
from pathlib import Path

from docx.enum.section import WD_ORIENT
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.oxml.ns import qn
from docx.shared import Mm, Pt
from lxml import etree

from .util import create_element

_W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
_W = f"{{{_W_NS}}}"

# GOST-like A4 margins (same as styles._apply_common_page_and_body)
MARGIN_LEFT = Mm(30)
MARGIN_RIGHT = Mm(10)
MARGIN_TOP = Mm(20)
MARGIN_BOTTOM = Mm(20)

# mirea / pis_custom — 30/10/20/20; paco_custom — форма ПАЦО 20/20/20/20
MARGIN_PRESETS = {
    "mirea": (Mm(30), Mm(10), Mm(20), Mm(20)),
    "pis_custom": (Mm(30), Mm(10), Mm(20), Mm(20)),
}


def apply_margin_preset(preset: str) -> None:
    """Switch active page margins used by apply_section_geometry / content_size."""
    global MARGIN_LEFT, MARGIN_RIGHT, MARGIN_TOP, MARGIN_BOTTOM
    left, right, top, bottom = MARGIN_PRESETS.get(preset, MARGIN_PRESETS["mirea"])
    MARGIN_LEFT = left
    MARGIN_RIGHT = right
    MARGIN_TOP = top
    MARGIN_BOTTOM = bottom

A4_SHORT = Mm(210)
A4_LONG = Mm(297)


def is_landscape_section(section) -> bool:
    """True if section is (or should be treated as) landscape A4."""
    try:
        if section.orientation == WD_ORIENT.LANDSCAPE:
            return True
    except Exception:
        pass
    return int(section.page_width) > int(section.page_height)


def apply_section_geometry(section, *, landscape: bool) -> None:
    """
    Set orientation and page size.

    Set orientation first (python-docx may swap w/h on change), then force A4 dims.
    """
    target = WD_ORIENT.LANDSCAPE if landscape else WD_ORIENT.PORTRAIT
    try:
        section.orientation = target
    except Exception:
        pass

    if landscape:
        section.page_width = A4_LONG
        section.page_height = A4_SHORT
    else:
        section.page_width = A4_SHORT
        section.page_height = A4_LONG

    section.left_margin = MARGIN_LEFT
    section.right_margin = MARGIN_RIGHT
    section.top_margin = MARGIN_TOP
    section.bottom_margin = MARGIN_BOTTOM

    # Explicit orient on pgSz for Word.
    pg_sz = section._sectPr.find(qn("w:pgSz"))
    if pg_sz is None:
        pg_sz = section._sectPr._add_pgSz()
    if landscape:
        pg_sz.set(qn("w:orient"), "landscape")
        # Re-assert after XML tweak (some builds reshuffle).
        section.page_width = A4_LONG
        section.page_height = A4_SHORT
    else:
        if pg_sz.get(qn("w:orient")) is not None:
            del pg_sz.attrib[qn("w:orient")]
        section.page_width = A4_SHORT
        section.page_height = A4_LONG

    # Vertical align: center on landscape (figures/tables only); top on portrait
    sect_pr = section._sectPr
    for el in list(sect_pr.findall(qn("w:vAlign"))):
        sect_pr.remove(el)
    if landscape:
        sect_pr.append(create_element("w:vAlign", {"w:val": "center"}))

    # add_section copies template sectPr including w:pgNumType/@w:start — strip so
    # landscape / return-to-portrait sections continue numbering instead of restarting.
    clear_section_page_start(section)


def content_size(*, landscape: bool) -> tuple:
    """Return (max_height, max_width) usable content area for LayoutTracker."""
    if landscape:
        page_w, page_h = A4_LONG, A4_SHORT
    else:
        page_w, page_h = A4_SHORT, A4_LONG
    max_height = page_h - MARGIN_TOP - MARGIN_BOTTOM
    max_width = page_w - MARGIN_LEFT - MARGIN_RIGHT
    return max_height, max_width


def apply_centered_page_footer(section) -> None:
    """Centered PAGE field, Times New Roman 12 (body section style)."""
    footer = section.footer
    footer.is_linked_to_previous = False
    if not footer.paragraphs:
        footer.add_paragraph()
    paragraph = footer.paragraphs[0]
    paragraph.clear()
    paragraph.paragraph_format.first_line_indent = 0
    paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    run = paragraph.add_run()
    run.font.name = "Times New Roman"
    run.font.size = Pt(12)
    paragraph._p.append(create_element("w:fldSimple", {
        "w:instr": "PAGE \\* MERGEFORMAT",
    }))


def clear_section_footer(section) -> None:
    """Empty footer (no PAGE) — for title / assignment / TOC sections."""
    footer = section.footer
    footer.is_linked_to_previous = False
    for p in footer.paragraphs:
        p.clear()
    if not footer.paragraphs:
        footer.add_paragraph()


def clear_section_page_start(section) -> None:
    """Remove w:pgNumType/@w:start so PAGE continues from the previous section."""
    sect_pr = section._sectPr
    for child in list(sect_pr):
        if child.tag != qn("w:pgNumType"):
            continue
        if child.get(qn("w:start")) is not None:
            del child.attrib[qn("w:start")]
        # Empty pgNumType (no start/fmt) is useless — drop the element.
        if len(child.attrib) == 0:
            sect_pr.remove(child)


def ensure_continuous_page_numbers(document) -> None:
    """
    After docxcompose of title/assignment: no PAGE on early sections,
    continuous numbering (do not restart at 1 on body section).
    """
    sections = list(document.sections)
    if not sections:
        return

    # Strip w:pgNumType start everywhere so landscape / body sections stay continuous.
    for section in sections:
        clear_section_page_start(section)


def set_section_page_start(section, start: int) -> None:
    """Set w:pgNumType/@w:start so PAGE field begins at ``start`` in this section."""
    if start < 1:
        return
    sect_pr = section._sectPr
    pg = sect_pr.find(qn("w:pgNumType"))
    if pg is None:
        sect_pr.append(create_element("w:pgNumType", {"w:start": str(int(start))}))
    else:
        pg.set(qn("w:start"), str(int(start)))


def apply_page_number_start(
    document,
    start: int | None,
    *,
    front_sections: int = 0,
) -> None:
    """
    Apply page-number offset for the body.

    ``start`` is None / <1 — leave continuous numbering (strip explicit starts).
    ``start`` >= 1 — first section after front matter restarts PAGE at that number;
    later body sections stay continuous from there.
    """
    ensure_continuous_page_numbers(document)
    if start is None or int(start) < 1:
        return
    sections = list(document.sections)
    idx = max(0, int(front_sections))
    if idx >= len(sections):
        return
    set_section_page_start(sections[idx], int(start))


def _iter_sect_pr(document_root) -> list:
    """Return sectPr elements in document order (inline + final body sectPr)."""
    body = document_root.find(f"{_W}body")
    if body is None:
        return []
    sects: list = []
    for child in body:
        tag = child.tag
        if tag == f"{_W}sectPr":
            sects.append(child)
            continue
        if tag != f"{_W}p":
            continue
        p_pr = child.find(f"{_W}pPr")
        if p_pr is None:
            continue
        sect = p_pr.find(f"{_W}sectPr")
        if sect is not None:
            sects.append(sect)
    return sects


def _strip_pg_num_start_xml(sect_pr) -> None:
    for child in list(sect_pr):
        if child.tag != f"{_W}pgNumType":
            continue
        if child.get(f"{_W}start") is not None:
            del child.attrib[f"{_W}start"]
        if len(child.attrib) == 0:
            sect_pr.remove(child)


def _set_pg_num_start_xml(sect_pr, start: int) -> None:
    pg = sect_pr.find(f"{_W}pgNumType")
    if pg is None:
        pg = etree.SubElement(sect_pr, f"{_W}pgNumType")
    pg.set(f"{_W}start", str(int(start)))


def patch_docx_page_starts(
    path: str | Path,
    start: int | None = None,
    *,
    front_sections: int = 0,
) -> None:
    """
    Strip (and optionally re-apply) w:pgNumType/@w:start inside word/document.xml.

    Used after Word COM Save, which often rewrites start=\"1\" onto every section.
    Zip/lxml only — does not round-trip through python-docx (preserves
    continuation captions and other Word-written parts).
    """
    abs_path = os.path.abspath(str(path))
    if not os.path.isfile(abs_path):
        return

    with zipfile.ZipFile(abs_path, "r") as zin:
        names = zin.namelist()
        if "word/document.xml" not in names:
            return
        doc_bytes = zin.read("word/document.xml")
        other = {n: zin.read(n) for n in names if n != "word/document.xml"}

    root = etree.fromstring(doc_bytes)
    sects = _iter_sect_pr(root)
    for sect in sects:
        _strip_pg_num_start_xml(sect)

    if start is not None and int(start) >= 1:
        idx = max(0, int(front_sections))
        if idx < len(sects):
            _set_pg_num_start_xml(sects[idx], int(start))

    new_doc = etree.tostring(
        root,
        xml_declaration=True,
        encoding="UTF-8",
        standalone=True,
    )

    parent = os.path.dirname(abs_path) or "."
    fd, tmp_name = tempfile.mkstemp(suffix=".docx", dir=parent)
    os.close(fd)
    try:
        with zipfile.ZipFile(tmp_name, "w", compression=zipfile.ZIP_DEFLATED) as zout:
            for name, data in other.items():
                zout.writestr(name, data)
            zout.writestr("word/document.xml", new_doc)
        os.replace(tmp_name, abs_path)
    except Exception:
        try:
            os.remove(tmp_name)
        except OSError:
            pass
        raise