1"""Add comments to DOCX documents.
2 
3Usage:
4    python comment.py unpacked/ 0 "Comment text"
5    python comment.py unpacked/ 1 "Reply text" --parent 0
6 
7Text should be pre-escaped XML (e.g., & for &, ’ for smart quotes).
8 
9After running, add markers to document.xml:
10  <w:commentRangeStart w:id="0"/>
11  ... commented content ...
12  <w:commentRangeEnd w:id="0"/>
13  <w:r><w:rPr><w:rStyle w:val="CommentReference"/></w:rPr><w:commentReference w:id="0"/></w:r>
14"""
15 
16import argparse
17import random
18import shutil
19import sys
20from datetime import datetime, timezone
21from pathlib import Path
22 
23import defusedxml.minidom
24 
25TEMPLATE_DIR = Path(__file__).parent / "templates"
26NS = {
27    "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
28    "w14": "http://schemas.microsoft.com/office/word/2010/wordml",
29    "w15": "http://schemas.microsoft.com/office/word/2012/wordml",
30    "w16cid": "http://schemas.microsoft.com/office/word/2016/wordml/cid",
31    "w16cex": "http://schemas.microsoft.com/office/word/2018/wordml/cex",
32}
33 
34COMMENT_XML = """\
35<w:comment w:id="{id}" w:author="{author}" w:date="{date}" w:initials="{initials}">
36  <w:p w14:paraId="{para_id}" w14:textId="77777777">
37    <w:r>
38      <w:rPr><w:rStyle w:val="CommentReference"/></w:rPr>
39      <w:annotationRef/>
40    </w:r>
41    <w:r>
42      <w:rPr>
43        <w:color w:val="000000"/>
44        <w:sz w:val="20"/>
45        <w:szCs w:val="20"/>
46      </w:rPr>
47      <w:t>{text}</w:t>
48    </w:r>
49  </w:p>
50</w:comment>"""
51 
52COMMENT_MARKER_TEMPLATE = """
53Add to document.xml (markers must be direct children of w:p, never inside w:r):
54  <w:commentRangeStart w:id="{cid}"/>
55  <w:r>...</w:r>
56  <w:commentRangeEnd w:id="{cid}"/>
57  <w:r><w:rPr><w:rStyle w:val="CommentReference"/></w:rPr><w:commentReference w:id="{cid}"/></w:r>"""
58 
59REPLY_MARKER_TEMPLATE = """
60Nest markers inside parent {pid}'s markers (markers must be direct children of w:p, never inside w:r):
61  <w:commentRangeStart w:id="{pid}"/><w:commentRangeStart w:id="{cid}"/>
62  <w:r>...</w:r>
63  <w:commentRangeEnd w:id="{cid}"/><w:commentRangeEnd w:id="{pid}"/>
64  <w:r><w:rPr><w:rStyle w:val="CommentReference"/></w:rPr><w:commentReference w:id="{pid}"/></w:r>
65  <w:r><w:rPr><w:rStyle w:val="CommentReference"/></w:rPr><w:commentReference w:id="{cid}"/></w:r>"""
66 
67 
68def _generate_hex_id() -> str:
69    return f"{random.randint(0, 0x7FFFFFFE):08X}"
70 
71 
72SMART_QUOTE_ENTITIES = {
73    "\u201c": "&#x201C;",  
74    "\u201d": "&#x201D;",  
75    "\u2018": "&#x2018;",  
76    "\u2019": "&#x2019;",  
77}
78 
79 
80def _encode_smart_quotes(text: str) -> str:
81    for char, entity in SMART_QUOTE_ENTITIES.items():
82        text = text.replace(char, entity)
83    return text
84 
85 
86def _append_xml(xml_path: Path, root_tag: str, content: str) -> None:
87    dom = defusedxml.minidom.parseString(xml_path.read_text(encoding="utf-8"))
88    root = dom.getElementsByTagName(root_tag)[0]
89    ns_attrs = " ".join(f'xmlns:{k}="{v}"' for k, v in NS.items())
90    wrapper_dom = defusedxml.minidom.parseString(f"<root {ns_attrs}>{content}</root>")
91    for child in wrapper_dom.documentElement.childNodes:  
92        if child.nodeType == child.ELEMENT_NODE:
93            root.appendChild(dom.importNode(child, True))
94    output = _encode_smart_quotes(dom.toxml(encoding="UTF-8").decode("utf-8"))
95    xml_path.write_text(output, encoding="utf-8")
96 
97 
98def _find_para_id(comments_path: Path, comment_id: int) -> str | None:
99    dom = defusedxml.minidom.parseString(comments_path.read_text(encoding="utf-8"))
100    for c in dom.getElementsByTagName("w:comment"):
101        if c.getAttribute("w:id") == str(comment_id):
102            for p in c.getElementsByTagName("w:p"):
103                if pid := p.getAttribute("w14:paraId"):
104                    return pid
105    return None
106 
107 
108def _get_next_rid(rels_path: Path) -> int:
109    dom = defusedxml.minidom.parseString(rels_path.read_text(encoding="utf-8"))
110    max_rid = 0
111    for rel in dom.getElementsByTagName("Relationship"):
112        rid = rel.getAttribute("Id")
113        if rid and rid.startswith("rId"):
114            try:
115                max_rid = max(max_rid, int(rid[3:]))
116            except ValueError:
117                pass
118    return max_rid + 1
119 
120 
121def _has_relationship(rels_path: Path, target: str) -> bool:
122    dom = defusedxml.minidom.parseString(rels_path.read_text(encoding="utf-8"))
123    for rel in dom.getElementsByTagName("Relationship"):
124        if rel.getAttribute("Target") == target:
125            return True
126    return False
127 
128 
129def _has_content_type(ct_path: Path, part_name: str) -> bool:
130    dom = defusedxml.minidom.parseString(ct_path.read_text(encoding="utf-8"))
131    for override in dom.getElementsByTagName("Override"):
132        if override.getAttribute("PartName") == part_name:
133            return True
134    return False
135 
136 
137def _ensure_comment_relationships(unpacked_dir: Path) -> None:
138    rels_path = unpacked_dir / "word" / "_rels" / "document.xml.rels"
139    if not rels_path.exists():
140        return
141 
142    if _has_relationship(rels_path, "comments.xml"):
143        return  
144 
145    dom = defusedxml.minidom.parseString(rels_path.read_text(encoding="utf-8"))
146    root = dom.documentElement
147    next_rid = _get_next_rid(rels_path)
148 
149    rels = [
150        (
151            "http://schemas.openxmlformats.org/officeDocument/2006/relationships/comments",
152            "comments.xml",
153        ),
154        (
155            "http://schemas.microsoft.com/office/2011/relationships/commentsExtended",
156            "commentsExtended.xml",
157        ),
158        (
159            "http://schemas.microsoft.com/office/2016/09/relationships/commentsIds",
160            "commentsIds.xml",
161        ),
162        (
163            "http://schemas.microsoft.com/office/2018/08/relationships/commentsExtensible",
164            "commentsExtensible.xml",
165        ),
166    ]
167 
168    for rel_type, target in rels:
169        rel = dom.createElement("Relationship")
170        rel.setAttribute("Id", f"rId{next_rid}")
171        rel.setAttribute("Type", rel_type)
172        rel.setAttribute("Target", target)
173        root.appendChild(rel)  
174        next_rid += 1
175 
176    rels_path.write_bytes(dom.toxml(encoding="UTF-8"))
177 
178 
179def _ensure_comment_content_types(unpacked_dir: Path) -> None:
180    ct_path = unpacked_dir / "[Content_Types].xml"
181    if not ct_path.exists():
182        return
183 
184    if _has_content_type(ct_path, "/word/comments.xml"):
185        return  
186 
187    dom = defusedxml.minidom.parseString(ct_path.read_text(encoding="utf-8"))
188    root = dom.documentElement
189 
190    overrides = [
191        (
192            "/word/comments.xml",
193            "application/vnd.openxmlformats-officedocument.wordprocessingml.comments+xml",
194        ),
195        (
196            "/word/commentsExtended.xml",
197            "application/vnd.openxmlformats-officedocument.wordprocessingml.commentsExtended+xml",
198        ),
199        (
200            "/word/commentsIds.xml",
201            "application/vnd.openxmlformats-officedocument.wordprocessingml.commentsIds+xml",
202        ),
203        (
204            "/word/commentsExtensible.xml",
205            "application/vnd.openxmlformats-officedocument.wordprocessingml.commentsExtensible+xml",
206        ),
207    ]
208 
209    for part_name, content_type in overrides:
210        override = dom.createElement("Override")
211        override.setAttribute("PartName", part_name)
212        override.setAttribute("ContentType", content_type)
213        root.appendChild(override)  
214 
215    ct_path.write_bytes(dom.toxml(encoding="UTF-8"))
216 
217 
218def add_comment(
219    unpacked_dir: str,
220    comment_id: int,
221    text: str,
222    author: str = "Claude",
223    initials: str = "C",
224    parent_id: int | None = None,
225) -> tuple[str, str]:
226    word = Path(unpacked_dir) / "word"
227    if not word.exists():
228        return "", f"Error: {word} not found"
229 
230    para_id, durable_id = _generate_hex_id(), _generate_hex_id()
231    ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
232 
233    comments = word / "comments.xml"
234    first_comment = not comments.exists()
235    if first_comment:
236        shutil.copy(TEMPLATE_DIR / "comments.xml", comments)
237        _ensure_comment_relationships(Path(unpacked_dir))
238        _ensure_comment_content_types(Path(unpacked_dir))
239    _append_xml(
240        comments,
241        "w:comments",
242        COMMENT_XML.format(
243            id=comment_id,
244            author=author,
245            date=ts,
246            initials=initials,
247            para_id=para_id,
248            text=text,  
249        ),
250    )
251 
252    ext = word / "commentsExtended.xml"
253    if not ext.exists():
254        shutil.copy(TEMPLATE_DIR / "commentsExtended.xml", ext)
255    if parent_id is not None:
256        parent_para = _find_para_id(comments, parent_id)
257        if not parent_para:
258            return "", f"Error: Parent comment {parent_id} not found"
259        _append_xml(
260            ext,
261            "w15:commentsEx",
262            f'<w15:commentEx w15:paraId="{para_id}" w15:paraIdParent="{parent_para}" w15:done="0"/>',
263        )
264    else:
265        _append_xml(
266            ext,
267            "w15:commentsEx",
268            f'<w15:commentEx w15:paraId="{para_id}" w15:done="0"/>',
269        )
270 
271    ids = word / "commentsIds.xml"
272    if not ids.exists():
273        shutil.copy(TEMPLATE_DIR / "commentsIds.xml", ids)
274    _append_xml(
275        ids,
276        "w16cid:commentsIds",
277        f'<w16cid:commentId w16cid:paraId="{para_id}" w16cid:durableId="{durable_id}"/>',
278    )
279 
280    extensible = word / "commentsExtensible.xml"
281    if not extensible.exists():
282        shutil.copy(TEMPLATE_DIR / "commentsExtensible.xml", extensible)
283    _append_xml(
284        extensible,
285        "w16cex:commentsExtensible",
286        f'<w16cex:commentExtensible w16cex:durableId="{durable_id}" w16cex:dateUtc="{ts}"/>',
287    )
288 
289    action = "reply" if parent_id is not None else "comment"
290    return para_id, f"Added {action} {comment_id} (para_id={para_id})"
291 
292 
293if __name__ == "__main__":
294    p = argparse.ArgumentParser(description="Add comments to DOCX documents")
295    p.add_argument("unpacked_dir", help="Unpacked DOCX directory")
296    p.add_argument("comment_id", type=int, help="Comment ID (must be unique)")
297    p.add_argument("text", help="Comment text")
298    p.add_argument("--author", default="Claude", help="Author name")
299    p.add_argument("--initials", default="C", help="Author initials")
300    p.add_argument("--parent", type=int, help="Parent comment ID (for replies)")
301    args = p.parse_args()
302 
303    para_id, msg = add_comment(
304        args.unpacked_dir,
305        args.comment_id,
306        args.text,
307        args.author,
308        args.initials,
309        args.parent,
310    )
311    print(msg)
312    if "Error" in msg:
313        sys.exit(1)
314    cid = args.comment_id
315    if args.parent is not None:
316        print(REPLY_MARKER_TEMPLATE.format(pid=args.parent, cid=cid))
317    else:
318        print(COMMENT_MARKER_TEMPLATE.format(cid=cid))
319