"""Default post-processing for md→docx conversion.

This script is loaded dynamically by mytoolkit.convert after pandoc has
produced the docx. The entry point is ``postprocess(docx_path, reference_doc_path)``.
"""

from pathlib import Path


def postprocess(docx_path: Path, reference_doc_path: str | None = None) -> None:
    """Post-process pandoc-generated docx for Chinese academic typography.

    Applies the following fixes:
    - Center image paragraphs and auto-generate captions (图 1, 图 2, ...)
    - Inject Caption style into styles.xml
    - Apply paragraph-level typography from reference-doc Normal style
      (first-line indent, spacing, line spacing, justification)
    - Convert BodyText paragraphs to Normal and apply same typography
    - Add heading spacing before/after Heading 1-4
    - Add 6pt spacing after reference entries (detected after "参考文献")
    - Three-line table borders (1.5pt top/bottom, 0.75pt header separator)
    - Center tables on the page
    - Remove table styles that override custom borders
    - Remove first-line indent inside table cells
    - Add 6pt spacing before paragraphs following tables
    - Convert '---' horizontal rules to HorizontalRule style (no indent)
    - Remove first-line indent from empty paragraphs
    - Inject HorizontalRule style into styles.xml
    """
    try:
        import lxml.etree as ET
    except ImportError:
        return

    import zipfile
    import os
    import copy
    from tempfile import TemporaryDirectory

    W = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
    wtag = lambda t: "{" + W + "}" + t

    def _paragraph_text(p):
        parts = []
        for r in p.findall(wtag("r")):
            t = r.find(wtag("t"))
            if t is not None and t.text:
                parts.append(t.text)
        return "".join(parts)

    def _copy_ppr_children(src_ppr, dst_ppr):
        for child in src_ppr:
            tag_name = child.tag.split("}")[-1]
            # Remove existing same-tag elements
            for existing in list(dst_ppr.findall(wtag(tag_name))):
                dst_ppr.remove(existing)
            dst_ppr.append(copy.deepcopy(child))

    with TemporaryDirectory() as workdir:
        with zipfile.ZipFile(docx_path, "r") as z:
            z.extractall(workdir)

        # --- Load reference-doc Normal style paragraph properties ---
        normal_ppr = None
        if reference_doc_path and Path(reference_doc_path).exists():
            try:
                with zipfile.ZipFile(reference_doc_path, "r") as z:
                    styles_content = z.read("word/styles.xml")
                ref_styles_root = ET.fromstring(styles_content)
                for style in ref_styles_root.findall(wtag("style")):
                    if style.get(wtag("styleId")) == "Normal":
                        normal_ppr = style.find(wtag("pPr"))
                        break
            except Exception:
                normal_ppr = None

        # Heading spacing in twips (1pt = 20 twips)
        heading_spacing = {
            "Heading1": {"before": "360", "after": "120"},
            "Heading2": {"before": "280", "after": "120"},
            "Heading3": {"before": "200", "after": "80"},
            "Heading4": {"before": "160", "after": "80"},
        }

        doc_xml = os.path.join(workdir, "word", "document.xml")
        tree = ET.parse(doc_xml)
        root = tree.getroot()
        body = root.find(".//" + wtag("body"))
        if body is None:
            return

        # --- 1. Table borders (three-line) + centering ---
        for tbl in body.findall(".//" + wtag("tbl")):
            tblPr = tbl.find(wtag("tblPr"))
            if tblPr is None:
                tblPr = ET.SubElement(tbl, wtag("tblPr"))
                tbl.insert(0, tblPr)
            # Remove table style so borders are not overridden
            for ts in list(tblPr.findall(wtag("tblStyle"))):
                tblPr.remove(ts)
            # Center the table on the page.
            for existing in list(tblPr.findall(wtag("jc"))):
                tblPr.remove(existing)
            jc = ET.SubElement(tblPr, wtag("jc"))
            jc.set(wtag("val"), "center")
            for existing in list(tblPr.findall(wtag("tblBorders"))):
                tblPr.remove(existing)
            borders = ET.SubElement(tblPr, wtag("tblBorders"))
            for pos, sz in [("top", "12"), ("bottom", "12")]:
                b = ET.SubElement(borders, wtag(pos))
                b.set(wtag("val"), "single")
                b.set(wtag("sz"), sz)
                b.set(wtag("space"), "0")
                b.set(wtag("color"), "auto")
            for pos in ["left", "right", "insideH", "insideV"]:
                b = ET.SubElement(borders, wtag(pos))
                b.set(wtag("val"), "none")
                b.set(wtag("sz"), "0")
                b.set(wtag("space"), "0")
                b.set(wtag("color"), "auto")
            # Header row bottom border
            rows = tbl.findall(wtag("tr"))
            if rows:
                for tc in rows[0].findall(wtag("tc")):
                    tcPr = tc.find(wtag("tcPr"))
                    if tcPr is None:
                        tcPr = ET.SubElement(tc, wtag("tcPr"))
                    for existing in list(tcPr.findall(wtag("tcBorders"))):
                        tcPr.remove(existing)
                    tcBorders = ET.SubElement(tcPr, wtag("tcBorders"))
                    for pos in ["top", "left", "right"]:
                        b = ET.SubElement(tcBorders, wtag(pos))
                        b.set(wtag("val"), "none")
                        b.set(wtag("sz"), "0")
                        b.set(wtag("space"), "0")
                        b.set(wtag("color"), "auto")
                    bottom = ET.SubElement(tcBorders, wtag("bottom"))
                    bottom.set(wtag("val"), "single")
                    bottom.set(wtag("sz"), "6")
                    bottom.set(wtag("space"), "0")
                    bottom.set(wtag("color"), "auto")

        # --- 2. Image captions and centering ---
        def _set_paragraph_alignment(p, align):
            pPr = p.find(wtag("pPr"))
            if pPr is None:
                pPr = ET.SubElement(p, wtag("pPr"))
                p.insert(0, pPr)
            for jc in list(pPr.findall(wtag("jc"))):
                pPr.remove(jc)
            jc = ET.SubElement(pPr, wtag("jc"))
            jc.set(wtag("val"), align)

        def _has_drawing(p):
            return p.find(".//" + wtag("drawing")) is not None

        drawing_ns = "http://schemas.openxmlformats.org/drawingml/2006/main"
        image_count = 0
        # Iterate over a snapshot because we insert caption paragraphs and may remove ImageCaption.
        for p in list(body.findall(wtag("p"))):
            if not _has_drawing(p):
                continue
            # Center the image paragraph.
            _set_paragraph_alignment(p, "center")
            image_count += 1

            # Gather alt text: prefer pandoc-generated ImageCaption, fallback to drawing docPr.
            alt = ""
            p_idx = list(body).index(p)
            next_p = body[p_idx + 1] if p_idx + 1 < len(list(body)) else None
            if next_p is not None and next_p.tag == wtag("p"):
                next_style = next_p.find(".//" + wtag("pStyle"))
                if next_style is not None and next_style.get(wtag("val")) == "ImageCaption":
                    parts = []
                    for r in next_p.findall(wtag("r")):
                        t = r.find(wtag("t"))
                        if t is not None and t.text:
                            parts.append(t.text)
                    if parts:
                        alt = "".join(parts).strip()
                    body.remove(next_p)

            if not alt:
                drawing = p.find(".//" + wtag("drawing"))
                if drawing is not None:
                    docPr = drawing.find(f".//{{{drawing_ns}}}docPr")
                    if docPr is not None:
                        alt = docPr.get("descr", "").strip()

            caption_text = f"图 {image_count}"
            if alt:
                caption_text = f"图 {image_count}  {alt}"

            # Insert caption paragraph after the image paragraph.
            caption_p = ET.Element(wtag("p"))
            caption_pPr = ET.SubElement(caption_p, wtag("pPr"))
            ET.SubElement(caption_pPr, wtag("pStyle")).set(wtag("val"), "Caption")
            _set_paragraph_alignment(caption_p, "center")
            ET.SubElement(caption_pPr, wtag("ind")).set(wtag("firstLine"), "0")
            spacing = ET.SubElement(caption_pPr, wtag("spacing"))
            spacing.set(wtag("before"), "60")
            spacing.set(wtag("after"), "120")
            r = ET.SubElement(caption_p, wtag("r"))
            rPr = ET.SubElement(r, wtag("rPr"))
            ET.SubElement(rPr, wtag("sz")).set(wtag("val"), "22")
            ET.SubElement(rPr, wtag("szCs")).set(wtag("val"), "22")
            t = ET.SubElement(r, wtag("t"))
            t.text = caption_text
            # Recompute index because we may have removed the next paragraph.
            p_idx = list(body).index(p)
            body.insert(p_idx + 1, caption_p)

        # --- 3. Remove indent in table cells ---
        for tc in body.findall(".//" + wtag("tc")):
            for p in tc.findall(wtag("p")):
                pPr = p.find(wtag("pPr"))
                if pPr is None:
                    pPr = ET.SubElement(p, wtag("pPr"))
                for ind in list(pPr.findall(wtag("ind"))):
                    pPr.remove(ind)
                ET.SubElement(pPr, wtag("ind")).set(wtag("firstLine"), "0")
                spacing = pPr.find(wtag("spacing"))
                if spacing is not None:
                    spacing.set(wtag("before"), "0")

        # --- 4. Horizontal rule + empty paragraphs ---
        for p in body.findall(wtag("p")):
            p_str = ET.tostring(p, encoding="unicode")
            is_hr = "o:hr=" in p_str and "o:hrstd=" in p_str
            has_text = False
            for r in p.findall(wtag("r")):
                t = r.find(wtag("t"))
                if t is not None and t.text and t.text.strip():
                    has_text = True
                    break
            if is_hr:
                pPr = p.find(wtag("pPr"))
                if pPr is None:
                    pPr = ET.SubElement(p, wtag("pPr"))
                    p.insert(0, pPr)
                for ps in list(pPr.findall(wtag("pStyle"))):
                    pPr.remove(ps)
                ET.SubElement(pPr, wtag("pStyle")).set(wtag("val"), "HorizontalRule")
                for ind in list(pPr.findall(wtag("ind"))):
                    pPr.remove(ind)
                ET.SubElement(pPr, wtag("ind")).set(wtag("firstLine"), "0")
                spacing = pPr.find(wtag("spacing"))
                if spacing is not None:
                    spacing.set(wtag("before"), "0")
            elif not has_text:
                pPr = p.find(wtag("pPr"))
                if pPr is None:
                    pPr = ET.SubElement(p, wtag("pPr"))
                for ind in list(pPr.findall(wtag("ind"))):
                    pPr.remove(ind)
                ET.SubElement(pPr, wtag("ind")).set(wtag("firstLine"), "0")
                spacing = pPr.find(wtag("spacing"))
                if spacing is not None:
                    spacing.set(wtag("before"), "0")

        # --- 5. Apply paragraph-level typography from reference doc ---
        in_references = False
        for p in body.findall(wtag("p")):
            # Skip paragraphs inside table cells (handled separately)
            parent = p.getparent()
            if parent is not None and parent.tag == wtag("tc"):
                continue

            pPr = p.find(wtag("pPr"))
            if pPr is None:
                pPr = ET.SubElement(p, wtag("pPr"))
                p.insert(0, pPr)

            pStyle = pPr.find(wtag("pStyle"))
            style_id = pStyle.get(wtag("val")) if pStyle is not None else "Normal"

            # Detect references section by heading text
            if style_id.startswith("Heading") and "参考文献" in _paragraph_text(p):
                in_references = True
                # Apply heading spacing
                hs = heading_spacing.get(style_id)
                if hs is not None:
                    spacing = pPr.find(wtag("spacing"))
                    if spacing is None:
                        spacing = ET.SubElement(pPr, wtag("spacing"))
                    spacing.set(wtag("before"), hs["before"])
                    spacing.set(wtag("after"), hs["after"])
                continue

            # Reference entry formatting
            if in_references:
                text = _paragraph_text(p).strip()
                if not text:
                    # Empty paragraph in refs: remove indent but keep normal spacing
                    for ind in list(pPr.findall(wtag("ind"))):
                        pPr.remove(ind)
                    ET.SubElement(pPr, wtag("ind")).set(wtag("firstLine"), "0")
                    continue
                # Convert to Normal and apply normal pPr
                if pStyle is None:
                    pStyle = ET.SubElement(pPr, wtag("pStyle"))
                pStyle.set(wtag("val"), "Normal")
                if normal_ppr is not None:
                    _copy_ppr_children(normal_ppr, pPr)
                spacing = pPr.find(wtag("spacing"))
                if spacing is None:
                    spacing = ET.SubElement(pPr, wtag("spacing"))
                spacing.set(wtag("after"), "120")  # 6pt after each reference
                continue

            # BodyText / FirstParagraph -> Normal
            if style_id in ("BodyText", "FirstParagraph"):
                if pStyle is None:
                    pStyle = ET.SubElement(pPr, wtag("pStyle"))
                pStyle.set(wtag("val"), "Normal")
                style_id = "Normal"

            # Apply Normal paragraph properties from reference doc
            if style_id in ("Normal", "FirstParagraph") and normal_ppr is not None:
                _copy_ppr_children(normal_ppr, pPr)

            # Apply heading spacing
            if style_id in heading_spacing:
                spacing = pPr.find(wtag("spacing"))
                if spacing is None:
                    spacing = ET.SubElement(pPr, wtag("spacing"))
                spacing.set(wtag("before"), heading_spacing[style_id]["before"])
                spacing.set(wtag("after"), heading_spacing[style_id]["after"])

        # --- 6. Add spacing after tables ---
        for tbl in body.findall(wtag("tbl")):
            tbl_idx = list(body).index(tbl)
            for i in range(tbl_idx + 1, len(list(body))):
                next_el = list(body)[i]
                if next_el.tag == wtag("p"):
                    pPr = next_el.find(wtag("pPr"))
                    if pPr is None:
                        pPr = ET.SubElement(next_el, wtag("pPr"))
                    spacing = pPr.find(wtag("spacing"))
                    if spacing is None:
                        spacing = ET.SubElement(pPr, wtag("spacing"))
                    spacing.set(wtag("before"), "120")
                    break

        # --- 7. Inject suppressAutoHyphens into all paragraphs ---
        for p in body.findall(wtag("p")):
            pPr = p.find(wtag("pPr"))
            if pPr is None:
                pPr = ET.SubElement(p, wtag("pPr"))
                p.insert(0, pPr)
            if pPr.find(wtag("suppressAutoHyphens")) is None:
                ET.SubElement(pPr, wtag("suppressAutoHyphens"))

        tree.write(doc_xml, xml_declaration=True, encoding="UTF-8", pretty_print=False)

        # --- 8. Inject HorizontalRule and Caption styles ---
        styles_xml = os.path.join(workdir, "word", "styles.xml")
        styles_tree = ET.parse(styles_xml)
        styles_root = styles_tree.getroot()
        normal_style = None
        for style in styles_root.findall(wtag("style")):
            if style.get(wtag("styleId")) == "Normal":
                normal_style = style
                break
        if normal_style is not None:
            # Remove existing custom styles to avoid duplicates.
            for style in list(styles_root.findall(wtag("style"))):
                if style.get(wtag("styleId")) in ("HorizontalRule", "Caption"):
                    styles_root.remove(style)

            # HorizontalRule style
            hr_style = ET.Element(wtag("style"))
            hr_style.set(wtag("type"), "paragraph")
            hr_style.set(wtag("styleId"), "HorizontalRule")
            ET.SubElement(hr_style, wtag("name")).set(wtag("val"), "Horizontal Rule")
            ET.SubElement(hr_style, wtag("basedOn")).set(wtag("val"), "Normal")
            pPr = ET.SubElement(hr_style, wtag("pPr"))
            ET.SubElement(pPr, wtag("ind")).set(wtag("firstLine"), "0")

            # Caption style: centered, 11pt, no first-line indent
            caption_style = ET.Element(wtag("style"))
            caption_style.set(wtag("type"), "paragraph")
            caption_style.set(wtag("styleId"), "Caption")
            ET.SubElement(caption_style, wtag("name")).set(wtag("val"), "Caption")
            ET.SubElement(caption_style, wtag("basedOn")).set(wtag("val"), "Normal")
            ET.SubElement(caption_style, wtag("qFormat"))
            cap_pPr = ET.SubElement(caption_style, wtag("pPr"))
            ET.SubElement(cap_pPr, wtag("jc")).set(wtag("val"), "center")
            ET.SubElement(cap_pPr, wtag("ind")).set(wtag("firstLine"), "0")
            ET.SubElement(cap_pPr, wtag("spacing")).set(wtag("after"), "120")
            cap_rPr = ET.SubElement(caption_style, wtag("rPr"))
            ET.SubElement(cap_rPr, wtag("sz")).set(wtag("val"), "22")
            ET.SubElement(cap_rPr, wtag("szCs")).set(wtag("val"), "22")

            idx = list(styles_root).index(normal_style)
            styles_root.insert(idx + 1, hr_style)
            styles_root.insert(idx + 2, caption_style)
            styles_tree.write(styles_xml, xml_declaration=True, encoding="UTF-8", pretty_print=False)

        with zipfile.ZipFile(docx_path, "w", zipfile.ZIP_DEFLATED) as zout:
            for rd, _, files in os.walk(workdir):
                for file in files:
                    fp = os.path.join(rd, file)
                    zout.write(fp, os.path.relpath(fp, workdir))