#!/usr/bin/env python3
"""将操作说明 docx 样式对齐健康评估操作说明参考文档（仅改样式，不改正文）。"""

import shutil
import zipfile
from pathlib import Path

from docx import Document
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Pt, RGBColor

REFERENCE_DOCX = Path(__file__).parent / "健康评估-操作说明.docx"
DEFAULT_TARGET_DOCX = Path(__file__).parent / "异常检测及风资源-操作说明.docx"
HEALTH_DOCX = Path(__file__).parent / "健康评估-操作说明.docx"

# 与交付样表一致：深蓝表头 + 浅蓝数据行 + 黑色边框
HEADING_FILL = "1F3864"
BODY_FILL = "EBF1F9"
HEADING_TEXT = RGBColor(0xFF, 0xFF, 0xFF)
BODY_TEXT = RGBColor(0x00, 0x00, 0x00)
BORDER_COLOR = "000000"

STYLE_REMAP = [
    ("Heading 3", "Heading 4"),
    ("Heading 2", "Heading 3"),
    ("Heading 1", "Heading 2"),
    ("Title", "Heading 1"),
]


def set_cell_shading(cell, fill_hex):
    tc_pr = cell._tc.get_or_add_tcPr()
    for child in tc_pr.findall(qn("w:shd")):
        tc_pr.remove(child)
    shd = OxmlElement("w:shd")
    shd.set(qn("w:val"), "clear")
    shd.set(qn("w:color"), "auto")
    shd.set(qn("w:fill"), fill_hex)
    tc_pr.append(shd)


def clear_run_font_override(run):
    r_pr = run._element.rPr
    if r_pr is None:
        return
    for tag in ("w:rFonts",):
        for node in r_pr.findall(qn(tag)):
            r_pr.remove(node)


def style_title_paragraph(paragraph):
    paragraph.paragraph_format.space_before = Pt(0)
    paragraph.paragraph_format.space_after = Pt(0)
    paragraph.paragraph_format.line_spacing = 1.15
    for run in paragraph.runs:
        clear_run_font_override(run)
        run.font.size = Pt(18)


def set_cell_border(cell, color=BORDER_COLOR, size="4"):
    """设置单元格四边黑色实线边框（size 单位为 1/8 pt）。"""
    tc_pr = cell._tc.get_or_add_tcPr()
    for child in tc_pr.findall(qn("w:tcBorders")):
        tc_pr.remove(child)
    borders = OxmlElement("w:tcBorders")
    for edge in ("top", "left", "bottom", "right"):
        element = OxmlElement(f"w:{edge}")
        element.set(qn("w:val"), "single")
        element.set(qn("w:sz"), size)
        element.set(qn("w:space"), "0")
        element.set(qn("w:color"), color)
        borders.append(element)
    tc_pr.append(borders)


def style_table_cell_runs(paragraph, *, bold=False, color=BODY_TEXT, align=WD_PARAGRAPH_ALIGNMENT.LEFT):
    paragraph.alignment = align
    for run in paragraph.runs:
        clear_run_font_override(run)
        run.font.bold = bold
        run.font.color.rgb = color
        run.font.size = Pt(11)


def style_all_tables(doc):
    for table in doc.tables:
        if not table.rows:
            continue
        table.style = "Table Grid"

        for row_idx, row in enumerate(table.rows):
            is_header = row_idx == 0
            for cell in row.cells:
                set_cell_border(cell)
                set_cell_shading(cell, HEADING_FILL if is_header else BODY_FILL)
                for paragraph in cell.paragraphs:
                    if is_header:
                        style_table_cell_runs(
                            paragraph,
                            bold=True,
                            color=HEADING_TEXT,
                            align=WD_PARAGRAPH_ALIGNMENT.CENTER,
                        )
                    else:
                        style_table_cell_runs(
                            paragraph,
                            bold=False,
                            color=BODY_TEXT,
                            align=WD_PARAGRAPH_ALIGNMENT.LEFT,
                        )


def remap_heading_styles(doc):
    for old_style, new_style in STYLE_REMAP:
        for paragraph in doc.paragraphs:
            if paragraph.style.name == old_style:
                paragraph.style = doc.styles[new_style]


def copy_reference_package_parts(target_docx):
    if not REFERENCE_DOCX.exists():
        raise FileNotFoundError(f"参考文档不存在: {REFERENCE_DOCX}")

    backup = target_docx.with_suffix(".docx.bak")
    shutil.copy2(target_docx, backup)

    temp_path = target_docx.with_suffix(".docx.tmp")
    replace_parts = {
        "word/styles.xml",
        "word/fontTable.xml",
        "word/theme/theme1.xml",
    }

    with zipfile.ZipFile(target_docx, "r") as src_zip:
        with zipfile.ZipFile(REFERENCE_DOCX, "r") as ref_zip:
            ref_parts = {name: ref_zip.read(name) for name in replace_parts}

        with zipfile.ZipFile(temp_path, "w", zipfile.ZIP_DEFLATED) as out_zip:
            for item in src_zip.infolist():
                data = ref_parts[item.filename] if item.filename in ref_parts else src_zip.read(item.filename)
                out_zip.writestr(item, data)

    temp_path.replace(target_docx)
    return backup


def restyle_docx(target_docx=DEFAULT_TARGET_DOCX, copy_reference=True):
    target_docx = Path(target_docx)
    backup = None
    if copy_reference:
        backup = copy_reference_package_parts(target_docx)
    doc = Document(str(target_docx))

    if copy_reference:
        remap_heading_styles(doc)
        for paragraph in doc.paragraphs:
            if paragraph.style.name == "Heading 1" and paragraph.text.strip():
                style_title_paragraph(paragraph)
                break

    style_all_tables(doc)
    doc.save(str(target_docx))
    print(f"Restyled: {target_docx}")
    if backup:
        print(f"Backup: {backup}")


def restyle_tables_only(target_docx=DEFAULT_TARGET_DOCX):
    target_docx = Path(target_docx)
    doc = Document(str(target_docx))
    style_all_tables(doc)
    doc.save(str(target_docx))
    print(f"Table style applied: {target_docx}")


if __name__ == "__main__":
    import sys

    args = sys.argv[1:]
    tables_only = False
    target = DEFAULT_TARGET_DOCX

    for arg in args:
        if arg == "--tables-only":
            tables_only = True
        elif arg.endswith(".docx"):
            target = Path(arg)
            if not target.is_absolute():
                target = Path(__file__).parent / target.name

    if tables_only:
        restyle_tables_only(target)
    else:
        restyle_docx(target)