#!/usr/bin/env python3
"""Create a one-section blank working copy from the retained Kangbite DOCX."""

from __future__ import annotations

import argparse
import hashlib
import shutil
import tempfile
import zipfile
from pathlib import Path
from xml.etree import ElementTree as ET


EXPECTED_SHA256 = "3c3a53327faccaec22b753a8a2a6d55f399f83aff2c9d4791b7c8fceefbd7f13"
W = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
R = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"

ET.register_namespace("w", W)
ET.register_namespace("r", R)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def build(source: Path, output: Path) -> None:
    actual = sha256(source)
    if actual != EXPECTED_SHA256:
        raise SystemExit(
            f"Template hash mismatch: expected {EXPECTED_SHA256}, observed {actual}. "
            "Redistill the template before continuing."
        )
    if source.resolve() == output.resolve():
        raise SystemExit("Output must differ from the retained source template.")

    with tempfile.TemporaryDirectory(prefix="kangbite-word-template-") as tmp:
        unpacked = Path(tmp) / "package"
        with zipfile.ZipFile(source) as archive:
            archive.extractall(unpacked)

        document_path = unpacked / "word" / "document.xml"
        tree = ET.parse(document_path)
        root = tree.getroot()
        body = root.find(f"{{{W}}}body")
        if body is None:
            raise SystemExit("Invalid DOCX: word/document.xml has no w:body.")

        final_sect_pr = body.find(f"{{{W}}}sectPr")
        if final_sect_pr is None:
            raise SystemExit("Invalid template: final w:sectPr is missing.")

        for child in list(body):
            if child is not final_sect_pr:
                body.remove(child)

        for ref in list(final_sect_pr):
            if ref.tag in {f"{{{W}}}headerReference", f"{{{W}}}footerReference"}:
                final_sect_pr.remove(ref)

        header_ref = ET.Element(f"{{{W}}}headerReference")
        header_ref.set(f"{{{W}}}type", "default")
        header_ref.set(f"{{{R}}}id", "rId9")
        final_sect_pr.insert(0, header_ref)

        footer_ref = ET.Element(f"{{{W}}}footerReference")
        footer_ref.set(f"{{{W}}}type", "default")
        footer_ref.set(f"{{{R}}}id", "rId10")
        final_sect_pr.insert(1, footer_ref)

        paragraph = ET.Element(f"{{{W}}}p")
        body.insert(0, paragraph)
        tree.write(document_path, encoding="UTF-8", xml_declaration=True)

        output.parent.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive:
            for path in sorted(unpacked.rglob("*")):
                if path.is_file():
                    archive.write(path, path.relative_to(unpacked).as_posix())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    build(args.source, args.output)
    print(f"Prepared working copy: {args.output}")


if __name__ == "__main__":
    main()