"""
DOCX template rendering utilities for printing.

This repo's print templates (e.g. `print-template.docx`) use placeholder tokens like:
- `#{打印-销售单（8*9）.客户名称}`
- `{备注}`

Important: Microsoft Word / WPS may split a token across multiple "runs" (w:r / w:t).
So we must replace placeholders on the concatenated text of a paragraph/cell, then
write back to runs to avoid leaving raw tokens in the output.
"""

from __future__ import annotations

import copy
import re
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from typing import Any

from docx.document import Document as _DocumentType
from docx.table import Table, _Cell
from docx.text.paragraph import Paragraph


_RE_HASH_PLACEHOLDER = re.compile(r"#\{(.{1,300}?)\}", re.S)
# Plain `{xxx}` but not `#{xxx}` (negative lookbehind on '#')
_RE_BRACE_PLACEHOLDER = re.compile(r"(?<!#)\{([^{}\r\n]{1,200})\}")


def _normalize_placeholder_inner(s: str) -> str:
    # In real templates, tokens may be split with newlines/spaces across runs.
    # Normalizing by removing all whitespace makes keys stable for mapping.
    return re.sub(r"\s+", "", (s or "")).strip()


def _normalize_context(context: Mapping[str, Any]) -> dict[str, str]:
    """
    Normalize context keys by stripping whitespace, and stringify values.
    """
    m: dict[str, str] = {}
    for k, v in (context or {}).items():
        kk = _normalize_placeholder_inner(str(k))
        m[kk] = "" if v is None else str(v)
    return m


@dataclass(frozen=True)
class DocxRenderStats:
    paragraphs_touched: int
    placeholders_replaced_or_cleared: int


def _iter_paragraphs_in_table(table: Table) -> Iterable[Paragraph]:
    for row in table.rows:
        for cell in row.cells:
            yield from _iter_paragraphs_in_cell(cell)


def _iter_paragraphs_in_cell(cell: _Cell) -> Iterable[Paragraph]:
    for p in cell.paragraphs:
        yield p
    for t in cell.tables:
        yield from _iter_paragraphs_in_table(t)


def _iter_all_paragraphs_in_doc(d: _DocumentType) -> Iterable[Paragraph]:
    for p in d.paragraphs:
        yield p
    for t in d.tables:
        yield from _iter_paragraphs_in_table(t)

    # headers / footers may also contain placeholders
    for sec in d.sections:
        for p in sec.header.paragraphs:
            yield p
        for t in sec.header.tables:
            yield from _iter_paragraphs_in_table(t)
        for p in sec.footer.paragraphs:
            yield p
        for t in sec.footer.tables:
            yield from _iter_paragraphs_in_table(t)


def _replace_placeholders_in_paragraphs(
    *,
    paragraphs: Iterable[Paragraph],
    context: Mapping[str, Any],
    clear_unresolved: bool,
) -> DocxRenderStats:
    ctx = _normalize_context(context)
    paragraphs_touched = 0
    placeholders_done = 0

    def replace_in_text(text: str) -> tuple[str, int]:
        changed = 0

        def repl_hash(m: re.Match[str]) -> str:
            nonlocal changed
            raw_inner = m.group(1)
            key = _normalize_placeholder_inner(raw_inner)
            if key in ctx:
                changed += 1
                return ctx[key]
            if clear_unresolved:
                changed += 1
                return ""
            return m.group(0)

        def repl_brace(m: re.Match[str]) -> str:
            nonlocal changed
            raw_inner = m.group(1)
            key = _normalize_placeholder_inner(raw_inner)
            if key in ctx:
                changed += 1
                return ctx[key]
            if clear_unresolved:
                changed += 1
                return ""
            return m.group(0)

        new_text = _RE_HASH_PLACEHOLDER.sub(repl_hash, text)
        new_text = _RE_BRACE_PLACEHOLDER.sub(repl_brace, new_text)
        return new_text, changed

    def write_back_runs_keep_styles(runs, new_text: str) -> None:
        """
        Best-effort to keep run styles:
        - Keep the same number of runs
        - Distribute new text by original run lengths (last run gets the remainder)
        """
        orig_lens = [len(r.text or "") for r in runs]
        total_len = sum(orig_lens)
        if total_len <= 0:
            runs[0].text = new_text
            for r in runs[1:]:
                r.text = ""
            return

        # If new text is shorter, later runs become empty; if longer, last run grows.
        pos = 0
        for idx, r in enumerate(runs):
            if idx == len(runs) - 1:
                r.text = new_text[pos:]
                pos = len(new_text)
            else:
                n = orig_lens[idx]
                r.text = new_text[pos : pos + n]
                pos += n
        # Any remaining runs after pos are already set (some may become "").

    def replace_in_paragraph(p: Paragraph) -> int:
        nonlocal paragraphs_touched, placeholders_done
        runs = list(p.runs)
        if not runs:
            return 0
        text = "".join(r.text or "" for r in runs)
        # quick filter to skip most paragraphs
        if "#{" not in text and "}" not in text and "{" not in text:
            return 0

        new_text, changed = replace_in_text(text)
        if changed <= 0 or new_text == text:
            return 0

        write_back_runs_keep_styles(runs, new_text)
        paragraphs_touched += 1
        placeholders_done += changed
        return changed

    for p in paragraphs:
        replace_in_paragraph(p)

    return DocxRenderStats(
        paragraphs_touched=int(paragraphs_touched),
        placeholders_replaced_or_cleared=int(placeholders_done),
    )


def extract_docx_placeholders(*, template_path: str) -> set[str]:
    """
    Extract placeholder keys from a .docx file (document + header/footer parts if present).

    Returned keys are normalized (all whitespace removed) so they can be used as mapping keys.
    """
    import zipfile
    import xml.etree.ElementTree as ET

    ns = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
    keys: set[str] = set()

    def collect_from_xml(xml_text: str) -> None:
        try:
            root = ET.fromstring(xml_text)
        except Exception:
            return
        texts = [(el.text or "") for el in root.findall(".//w:t", ns)]
        full = "".join(texts)
        for inner in _RE_HASH_PLACEHOLDER.findall(full):
            k = _normalize_placeholder_inner(inner)
            if k:
                keys.add(k)
        for inner in _RE_BRACE_PLACEHOLDER.findall(full):
            k = _normalize_placeholder_inner(inner)
            if k:
                keys.add(k)

    with zipfile.ZipFile(template_path, "r") as zf:
        for name in zf.namelist():
            if not name.startswith("word/"):
                continue
            if not name.endswith(".xml"):
                continue
            # Focus on common content parts; still safe if more exist.
            if not (
                name == "word/document.xml"
                or name.startswith("word/header")
                or name.startswith("word/footer")
            ):
                continue
            try:
                xml = zf.read(name).decode("utf-8", errors="replace")
            except Exception:
                continue
            collect_from_xml(xml)

    return keys


def render_docx_template(
    *,
    template_path: str,
    output_path: str,
    context: Mapping[str, Any],
    clear_unresolved: bool = True,
) -> DocxRenderStats:
    """
    Render a .docx template with placeholder replacement.

    - Replace `#{...}` and `{...}` tokens using `context` (key = inner token).
    - If a token is not found in context:
      - cleared to empty string when `clear_unresolved=True`
      - kept as-is when `clear_unresolved=False`

    Notes:
    - This function intentionally handles tokens split across runs by replacing on
      concatenated text and writing back to runs (first run gets full text; others cleared).
    """
    from docx import Document

    doc: _DocumentType = Document(template_path)
    stats = _replace_placeholders_in_paragraphs(
        paragraphs=_iter_all_paragraphs_in_doc(doc),
        context=context,
        clear_unresolved=bool(clear_unresolved),
    )

    doc.save(output_path)
    return stats


def render_docx_template_pages(
    *,
    template_path: str,
    output_path: str,
    pages_contexts: list[Mapping[str, Any]],
    clear_unresolved: bool = True,
) -> DocxRenderStats:
    """
    Render multiple pages from a single-page docx template by duplicating the template
    table(s) for each page and filling placeholders per-page.

    Assumption (true for this repo's `print-template.docx`):
    - The whole page layout is in the document body tables.
    - Each page uses the same placeholder names; we fill them page-by-page.
    """
    if not pages_contexts:
        raise ValueError("pages_contexts 不能为空")

    from docx import Document

    doc: _DocumentType = Document(template_path)
    template_tables = list(doc.tables)
    if not template_tables:
        raise ValueError("模板 docx 不包含任何 table，无法按页复制")

    pages_tables: list[list[Table]] = []
    # page 1 uses existing tables
    pages_tables.append(list(template_tables))

    # subsequent pages: page break + deep copy template tables
    for _i in range(2, len(pages_contexts) + 1):
        doc.add_page_break()
        new_tables: list[Table] = []
        for t in template_tables:
            new_tbl_el = copy.deepcopy(t._tbl)  # pylint: disable=protected-access
            doc._body._body.append(new_tbl_el)  # type: ignore[attr-defined]  # pylint: disable=protected-access
            new_tables.append(Table(new_tbl_el, doc._body))  # type: ignore[arg-type]
        pages_tables.append(new_tables)

    total_stats = DocxRenderStats(paragraphs_touched=0, placeholders_replaced_or_cleared=0)

    for page_idx, (ctx, tbls) in enumerate(zip(pages_contexts, pages_tables), start=1):
        page_paragraphs: list[Paragraph] = []
        for t in tbls:
            page_paragraphs.extend(list(_iter_paragraphs_in_table(t)))
        st = _replace_placeholders_in_paragraphs(
            paragraphs=page_paragraphs,
            context=ctx,
            clear_unresolved=bool(clear_unresolved),
        )
        total_stats = DocxRenderStats(
            paragraphs_touched=total_stats.paragraphs_touched + st.paragraphs_touched,
            placeholders_replaced_or_cleared=total_stats.placeholders_replaced_or_cleared + st.placeholders_replaced_or_cleared,
        )

    doc.save(output_path)
    return total_stats