Fixes

Author: Xpitfire

singularity/DocumentExport/DocumentDocxWriter.swift

@@ -7,6 +7,8 @@ enum DocumentDocxWriterError: Error {
 }
 
 struct DocumentDocxWriter {
+    private let wordNamespace = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
+
     func write(envelope: DocumentContentEnvelope, templateURL: URL, outputURL: URL, headingStyle: String, bodyStyle: String, tableStyle: String) throws {
         let fm = FileManager.default
         guard fm.fileExists(atPath: templateURL.path) else { throw DocumentDocxWriterError.templateMissing }
@@ -15,31 +17,93 @@ struct DocumentDocxWriter {
         try fm.createDirectory(at: unzipDir, withIntermediateDirectories: true)
         try unzip(docx: templateURL, to: unzipDir)
         let documentPath = unzipDir.appendingPathComponent("word/document.xml")
-        guard let data = try? Data(contentsOf: documentPath), var documentXML = String(data: data, encoding: .utf8) else {
+        guard let data = try? Data(contentsOf: documentPath), let documentXML = String(data: data, encoding: .utf8) else {
             throw DocumentDocxWriterError.invalidDocument
         }
+        let updatedXML = processDocumentXML(documentXML, envelope: envelope, headingStyle: headingStyle, bodyStyle: bodyStyle, tableStyle: tableStyle)
+        try updatedXML.data(using: .utf8)?.write(to: documentPath)
+        try zipFolder(at: unzipDir, to: workRoot.appendingPathComponent("output.docx"))
+        let finalURL = outputURL
+        let parent = finalURL.deletingLastPathComponent()
+        try fm.createDirectory(at: parent, withIntermediateDirectories: true)
+        if fm.fileExists(atPath: finalURL.path) {
+            try fm.removeItem(at: finalURL)
+        }
+        try fm.moveItem(at: workRoot.appendingPathComponent("output.docx"), to: finalURL)
+        try? fm.removeItem(at: workRoot)
+    }
+
+    private func processDocumentXML(_ xmlString: String, envelope: DocumentContentEnvelope, headingStyle: String, bodyStyle: String, tableStyle: String) -> String {
+        guard let document = try? XMLDocument(xmlString: xmlString, options: [.nodePreserveAll]) else {
+            return legacyProcessDocumentXML(xmlString, envelope: envelope, headingStyle: headingStyle, bodyStyle: bodyStyle, tableStyle: tableStyle)
+        }
+        var unmatched: [(replacement: String, key: String)] = []
+        for section in envelope.sections {
+            let replacement = renderSection(section, headingStyle: headingStyle, bodyStyle: bodyStyle, tableStyle: tableStyle)
+            if !replacePlaceholder(in: document, key: section.definition.key, replacement: replacement) {
+                unmatched.append((replacement, section.definition.key))
+            }
+        }
+        if !unmatched.isEmpty,
+           let body = (try? document.nodes(forXPath: "//*[local-name()='body']"))?.first as? XMLElement {
+            for entry in unmatched {
+                if let nodes = try? nodes(from: entry.replacement) {
+                    for node in nodes {
+                        body.addChild(node)
+                    }
+                }
+                removePlaceholderText(key: entry.key, in: document)
+            }
+        }
+        return document.xmlString(options: [.nodePreserveAll])
+    }
+
+    private func legacyProcessDocumentXML(_ xmlString: String, envelope: DocumentContentEnvelope, headingStyle: String, bodyStyle: String, tableStyle: String) -> String {
+        var documentXML = xmlString
         var appendQueue: [String] = []
         for section in envelope.sections {
-            let replacement = xml(for: section, headingStyle: headingStyle, bodyStyle: bodyStyle, tableStyle: tableStyle)
+            let replacement = renderSection(section, headingStyle: headingStyle, bodyStyle: bodyStyle, tableStyle: tableStyle)
             if !replacePlaceholder(key: section.definition.key, in: &documentXML, replacement: replacement) {
                 removePlaceholderOccurrences(key: section.definition.key, in: &documentXML)
                 appendQueue.append(replacement)
             }
         }
-        if !appendQueue.isEmpty {
-            guard let range = documentXML.range(of: "</w:body>") else { throw DocumentDocxWriterError.invalidDocument }
+        if !appendQueue.isEmpty, let range = documentXML.range(of: "</w:body>") {
             documentXML.replaceSubrange(range, with: appendQueue.joined() + "</w:body>")
         }
-        try documentXML.data(using: .utf8)?.write(to: documentPath)
-        try zipFolder(at: unzipDir, to: workRoot.appendingPathComponent("output.docx"))
-        let finalURL = outputURL
-        let parent = finalURL.deletingLastPathComponent()
-        try fm.createDirectory(at: parent, withIntermediateDirectories: true)
-        if fm.fileExists(atPath: finalURL.path) {
-            try fm.removeItem(at: finalURL)
+        return documentXML
+    }
+
+    private func replacePlaceholder(in document: XMLDocument, key: String, replacement: String) -> Bool {
+        let xpath = "//*[local-name()='p' and contains(., '{{\(key)}}')]"
+        guard let matches = try? document.nodes(forXPath: xpath), !matches.isEmpty else { return false }
+        for case let element as XMLElement in matches {
+            guard let parent = element.parent as? XMLElement else { continue }
+            let insertionIndex = parent.children?.firstIndex(of: element) ?? parent.childCount
+            parent.removeChild(at: insertionIndex)
+            if let nodes = try? nodes(from: replacement) {
+                for (offset, node) in nodes.enumerated() {
+                    parent.insertChild(node, at: insertionIndex + offset)
+                }
+            }
+        }
+        return true
+    }
+
+    private func nodes(from fragment: String) throws -> [XMLNode] {
+        let wrapped = "<root xmlns:w=\"\(wordNamespace)\">\(fragment)</root>"
+        let doc = try XMLDocument(xmlString: wrapped, options: [.nodePreserveAll])
+        guard let children = doc.rootElement()?.children else { return [] }
+        return children.compactMap { $0.copy() as? XMLNode }
+    }
+
+    private func removePlaceholderText(key: String, in document: XMLDocument) {
+        guard let textNodes = try? document.nodes(forXPath: "//*[local-name()='t']") else { return }
+        let token = "{{\(key)}}"
+        for case let node as XMLNode in textNodes {
+            guard let text = node.stringValue, text.contains(token) else { continue }
+            node.stringValue = text.replacingOccurrences(of: token, with: "")
         }
-        try fm.moveItem(at: workRoot.appendingPathComponent("output.docx"), to: finalURL)
-        try? fm.removeItem(at: workRoot)
     }
 
     private func headingStyleFor(section: DocumentSection, defaultHeadingStyle: String, bodyStyle: String) -> String {
@@ -86,7 +150,7 @@ struct DocumentDocxWriter {
     private func paragraphXML(text: String, style: String) -> String {
         let escaped = escape(text)
         return """
-<w:p xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
+<w:p xmlns:w="\(wordNamespace)">
   <w:pPr>
     <w:pStyle w:val="\(style)"/>
   </w:pPr>
@@ -108,7 +172,7 @@ struct DocumentDocxWriter {
             return "<w:tr>\(cells)</w:tr>"
         }.joined()
         return """
-<w:tbl xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
+<w:tbl xmlns:w="\(wordNamespace)">
   <w:tblPr>
     <w:tblStyle w:val="\(style)"/>
     <w:tblLook w:val="04A0" w:firstRow="1" w:lastRow="0" w:firstColumn="1" w:lastColumn="0" w:noHBand="0" w:noVBand="1"/>
@@ -126,7 +190,7 @@ struct DocumentDocxWriter {
         let shading = isHeader ? "<w:shd w:val=\"clear\" w:color=\"auto\" w:fill=\"D9D9D9\"/>" : ""
         let bold = isHeader ? "<w:b/>" : ""
         return """
-    <w:tc>
+    <w:tc xmlns:w="\(wordNamespace)">
       <w:tcPr>
         <w:tcW w:w="5000" w:type="dxa"/>
         \(shading)
@@ -154,7 +218,7 @@ struct DocumentDocxWriter {
         return escaped
     }
 
-    private func xml(for section: DocumentSection, headingStyle: String, bodyStyle: String, tableStyle: String) -> String {
+    private func renderSection(_ section: DocumentSection, headingStyle: String, bodyStyle: String, tableStyle: String) -> String {
         switch section.value {
         case .text(let value):
             return paragraphsXML(text: value, style: section.definition.style ?? headingStyleFor(section: section, defaultHeadingStyle: headingStyle, bodyStyle: bodyStyle))