from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile
import shutil
import tempfile
import xml.etree.ElementTree as ET


SOURCE_DOC = Path(r"C:\Users\ASUS\Downloads\Penjelasan_Rinci_KMeans_Rekomendasi_JMW.docx")
APPEND_DOC = Path("docs/penjelasan-lengkap-kmeans-rekomendasi-budget-update.docx")
OUT_DOC = Path("docs/Penjelasan_Rinci_KMeans_Rekomendasi_JMW_Gabungan.docx")

WORD_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
ET.register_namespace("w", WORD_NS)


def body_children(docx_path):
    with ZipFile(docx_path, "r") as archive:
        xml = archive.read("word/document.xml")

    root = ET.fromstring(xml)
    body = root.find(f"{{{WORD_NS}}}body")

    children = list(body)
    sect_pr = None
    content = []

    for child in children:
        if child.tag == f"{{{WORD_NS}}}sectPr":
            sect_pr = child
        else:
            content.append(child)

    return root, body, content, sect_pr


def page_break_paragraph():
    paragraph = ET.Element(f"{{{WORD_NS}}}p")
    run = ET.SubElement(paragraph, f"{{{WORD_NS}}}r")
    br = ET.SubElement(run, f"{{{WORD_NS}}}br")
    br.set(f"{{{WORD_NS}}}type", "page")
    return paragraph


def merge_docs():
    OUT_DOC.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as tmp:
        tmp_dir = Path(tmp)

        with ZipFile(SOURCE_DOC, "r") as archive:
            archive.extractall(tmp_dir)

        source_root, source_body, _, source_sect = body_children(SOURCE_DOC)
        _, _, append_content, _ = body_children(APPEND_DOC)

        for child in list(source_body):
            source_body.remove(child)

        _, _, source_content, _ = body_children(SOURCE_DOC)
        for child in source_content:
            source_body.append(child)

        source_body.append(page_break_paragraph())

        for child in append_content:
            source_body.append(child)

        if source_sect is not None:
            source_body.append(source_sect)

        document_path = tmp_dir / "word" / "document.xml"
        ET.ElementTree(source_root).write(document_path, encoding="utf-8", xml_declaration=True)

        if OUT_DOC.exists():
            OUT_DOC.unlink()

        with ZipFile(OUT_DOC, "w", ZIP_DEFLATED) as archive:
            for path in tmp_dir.rglob("*"):
                if path.is_file():
                    archive.write(path, path.relative_to(tmp_dir).as_posix())

    print(OUT_DOC.resolve())


if __name__ == "__main__":
    merge_docs()