enkly-ai/pkg/officekit/xlsx.go

package officekit

import (
	"encoding/xml"
	"fmt"
	"sort"
	"strconv"
	"strings"
)

// Sheet is one worksheet's data, as the model supplies it to
// create_spreadsheet: a name and a grid of string cell values (numbers
// are written as plain decimal strings; anything that parses as a
// number is stored as a numeric cell so Excel treats it as one).
type Sheet struct {
	Name string
	Rows [][]string
}

const xlsxWorkbookPart = "xl/workbook.xml"
const xlsxWorkbookRelsPart = "xl/_rels/workbook.xml.rels"

// xlsxSheetPart returns the conventional worksheet part path for the
// Nth (1-based) sheet BuildXlsx writes — sheet1.xml, sheet2.xml, ...
func xlsxSheetPart(n int) string { return fmt.Sprintf("xl/worksheets/sheet%d.xml", n) }

// --- typed XML for the parts BuildXlsx/edit operations actually
// generate from user-supplied content (worksheet XML, plus the small
// workbook/rels indirection needed to support more than one sheet);
// every other part is static boilerplate with no user content. ---

type sWorksheet struct {
	XMLName   xml.Name   `xml:"worksheet"`
	Xmlns     string     `xml:"xmlns,attr"`
	SheetData sSheetData `xml:"sheetData"`
}

type sSheetData struct {
	Rows []sRow `xml:"row"`
}

type sRow struct {
	R     int     `xml:"r,attr"`
	Cells []sCell `xml:"c"`
}

// sCell's field order (after the r/t attrs) follows OOXML's CT_Cell
// schema sequence — f (formula), v (value), is (inline string) — so a
// formula cell with a cached value (f + v both set, as real Excel
// always writes them) round-trips as a schema-valid element.
type sCell struct {
	R       string  `xml:"r,attr"`
	Type    string  `xml:"t,attr,omitempty"` // "inlineStr" for text; omitted (numeric/formula) otherwise
	Formula *string `xml:"f,omitempty"`      // formula cells ("=SUM(...)" minus the leading "=")
	Value   *string `xml:"v,omitempty"`      // numeric cells, and a formula's last-known/cached result
	Is      *sIs    `xml:"is,omitempty"`     // inline-string cells
}

type sIs struct {
	T string `xml:"t"`
}

// sWorkbook is the minimal shape of xl/workbook.xml this package
// cares about: just enough to resolve a sheet name/1-based index to
// its relationship id, so ReadXlsxCells/SetXlsxCells can find the
// right worksheet part even in a workbook with more than one sheet
// (including ones not authored by BuildXlsx, e.g. a user upload).
type sWorkbook struct {
	XMLName xml.Name  `xml:"workbook"`
	Sheets  []sWbName `xml:"sheets>sheet"`
}

type sWbName struct {
	Name string `xml:"name,attr"`
	RID  string `xml:"id,attr"` // r:id — encoding/xml matches the local name regardless of the "r" namespace prefix.
}

// sRelationships is the minimal shape of a .rels part this package
// needs: a relationship id -> target path.
type sRelationships struct {
	XMLName xml.Name        `xml:"Relationships"`
	Rel     []sRelationship `xml:"Relationship"`
}

type sRelationship struct {
	ID     string `xml:"Id,attr"`
	Target string `xml:"Target,attr"`
}

// colLetter converts a 0-based column index to its spreadsheet letter
// (0 -> A, 25 -> Z, 26 -> AA, ...).
func colLetter(i int) string {
	s := ""
	for {
		s = string(rune('A'+(i%26))) + s
		i = i/26 - 1
		if i < 0 {
			break
		}
	}
	return s
}

func isNumeric(s string) bool {
	if s == "" {
		return false
	}
	_, err := strconv.ParseFloat(s, 64)
	return err == nil
}

func rowsToSheetXML(rows [][]string) ([]byte, error) {
	sd := sSheetData{}
	for ri, row := range rows {
		r := sRow{R: ri + 1}
		for ci, val := range row {
			cell := sCell{R: colLetter(ci) + strconv.Itoa(ri+1)}
			setCellValue(&cell, val)
			r.Cells = append(r.Cells, cell)
		}
		sd.Rows = append(sd.Rows, r)
	}
	ws := sWorksheet{Xmlns: "http://schemas.openxmlformats.org/spreadsheetml/2006/main", SheetData: sd}
	body, err := xml.Marshal(ws)
	if err != nil {
		return nil, fmt.Errorf("officekit: marshal worksheet xml: %w", err)
	}
	return []byte(xml.Header + string(body)), nil
}

// BuildXlsx renders one or more sheets into a complete, valid .xlsx
// workbook — every sheet given is written as its own worksheet part
// and listed in xl/workbook.xml in order, each with a matching
// relationship in xl/_rels/workbook.xml.rels. An unnamed sheet gets
// the conventional "SheetN" name.
func BuildXlsx(sheets []Sheet) ([]byte, error) {
	if len(sheets) == 0 {
		sheets = []Sheet{{Name: "Sheet1"}}
	}

	parts := []zipPart{
		{Name: "_rels/.rels", Body: xlsxRootRels},
		{Name: "xl/styles.xml", Body: xlsxStyles},
		{Name: "docProps/core.xml", Body: officeCoreProps},
	}

	var sheetEntries, relEntries, overrides strings.Builder
	for i, sh := range sheets {
		idx := i + 1
		name := sh.Name
		if name == "" {
			name = fmt.Sprintf("Sheet%d", idx)
		}
		sheetXML, err := rowsToSheetXML(sh.Rows)
		if err != nil {
			return nil, err
		}
		parts = append(parts, zipPart{Name: xlsxSheetPart(idx), Body: string(sheetXML)})
		fmt.Fprintf(&sheetEntries, `<sheet name="%s" sheetId="%d" r:id="rId%d"/>`, xmlEscapeAttr(name), idx, idx)
		fmt.Fprintf(&relEntries, `<Relationship Id="rId%d" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" Target="worksheets/sheet%d.xml"/>`, idx, idx)
		fmt.Fprintf(&overrides, `<Override PartName="/xl/worksheets/sheet%d.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>`, idx)
	}
	// The styles relationship id just needs to be distinct from every
	// sheet's rId above; one past the last sheet always is.
	fmt.Fprintf(&relEntries, `<Relationship Id="rId%d" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>`, len(sheets)+1)

	workbookXML := xml.Header + `<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"><sheets>` +
		sheetEntries.String() + `</sheets></workbook>`
	workbookRels := xml.Header + `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">` +
		relEntries.String() + `</Relationships>`
	contentTypes := 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="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>` +
		overrides.String() +
		`<Override PartName="/xl/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.styles+xml"/>` +
		`<Override PartName="/docProps/core.xml" ContentType="application/vnd.openxmlformats-package.core-properties+xml"/></Types>`

	parts = append(parts,
		zipPart{Name: "[Content_Types].xml", Body: contentTypes},
		zipPart{Name: xlsxWorkbookPart, Body: workbookXML},
		zipPart{Name: xlsxWorkbookRelsPart, Body: workbookRels},
	)
	return buildZip(parts)
}

// xmlEscapeAttr escapes the handful of characters that matter inside
// an XML attribute value. Used only for the places we interpolate
// user text into a hand-written template string (sheet names in
// workbook.xml) rather than through encoding/xml.
func xmlEscapeAttr(s string) string {
	r := strings.NewReplacer(`&`, "&amp;", `<`, "&lt;", `>`, "&gt;", `"`, "&quot;")
	return r.Replace(s)
}

// ReadXlsxSheetNames returns every sheet name in the workbook, in
// order — used by read_spreadsheet to tell the model what else is in
// a multi-sheet file it only asked to read one sheet of.
func ReadXlsxSheetNames(data []byte) ([]string, error) {
	wb, err := readXlsxWorkbook(data)
	if err != nil {
		return nil, err
	}
	names := make([]string, len(wb.Sheets))
	for i, s := range wb.Sheets {
		names[i] = s.Name
	}
	return names, nil
}

func readXlsxWorkbook(data []byte) (*sWorkbook, error) {
	raw, ok, err := readZipPart(data, xlsxWorkbookPart)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("officekit: %s not found (not a valid .xlsx)", xlsxWorkbookPart)
	}
	var wb sWorkbook
	if err := xml.Unmarshal(raw, &wb); err != nil {
		return nil, fmt.Errorf("officekit: parse workbook xml: %w", err)
	}
	return &wb, nil
}

// resolveXlsxSheetPart maps a sheet reference (empty for the first
// sheet, a 1-based index as a string, or a sheet name) to that sheet's
// actual worksheet part path, by reading xl/workbook.xml and
// xl/_rels/workbook.xml.rels — not by assuming sheet1.xml is sheet
// one, since that isn't guaranteed for a file this package didn't
// write itself (a user-uploaded workbook whose sheets were reordered
// or renumbered in Excel). Falls back to the sheet1.xml convention
// only if the workbook's own bookkeeping can't be parsed at all.
func resolveXlsxSheetPart(data []byte, sheetRef string) (partPath, sheetName string, err error) {
	wb, err := readXlsxWorkbook(data)
	if err != nil || len(wb.Sheets) == 0 {
		return xlsxSheetPart(1), "", nil
	}

	var target *sWbName
	switch {
	case sheetRef == "":
		target = &wb.Sheets[0]
	default:
		if n, convErr := strconv.Atoi(sheetRef); convErr == nil {
			if n < 1 || n > len(wb.Sheets) {
				return "", "", fmt.Errorf("officekit: sheet %d out of range (workbook has %d sheet(s))", n, len(wb.Sheets))
			}
			target = &wb.Sheets[n-1]
		} else {
			for i := range wb.Sheets {
				if strings.EqualFold(wb.Sheets[i].Name, sheetRef) {
					target = &wb.Sheets[i]
					break
				}
			}
			if target == nil {
				return "", "", fmt.Errorf("officekit: no sheet named %q", sheetRef)
			}
		}
	}

	relsRaw, ok, err := readZipPart(data, xlsxWorkbookRelsPart)
	if err != nil || !ok {
		return xlsxSheetPart(1), target.Name, nil
	}
	var rels sRelationships
	if err := xml.Unmarshal(relsRaw, &rels); err != nil {
		return xlsxSheetPart(1), target.Name, nil
	}
	for _, r := range rels.Rel {
		if r.ID == target.RID {
			t := strings.TrimPrefix(r.Target, "/")
			if !strings.HasPrefix(t, "xl/") {
				t = "xl/" + t
			}
			return t, target.Name, nil
		}
	}
	return xlsxSheetPart(1), target.Name, nil
}

// ReadXlsxCells returns one worksheet's cells as a row-major grid of
// display strings, for the read_spreadsheet tool. sheet selects which
// one: empty for the first sheet, a 1-based index ("2"), or a sheet
// name (case-insensitive); the resolved name is returned alongside so
// callers can report it.
func ReadXlsxCells(data []byte, sheet string) ([][]string, string, error) {
	partPath, sheetName, err := resolveXlsxSheetPart(data, sheet)
	if err != nil {
		return nil, "", err
	}
	raw, ok, err := readZipPart(data, partPath)
	if err != nil {
		return nil, "", err
	}
	if !ok {
		return nil, "", fmt.Errorf("officekit: %s not found (not a valid .xlsx)", partPath)
	}
	var ws sWorksheet
	if err := xml.Unmarshal(raw, &ws); err != nil {
		return nil, "", fmt.Errorf("officekit: parse worksheet xml: %w", err)
	}
	var out [][]string
	for _, row := range ws.SheetData.Rows {
		var line []string
		for _, c := range row.Cells {
			switch {
			case c.Is != nil:
				line = append(line, c.Is.T)
			case c.Value != nil:
				// A real Excel-authored formula cell always caches its
				// last-computed value alongside the formula itself; show
				// that computed result rather than the formula text.
				line = append(line, *c.Value)
			case c.Formula != nil:
				// A formula cell this package just wrote has no cached
				// value (nothing here evaluates formulas) — show the
				// formula itself rather than a misleading blank.
				line = append(line, "="+*c.Formula)
			default:
				line = append(line, "")
			}
		}
		out = append(out, line)
	}
	return out, sheetName, nil
}

// SetXlsxCells downloads-parses-modifies-reuploads: applies a sparse
// set of cell writes (keyed by "A1"-style reference) to one worksheet
// (selected the same way as ReadXlsxCells) and returns the new file
// bytes. Cells outside the current grid extend it; existing rows/cells
// are preserved otherwise.
func SetXlsxCells(data []byte, sheet string, writes map[string]string) ([]byte, error) {
	partPath, _, err := resolveXlsxSheetPart(data, sheet)
	if err != nil {
		return nil, err
	}
	raw, ok, err := readZipPart(data, partPath)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("officekit: %s not found (not a valid .xlsx)", partPath)
	}
	var ws sWorksheet
	if err := xml.Unmarshal(raw, &ws); err != nil {
		return nil, fmt.Errorf("officekit: parse worksheet xml: %w", err)
	}

	// rowByNum maps a row number to its INDEX in ws.SheetData.Rows, not
	// a pointer into it. A pointer taken here (&ws.SheetData.Rows[i])
	// is only valid for the backing array that exists at the moment
	// it's taken; the very next `append` below that grows the slice
	// past its capacity allocates a new backing array and copies into
	// it, silently stranding any earlier-taken pointer on the old,
	// now-orphaned one. Writes applied through a stranded pointer
	// still "succeed" from the caller's point of view — no error, no
	// panic — but vanish, because they land on memory nothing else
	// reads. An index has no such failure mode: ws.SheetData.Rows[idx]
	// always re-resolves against whatever backing array the slice
	// header currently points to.
	rowByNum := map[int]int{}
	for i := range ws.SheetData.Rows {
		rowByNum[ws.SheetData.Rows[i].R] = i
	}

	// Writes are applied in ascending (row, column) order rather than
	// map iteration order (which Go deliberately randomizes). This
	// doesn't affect correctness — sortSheetData below restores final
	// order regardless — but it keeps behavior deterministic and easy
	// to reason about/debug when a batch of cell writes is logged or
	// replayed.
	refs := make([]string, 0, len(writes))
	for ref := range writes {
		refs = append(refs, ref)
	}
	sort.Slice(refs, func(a, b int) bool {
		ra, ca, _ := parseCellRef(refs[a])
		rb, cb, _ := parseCellRef(refs[b])
		if ra != rb {
			return ra < rb
		}
		return ca < cb
	})

	for _, ref := range refs {
		val := writes[ref]
		rowNum, _, err := parseCellRef(ref)
		if err != nil {
			return nil, err
		}
		idx, ok := rowByNum[rowNum]
		if !ok {
			ws.SheetData.Rows = append(ws.SheetData.Rows, sRow{R: rowNum})
			idx = len(ws.SheetData.Rows) - 1
			rowByNum[rowNum] = idx
		}
		row := &ws.SheetData.Rows[idx]
		replaced := false
		for i := range row.Cells {
			if row.Cells[i].R == ref {
				setCellValue(&row.Cells[i], val)
				replaced = true
				break
			}
		}
		if !replaced {
			c := sCell{R: ref}
			setCellValue(&c, val)
			row.Cells = append(row.Cells, c)
		}
	}

	// Excel is strict about row/cell order within xl/worksheets/sheetN.xml
	// (rows ascending by r, cells within a row ascending by column) —
	// writes above can append a row/cell out of order (e.g. writing "A1"
	// after "A5" already existed), which a lenient reader shrugs off but
	// Excel itself can flag as needing "repair". Restore ascending order
	// before marshaling so a sparse, arbitrarily-ordered set of writes
	// always produces a schema-valid worksheet.
	sortSheetData(&ws.SheetData)

	body, err := xml.Marshal(ws)
	if err != nil {
		return nil, fmt.Errorf("officekit: marshal worksheet xml: %w", err)
	}
	return replaceZipPart(data, partPath, []byte(xml.Header+string(body)))
}

// sortSheetData reorders rows ascending by row number and, within each
// row, cells ascending by column — the order SetXlsxCells' sparse,
// map-driven writes don't otherwise guarantee.
func sortSheetData(sd *sSheetData) {
	sort.Slice(sd.Rows, func(i, j int) bool { return sd.Rows[i].R < sd.Rows[j].R })
	for i := range sd.Rows {
		cells := sd.Rows[i].Cells
		sort.Slice(cells, func(a, b int) bool { return cellColumn(cells[a].R) < cellColumn(cells[b].R) })
	}
}

// cellColumn returns a cell reference's 0-based column index (e.g.
// "B12" -> 1), or 0 if ref doesn't parse as one — used only to order
// cells within a row for output, so a malformed ref sorts first rather
// than failing a write this deep into the process.
func cellColumn(ref string) int {
	_, col, err := parseCellRef(ref)
	if err != nil {
		return 0
	}
	return col
}

// setCellValue classifies val and sets exactly the sCell fields that
// representation needs, clearing the others: a leading "=" (with at
// least one character after it) stores a formula, same convention
// Excel itself uses to distinguish "=SUM(A1:A2)" from the literal text
// "=SUM(A1:A2)"; otherwise a value that parses as a plain number is
// numeric, and everything else is inline text.
func setCellValue(c *sCell, val string) {
	switch {
	case strings.HasPrefix(val, "=") && len(val) > 1:
		f := val[1:]
		c.Formula, c.Value, c.Is, c.Type = &f, nil, nil, ""
	case isNumeric(val):
		v := val
		c.Formula, c.Value, c.Is, c.Type = nil, &v, nil, ""
	default:
		c.Formula, c.Value = nil, nil
		c.Is, c.Type = &sIs{T: val}, "inlineStr"
	}
}

// parseCellRef splits "B12" into (row=12, col=1). Only used to locate
// a row bucket for SetXlsxCells; the column stays embedded in the "r"
// attribute string as-is.
func parseCellRef(ref string) (row, col int, err error) {
	i := 0
	for i < len(ref) && ref[i] >= 'A' && ref[i] <= 'Z' {
		col = col*26 + int(ref[i]-'A'+1)
		i++
	}
	if i == 0 || i == len(ref) {
		return 0, 0, fmt.Errorf("officekit: invalid cell reference %q (expected e.g. \"A1\")", ref)
	}
	row, err = strconv.Atoi(ref[i:])
	if err != nil || row < 1 {
		return 0, 0, fmt.Errorf("officekit: invalid cell reference %q (expected e.g. \"A1\")", ref)
	}
	return row, col - 1, nil
}

const xlsxRootRels = xml.Header + `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>
<Relationship Id="rId2" Type="http://schemas.openxmlformats.org/package/2006/relationships/metadata/core-properties" Target="docProps/core.xml"/>
</Relationships>`

const xlsxStyles = xml.Header + `<styleSheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">
<fonts count="1"><font><sz val="11"/><name val="Calibri"/></font></fonts>
<fills count="1"><fill><patternFill patternType="none"/></fill></fills>
<borders count="1"><border><left/><right/><top/><bottom/><diagonal/></border></borders>
<cellStyleXfs count="1"><xf numFmtId="0" fontId="0" fillId="0" borderId="0"/></cellStyleXfs>
<cellXfs count="1"><xf numFmtId="0" fontId="0" fillId="0" borderId="0" xfId="0"/></cellXfs>
</styleSheet>`