# -*- coding: utf-8 -*-
"""Restore delivery docx cover / document-control formatting (S0 blueprint style)."""
from __future__ import annotations

from pathlib import Path

from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Pt, RGBColor
from docx.table import Table
from docx.text.paragraph import Paragraph

BASE = Path(
    r"C:\Users\skygu\OneDrive\Projects\AIDOP\项目\项目管理\产互联项目管理\交付文档"
)
FONT = "微软雅黑"
BLUE = RGBColor(0x1F, 0x4E, 0x79)
HEADER_FILL = "1F3864"
TABLE_STYLE = "aff1"  # Table Grid in this template


def body_blocks(doc: Document):
    for child in doc.element.body:
        if child.tag.endswith("p"):
            yield "p", Paragraph(child, doc)
        elif child.tag.endswith("tbl"):
            yield "t", Table(child, doc)


def find_toc_element(doc: Document):
    for kind, obj in body_blocks(doc):
        if kind == "p" and obj.text.strip() == "目录":
            return obj._element
    return None


def set_run_font(run, *, size: Pt | None = None, bold: bool | None = None, color: RGBColor | None = None):
    run.font.name = FONT
    r_pr = run._element.get_or_add_rPr()
    r_fonts = r_pr.rFonts
    if r_fonts is None:
        r_fonts = OxmlElement("w:rFonts")
        r_pr.insert(0, r_fonts)
    r_fonts.set(qn("w:ascii"), FONT)
    r_fonts.set(qn("w:hAnsi"), FONT)
    r_fonts.set(qn("w:eastAsia"), FONT)
    if size is not None:
        run.font.size = size
    if bold is not None:
        run.bold = bold
    if color is not None:
        run.font.color.rgb = color


def format_title_paragraph(p: Paragraph, text: str, *, size: Pt, bold: bool = True, color: RGBColor = BLUE):
    p.text = text
    p.alignment = WD_ALIGN_PARAGRAPH.CENTER
    for run in p.runs:
        set_run_font(run, size=size, bold=bold, color=color)


def format_system_paragraph(p: Paragraph, text: str):
    p.text = text
    p.alignment = WD_ALIGN_PARAGRAPH.CENTER
    for run in p.runs:
        set_run_font(run, size=Pt(13), bold=False, color=RGBColor(0, 0, 0))


def apply_table_style(table: Table) -> None:
    tbl = table._tbl
    tbl_pr = tbl.tblPr
    if tbl_pr is None:
        tbl_pr = OxmlElement("w:tblPr")
        tbl.insert(0, tbl_pr)
    style = tbl_pr.find(qn("w:tblStyle"))
    if style is None:
        style = OxmlElement("w:tblStyle")
        tbl_pr.insert(0, style)
    style.set(qn("w:val"), TABLE_STYLE)
    jc = tbl_pr.find(qn("w:jc"))
    if jc is None:
        jc = OxmlElement("w:jc")
        jc.set(qn("w:val"), "center")
        tbl_pr.append(jc)
    else:
        jc.set(qn("w:val"), "center")


def shade_cell(cell, fill: str) -> None:
    tc_pr = cell._element.get_or_add_tcPr()
    old = tc_pr.find(qn("w:shd"))
    if old is not None:
        tc_pr.remove(old)
    shd = OxmlElement("w:shd")
    shd.set(qn("w:val"), "clear")
    shd.set(qn("w:color"), "auto")
    shd.set(qn("w:fill"), fill)
    tc_pr.append(shd)


def set_cell(
    cell,
    text: str,
    *,
    bold: bool = False,
    white: bool = False,
    size: Pt = Pt(10.5),
    center: bool = False,
):
    cell.text = ""
    p = cell.paragraphs[0]
    if center:
        p.alignment = WD_ALIGN_PARAGRAPH.CENTER
    run = p.add_run(text)
    set_run_font(run, size=size, bold=bold, color=RGBColor(255, 255, 255) if white else RGBColor(0, 0, 0))


def format_cover_table(table: Table) -> None:
    apply_table_style(table)
    for row in table.rows:
        set_cell(row.cells[0], row.cells[0].text.strip(), bold=True)
        set_cell(row.cells[1], row.cells[1].text.strip())


def format_change_table(table: Table) -> None:
    apply_table_style(table)
    header = table.rows[0]
    for cell in header.cells:
        shade_cell(cell, HEADER_FILL)
        set_cell(cell, cell.text.strip(), bold=True, white=True, center=True)
    for row in table.rows[1:]:
        for cell in row.cells:
            set_cell(cell, cell.text.strip(), center=True)


def ensure_leading_blank(doc: Document, toc_el) -> None:
    body = doc.element.body
    children = list(body)
    toc_pos = children.index(toc_el)
    if toc_pos == 0:
        blank = OxmlElement("w:p")
        body.insert(0, blank)
        return
    first = children[0]
    if first.tag.endswith("p"):
        p = Paragraph(first, doc)
        if p.text.strip():
            blank = OxmlElement("w:p")
            body.insert(0, blank)


def insert_page_break_before_doc_control(doc: Document, toc_el) -> None:
    from docx.enum.text import WD_BREAK

    for kind, obj in body_blocks(doc):
        if kind == "p" and obj.text.strip() == "文档控制":
            if obj.runs:
                obj.runs[0].add_break(WD_BREAK.PAGE)
            else:
                obj.add_run().add_break(WD_BREAK.PAGE)
            return


def restore_file(path: Path) -> None:
    doc = Document(path)
    toc_el = find_toc_element(doc)
    if toc_el is None:
        raise RuntimeError(f"未找到目录: {path.name}")

    front_paras: list[Paragraph] = []
    front_tables: list[Table] = []
    for kind, obj in body_blocks(doc):
        if obj._element is toc_el:
            break
        if kind == "p":
            t = obj.text.strip()
            if t:
                front_paras.append(obj)
        elif kind == "t":
            front_tables.append(obj)

    if len(front_paras) >= 3:
        format_title_paragraph(front_paras[0], front_paras[0].text.strip(), size=Pt(28))
        format_title_paragraph(front_paras[1], front_paras[1].text.strip(), size=Pt(22))
        format_system_paragraph(front_paras[2], front_paras[2].text.strip())

    for table in front_tables:
        if not table.rows:
            continue
        cols = len(table.columns)
        if cols == 2 and len(table.rows) == 5:
            format_cover_table(table)
        elif cols == 4 and len(table.rows) == 3:
            format_change_table(table)

    ensure_leading_blank(doc, toc_el)
    insert_page_break_before_doc_control(doc, toc_el)
    doc.save(path)


def main() -> None:
    for module in ["S0", "S1", "S2", "S3", "S4"]:
        for path in sorted((BASE / module).glob("*.docx")):
            restore_file(path)
            print(f"restored {module}/{path.name}")


if __name__ == "__main__":
    main()