enkly-ai/pkg/officekit/docx.go

package officekit

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// ParaStyle is the small, Claude-Code-skill-inspired vocabulary of
// paragraph styles the model is offered — deliberately not "any OOXML
// style name", so the model can never need to type XML to format
// something.
type ParaStyle string

const (
	StyleNormal   ParaStyle = "normal"
	StyleHeading1 ParaStyle = "heading1"
	StyleHeading2 ParaStyle = "heading2"
	StyleHeading3 ParaStyle = "heading3"
	StyleBullet   ParaStyle = "bullet"
)

// Paragraph is one structured paragraph of a Word document, as the
// model supplies it via the create_document/edit_document tool
// arguments — plain fields, never markup.
type Paragraph struct {
	Text      string
	Style     ParaStyle
	Bold      bool
	Italic    bool
	Underline bool
}

const docxDocumentPart = "word/document.xml"

// --- word/document.xml, built with encoding/xml so every run of text
// is correctly escaped no matter what characters it contains (a stray
// "<" or "&" in a paragraph's text becomes "&lt;"/"&amp;" the same way
// Go's XML encoder would escape it anywhere else — it can never end up
// as literal, parseable-looking markup). ---

// wDocument's XMLName tag uses Go's "namespace-URI local-name" form
// (space-separated, not a colon prefix — see the package doc comment
// in zipkit.go's header and this file's note above BuildDocx for why
// colon-prefixed tags like "w:document" don't round-trip). Marshal
// emits this as the default namespace (xmlns="...") on the root
// element; every child struct below uses a plain, unprefixed local
// name with no namespace component, which both (a) is valid XML —
// unprefixed descendants simply inherit the nearest enclosing default
// namespace — and (b) unmarshals back correctly, because an
// unnamespaced field tag matches an element by local name only,
// regardless of which namespace it actually resolved to.
type wDocument struct {
	XMLName xml.Name `xml:"http://schemas.openxmlformats.org/wordprocessingml/2006/main document"`
	Body    wBody    `xml:"body"`
}

type wBody struct {
	Paragraphs []wParagraph `xml:"p"`
	SectPr     wSectPr      `xml:"sectPr"`
}

type wSectPr struct {
	PgSz wPgSz `xml:"pgSz"`
}

type wPgSz struct {
	W int `xml:"w,attr"`
	H int `xml:"h,attr"`
}

type wParagraph struct {
	PPr *wPPr `xml:"pPr,omitempty"`
	Run wRun  `xml:"r"`
}

type wPPr struct {
	NumPr *wNumPr `xml:"numPr,omitempty"`
}

type wNumPr struct {
	Ilvl  wVal `xml:"ilvl"`
	NumId wVal `xml:"numId"`
}

type wVal struct {
	Val int `xml:"val,attr"`
}

type wRun struct {
	RPr  *wRPr `xml:"rPr,omitempty"`
	Text wText `xml:"t"`
}

// wRPr's field order follows OOXML's CT_RPr schema sequence (b, i,
// sz, u, ...) — Go's xml.Marshal emits struct fields in declaration
// order, and while real-world readers (Word included) tend to be
// lenient about run-property order, writing it correctly costs
// nothing and keeps generated files schema-valid, not just
// Word-openable.
type wRPr struct {
	Bold      *struct{} `xml:"b,omitempty"`
	Italic    *struct{} `xml:"i,omitempty"`
	Sz        *wSz      `xml:"sz,omitempty"`
	Underline *wU       `xml:"u,omitempty"`
}

type wSz struct {
	Val int `xml:"val,attr"`
}

type wU struct {
	Val string `xml:"val,attr"`
}

type wText struct {
	Space string `xml:"xml:space,attr"`
	Value string `xml:",chardata"`
}

// wParagraphFragment is the same content as wParagraph, but with its
// own explicit WordprocessingML namespace declaration on XMLName.
// It exists only for AppendDocxParagraph, which marshals ONE paragraph
// standalone and splices its bytes directly into an existing
// document.xml rather than re-marshaling the whole document (see that
// function's doc comment for why). A spliced fragment has no parent
// element to inherit a namespace from if the target document uses
// Word's own "w:"-prefixed convention instead of this package's
// default-namespace one — so this type redeclares the namespace
// locally on the fragment's own root element. XML lets any element do
// that regardless of its surrounding document's prefix convention, and
// every OOXML-consuming application (Word included) resolves element
// identity by namespace URI + local name, not by literal prefix text,
// so `<p xmlns="...wordprocessingml.../main">` and `<w:p>` are the same
// element as far as Word is concerned. wParagraph itself deliberately
// keeps no such tag, so BuildDocx's normal output — where every
// paragraph already inherits the root's default namespace — doesn't
// pay for a redundant xmlns on every single paragraph.
type wParagraphFragment struct {
	XMLName xml.Name `xml:"http://schemas.openxmlformats.org/wordprocessingml/2006/main p"`
	PPr     *wPPr    `xml:"pPr,omitempty"`
	Run     wRun     `xml:"r"`
}

func newParagraphFragment(p Paragraph) wParagraphFragment {
	wp := toWParagraph(p)
	return wParagraphFragment{PPr: wp.PPr, Run: wp.Run}
}

// runProps maps the model's small formatting vocabulary (bold, italic,
// underline, plus the size/boldness implied by a heading style) to a
// run's rPr, or nil if the run needs no rPr at all (the common case:
// a plain, unstyled normal/bullet paragraph).
func runProps(p Paragraph) *wRPr {
	bold, italic, underline := p.Bold, p.Italic, p.Underline
	var sz *wSz
	switch p.Style {
	case StyleHeading1:
		bold = true
		sz = &wSz{Val: 32} // 16pt
	case StyleHeading2:
		bold = true
		sz = &wSz{Val: 28} // 14pt
	case StyleHeading3:
		bold = true
		sz = &wSz{Val: 24} // 12pt
	}
	if !bold && !italic && !underline && sz == nil {
		return nil
	}
	rpr := &wRPr{Sz: sz}
	if bold {
		rpr.Bold = &struct{}{}
	}
	if italic {
		rpr.Italic = &struct{}{}
	}
	if underline {
		rpr.Underline = &wU{Val: "single"}
	}
	return rpr
}

func toWParagraph(p Paragraph) wParagraph {
	wp := wParagraph{
		Run: wRun{
			RPr:  runProps(p),
			Text: wText{Space: "preserve", Value: p.Text},
		},
	}
	if p.Style == StyleBullet {
		wp.PPr = &wPPr{NumPr: &wNumPr{Ilvl: wVal{Val: 0}, NumId: wVal{Val: 1}}}
	}
	return wp
}

// BuildDocx renders paragraphs into a complete, valid .docx file.
func BuildDocx(paragraphs []Paragraph) ([]byte, error) {
	body := wBody{SectPr: wSectPr{PgSz: wPgSz{W: 12240, H: 15840}}}
	for _, p := range paragraphs {
		body.Paragraphs = append(body.Paragraphs, toWParagraph(p))
	}
	doc := wDocument{Body: body}
	docXML, err := xml.Marshal(doc)
	if err != nil {
		return nil, fmt.Errorf("officekit: marshal document.xml: %w", err)
	}

	// A minimal numbering part is included (word/numbering.xml) only so
	// StyleBullet's w:numId="1" resolves to something rather than being
	// a dangling reference; every bullet paragraph shares the same
	// single bullet-list definition.
	parts := []zipPart{
		{Name: "[Content_Types].xml", Body: docxContentTypes},
		{Name: "_rels/.rels", Body: docxRootRels},
		{Name: "word/_rels/document.xml.rels", Body: docxDocumentRels},
		{Name: docxDocumentPart, Body: xml.Header + string(docXML)},
		{Name: "word/numbering.xml", Body: docxNumberingXML},
		{Name: "docProps/core.xml", Body: officeCoreProps},
		{Name: "docProps/app.xml", Body: docxAppProps},
	}
	return buildZip(parts)
}

// ReadDocxText extracts the plain text of a .docx file, one line per
// paragraph (bullets prefixed with "- "), for the read_document tool.
// This is what the model actually sees — never the underlying XML.
//
// This walks document.xml's raw token stream rather than unmarshaling
// into a struct that only models direct <w:body> children: a real
// document's paragraphs can be nested inside a table cell
// (<w:tbl>/<w:tr>/<w:tc>/<w:p>), and a struct-based Unmarshal into
// wBody.Paragraphs would silently skip every one of those, making a
// user-uploaded document with tables look emptier than it is. The
// token walk picks up every <w:p> wherever it occurs, in document
// order; it still doesn't reconstruct a table's row/column grid (see
// the docx skill's documented limitation) — each cell's paragraphs
// simply appear as their own lines.
func ReadDocxText(data []byte) (string, error) {
	raw, ok, err := readZipPart(data, docxDocumentPart)
	if err != nil {
		return "", err
	}
	if !ok {
		return "", fmt.Errorf("officekit: %s not found (not a valid .docx)", docxDocumentPart)
	}
	lines, err := walkParagraphText(raw)
	if err != nil {
		return "", fmt.Errorf("officekit: parse document.xml: %w", err)
	}
	return strings.Join(lines, "\n"), nil
}

// walkParagraphText returns one text line per <w:p> paragraph found
// anywhere in document.xml, in document order: the concatenation of
// every <w:t> run's text within it, with <w:tab/> rendered as a
// literal tab, prefixed "- " when the paragraph carries a <w:numPr>
// (list) property. Matches by local element name only (see the
// package-level namespace note on wDocument), so it works the same
// whether the source uses Word's "w:" prefix or this package's own
// default-namespace convention.
func walkParagraphText(raw []byte) ([]string, error) {
	dec := xml.NewDecoder(bytes.NewReader(raw))
	var (
		lines    []string
		cur      strings.Builder
		inPara   bool
		inText   bool
		isBullet bool
	)
	flushPara := func() {
		text := cur.String()
		if isBullet {
			text = "- " + text
		}
		lines = append(lines, text)
		cur.Reset()
		inPara, isBullet = false, false
	}
	for {
		tok, err := dec.Token()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		switch t := tok.(type) {
		case xml.StartElement:
			switch t.Name.Local {
			case "p":
				inPara = true
			case "t":
				inText = true
			case "tab":
				if inPara {
					cur.WriteByte('\t')
				}
			case "numPr":
				if inPara {
					isBullet = true
				}
			}
		case xml.EndElement:
			switch t.Name.Local {
			case "p":
				if inPara {
					flushPara()
				}
			case "t":
				inText = false
			}
		case xml.CharData:
			if inPara && inText {
				cur.Write(t)
			}
		}
	}
	return lines, nil
}

// AppendDocxParagraph adds one paragraph as the new last content
// element of the document body and returns the new file bytes.
//
// This splices the new paragraph's marshaled XML directly into
// document.xml's existing raw bytes, immediately before the body's
// trailing <w:sectPr> (OOXML requires section properties to be the
// body's last child) — or immediately before </w:body> if there's no
// sectPr at all — instead of the unmarshal-modify-marshal round trip
// this used to share with ReplaceDocxText. That round trip only ever
// modeled direct <w:body> children (paragraphs and the trailing
// sectPr), so re-marshaling it back would silently drop everything
// else a real document has: tables, images, headers, multiple runs
// per paragraph, tracked changes. Splicing means every byte of the
// original document this package doesn't model survives untouched.
func AppendDocxParagraph(data []byte, p Paragraph) ([]byte, error) {
	raw, ok, err := readZipPart(data, docxDocumentPart)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("officekit: %s not found (not a valid .docx)", docxDocumentPart)
	}
	paraXML, err := xml.Marshal(newParagraphFragment(p))
	if err != nil {
		return nil, fmt.Errorf("officekit: marshal paragraph: %w", err)
	}
	at := bodyInsertionPoint(raw)
	if at < 0 {
		return nil, fmt.Errorf("officekit: %s has no <w:body> to append to (not a valid .docx)", docxDocumentPart)
	}
	newDoc := make([]byte, 0, len(raw)+len(paraXML))
	newDoc = append(newDoc, raw[:at]...)
	newDoc = append(newDoc, paraXML...)
	newDoc = append(newDoc, raw[at:]...)
	return replaceZipPart(data, docxDocumentPart, newDoc)
}

// bodyInsertionPoint returns the byte offset in document.xml where new
// body content should be spliced in: right before the body's trailing
// <w:sectPr> if present, else right before </w:body>. Returns -1 if
// neither tag is found at all (not a recognizable Word document body).
func bodyInsertionPoint(doc []byte) int {
	for _, tag := range []string{"<w:sectPr", "<sectPr"} {
		if i := lastTagIndex(doc, tag); i >= 0 {
			return i
		}
	}
	for _, tag := range []string{"</w:body>", "</body>"} {
		if i := bytes.Index(doc, []byte(tag)); i >= 0 {
			return i
		}
	}
	return -1
}

// lastTagIndex finds the last byte offset of an opening tag whose name
// is exactly tag (e.g. "<w:sectPr"), rejecting a longer tag name that
// merely starts with the same bytes — such as the tracked-changes
// element "<w:sectPrChange", which would otherwise look like a match.
func lastTagIndex(doc []byte, tag string) int {
	needle := []byte(tag)
	end := len(doc)
	for {
		i := bytes.LastIndex(doc[:end], needle)
		if i < 0 {
			return -1
		}
		after := i + len(needle)
		if after >= len(doc) {
			return i
		}
		switch doc[after] {
		case ' ', '\t', '\n', '\r', '>', '/':
			return i
		}
		end = i // false positive (e.g. sectPrChange); keep searching earlier
	}
}

// wTextElementRe matches one WordprocessingML text run's element —
// <w:t>...</w:t>, the unprefixed <t>...</t> this package's own output
// uses, or any other single-letter-prefixed variant — capturing the
// opening tag, the raw (still XML-escaped) inner text, and the closing
// tag separately, so ReplaceDocxText only ever rewrites the inner text
// and leaves the tag itself, and everything outside it, byte-for-byte
// untouched. The boundary check after the local name (next character
// must be whitespace, "/", or ">") excludes longer names that merely
// start with "t", such as "tab", "tbl", or "tc".
var wTextElementRe = regexp.MustCompile(`(?s)(<(?:[A-Za-z][\w.-]*:)?t(?:\s[^>]*)?>)(.*?)(</(?:[A-Za-z][\w.-]*:)?t>)`)

// ReplaceDocxText replaces every exact, literal occurrence of find
// with replace inside the document's text runs — including runs
// nested inside a table cell, which the old paragraph-struct-based
// approach never saw at all — and returns the new file bytes plus how
// many text runs were changed. It rewrites only the matched <w:t>
// elements' inner text in document.xml's raw bytes; every other part
// of the document (formatting, tables, images, headers) passes through
// completely untouched.
//
// Because Word often splits a single sentence across more than one run
// (spell-check and revision tracking both do this), a phrase that
// spans a run boundary may not be found even though it reads as one
// phrase in Word — read_document first, and prefer a shorter, more
// literal match string when a longer one doesn't seem to apply.
func ReplaceDocxText(data []byte, find, replace string) ([]byte, int, error) {
	raw, ok, err := readZipPart(data, docxDocumentPart)
	if err != nil {
		return nil, 0, err
	}
	if !ok {
		return nil, 0, fmt.Errorf("officekit: %s not found (not a valid .docx)", docxDocumentPart)
	}

	changed := 0
	newDoc := wTextElementRe.ReplaceAllFunc(raw, func(match []byte) []byte {
		groups := wTextElementRe.FindSubmatch(match)
		open, inner, closeTag := groups[1], groups[2], groups[3]
		decoded := xmlUnescapeText(inner)
		if !strings.Contains(decoded, find) {
			return match
		}
		changed++
		replaced := strings.ReplaceAll(decoded, find, replace)
		out := make([]byte, 0, len(open)+len(replaced)+len(closeTag))
		out = append(out, open...)
		out = append(out, []byte(xmlEscapeText(replaced))...)
		out = append(out, closeTag...)
		return out
	})
	if changed == 0 {
		return nil, 0, fmt.Errorf("officekit: %q not found in any text run", find)
	}
	out, err := replaceZipPart(data, docxDocumentPart, newDoc)
	return out, changed, err
}

// xmlUnescapeText reverses XML's five predefined entities plus numeric
// character references (&#NNN; / &#xHH;) — the inverse of
// xmlEscapeText (pptx.go) — so ReplaceDocxText can match and replace
// against the plain text the model actually passed, not its
// XML-escaped form. An unknown or malformed entity is left verbatim
// rather than erroring: this only ever feeds a literal-substring
// comparison, so there's nothing to fail loudly about.
func xmlUnescapeText(b []byte) string {
	s := string(b)
	if !strings.ContainsRune(s, '&') {
		return s
	}
	var out strings.Builder
	for {
		i := strings.IndexByte(s, '&')
		if i < 0 {
			out.WriteString(s)
			break
		}
		out.WriteString(s[:i])
		rest := s[i+1:]
		semi := strings.IndexByte(rest, ';')
		if semi < 0 {
			out.WriteByte('&')
			s = rest
			continue
		}
		entity, tail := rest[:semi], rest[semi+1:]
		switch {
		case entity == "amp":
			out.WriteByte('&')
		case entity == "lt":
			out.WriteByte('<')
		case entity == "gt":
			out.WriteByte('>')
		case entity == "apos":
			out.WriteByte('\'')
		case entity == "quot":
			out.WriteByte('"')
		case strings.HasPrefix(entity, "#x"), strings.HasPrefix(entity, "#X"):
			if n, err := strconv.ParseInt(entity[2:], 16, 32); err == nil {
				out.WriteRune(rune(n))
			} else {
				out.WriteByte('&')
				out.WriteString(entity)
				out.WriteByte(';')
			}
		case strings.HasPrefix(entity, "#"):
			if n, err := strconv.ParseInt(entity[1:], 10, 32); err == nil {
				out.WriteRune(rune(n))
			} else {
				out.WriteByte('&')
				out.WriteString(entity)
				out.WriteByte(';')
			}
		default:
			out.WriteByte('&')
			out.WriteString(entity)
			out.WriteByte(';')
		}
		s = tail
	}
	return out.String()
}

const docxContentTypes = xml.Header + `<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
<Override PartName="/word/numbering.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.numbering+xml"/>
<Override PartName="/docProps/core.xml" ContentType="application/vnd.openxmlformats-package.core-properties+xml"/>
<Override PartName="/docProps/app.xml" ContentType="application/vnd.openxmlformats-officedocument.extended-properties+xml"/>
</Types>`

const docxRootRels = xml.Header + `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
<Relationship Id="rId2" Type="http://schemas.openxmlformats.org/package/2006/relationships/metadata/core-properties" Target="docProps/core.xml"/>
<Relationship Id="rId3" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/extended-properties" Target="docProps/app.xml"/>
</Relationships>`

const docxDocumentRels = xml.Header + `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/numbering" Target="numbering.xml"/>
</Relationships>`

const docxNumberingXML = xml.Header + `<w:numbering xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
<w:abstractNum w:abstractNumId="0">
<w:lvl w:ilvl="0"><w:start w:val="1"/><w:numFmt w:val="bullet"/><w:lvlText w:val=""/><w:pPr><w:ind w:left="720" w:hanging="360"/></w:pPr></w:lvl>
</w:abstractNum>
<w:num w:numId="1"><w:abstractNumId w:val="0"/></w:num>
</w:numbering>`

const docxAppProps = xml.Header + `<Properties xmlns="http://schemas.openxmlformats.org/officeDocument/2006/extended-properties"><Application>enkly</Application></Properties>`

// officeCoreProps is shared verbatim by docx/xlsx/pptx — the schema is
// package-level (docProps/core.xml), not format-specific.
const officeCoreProps = xml.Header + `<cp:coreProperties xmlns:cp="http://schemas.openxmlformats.org/package/2006/metadata/core-properties" xmlns:dc="http://purl.org/dc/elements/1.1/"><dc:creator>enkly agent</dc:creator></cp:coreProperties>`