dataxl/cmd/dataxl/table.go

package main

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"html"
	"sort"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/xuri/excelize/v2"
)

// table is the common representation for CSV, TSV, XLSX and Markdown. Rows are padded
// to Header width when read, which keeps subsequent conversions rectangular.
type table struct {
	Header []string
	Rows   [][]string
}

func parseTable(input []byte, format, sheet string) (table, error) {
	switch format {
	case "csv":
		return readDelimited(input, ',')
	case "tsv":
		return readDelimited(input, '\t')
	case "xlsx":
		f, err := excelize.OpenReader(bytes.NewReader(input))
		if err != nil {
			return table{}, err
		}
		defer f.Close()
		rows, err := f.GetRows(sheet)
		if err != nil {
			return table{}, err
		}
		return rowsToTable(rows)
	case "markdown":
		return readMarkdown(input)
	default:
		return table{}, fmt.Errorf("format %q is not tabular", format)
	}
}

func readDelimited(input []byte, comma rune) (table, error) {
	// Spreadsheet applications commonly prefix UTF-8 CSV/TSV exports with a
	// BOM. It is an encoding marker, not part of the first column name.
	input = bytes.TrimPrefix(input, []byte{0xEF, 0xBB, 0xBF})
	if !utf8.Valid(input) {
		return table{}, fmt.Errorf("delimited input is not valid UTF-8")
	}
	r := csv.NewReader(bytes.NewReader(input))
	r.Comma = comma
	r.FieldsPerRecord = -1
	rows, err := r.ReadAll()
	if err != nil {
		return table{}, err
	}
	return rowsToTable(rows)
}

func rowsToTable(rows [][]string) (table, error) {
	if len(rows) == 0 {
		return table{}, nil
	}
	width := len(rows[0])
	for rowIndex, row := range rows[1:] {
		if len(row) > width {
			return table{}, fmt.Errorf(
				"row %d has %d fields but the header has %d",
				rowIndex+2, len(row), width,
			)
		}
	}
	header := padRow(rows[0], width)
	body := make([][]string, 0, len(rows)-1)
	for _, row := range rows[1:] {
		body = append(body, padRow(row, width))
	}
	return table{Header: header, Rows: body}, nil
}

func padRow(row []string, width int) []string {
	out := make([]string, width)
	copy(out, row)
	return out
}

func encodeTable(t table, format, sheet string) ([]byte, error) {
	switch format {
	case "csv":
		return writeDelimited(t, ',')
	case "tsv":
		return writeDelimited(t, '\t')
	case "xlsx":
		return writeXLSX(t, sheet)
	case "markdown":
		return writeMarkdown(t)
	default:
		return nil, fmt.Errorf("format %q is not tabular", format)
	}
}

func readMarkdown(input []byte) (table, error) {
	input = bytes.TrimPrefix(input, []byte{0xEF, 0xBB, 0xBF})
	if !utf8.Valid(input) {
		return table{}, fmt.Errorf("markdown input is not valid UTF-8")
	}
	text := strings.ReplaceAll(string(input), "\r\n", "\n")
	text = strings.TrimSpace(text)
	if text == "" {
		return table{}, nil
	}
	lines := strings.Split(text, "\n")
	if len(lines) < 2 {
		return table{}, fmt.Errorf("markdown table requires a header and separator row")
	}
	header, err := parseMarkdownRow(lines[0])
	if err != nil {
		return table{}, fmt.Errorf("markdown header: %w", err)
	}
	separator, err := parseMarkdownRow(lines[1])
	if err != nil {
		return table{}, fmt.Errorf("markdown separator: %w", err)
	}
	if len(separator) != len(header) {
		return table{}, fmt.Errorf("markdown separator has %d fields but the header has %d", len(separator), len(header))
	}
	for i, cell := range separator {
		value := strings.TrimSpace(cell)
		value = strings.TrimPrefix(value, ":")
		value = strings.TrimSuffix(value, ":")
		if len(value) < 3 || strings.Trim(value, "-") != "" {
			return table{}, fmt.Errorf("markdown separator field %d is invalid", i+1)
		}
	}
	rows := make([][]string, 0, len(lines)-2)
	for i, line := range lines[2:] {
		if strings.TrimSpace(line) == "" {
			continue
		}
		row, err := parseMarkdownRow(line)
		if err != nil {
			return table{}, fmt.Errorf("markdown row %d: %w", i+3, err)
		}
		if len(row) > len(header) {
			return table{}, fmt.Errorf("row %d has %d fields but the header has %d", i+3, len(row), len(header))
		}
		rows = append(rows, padRow(row, len(header)))
	}
	return table{Header: header, Rows: rows}, nil
}

func parseMarkdownRow(line string) ([]string, error) {
	line = strings.TrimSpace(strings.TrimSuffix(line, "\r"))
	if strings.HasPrefix(line, "|") {
		line = line[1:]
	}
	if hasUnescapedTrailingPipe(line) {
		line = strings.TrimSpace(line[:len(line)-1])
	}
	var cells []string
	var cell strings.Builder
	escaped := false
	for _, r := range line {
		if escaped {
			cell.WriteRune('\\')
			cell.WriteRune(r)
			escaped = false
			continue
		}
		if r == '\\' {
			escaped = true
			continue
		}
		if r == '|' {
			cells = append(cells, decodeMarkdownCell(strings.Trim(cell.String(), " \t")))
			cell.Reset()
			continue
		}
		cell.WriteRune(r)
	}
	if escaped {
		return nil, fmt.Errorf("row ends with an incomplete escape")
	}
	cells = append(cells, decodeMarkdownCell(strings.Trim(cell.String(), " \t")))
	return cells, nil
}

func decodeMarkdownCell(value string) string {
	var b strings.Builder
	escaped := false
	for _, r := range value {
		if escaped {
			if r == '\\' || r == '|' {
				b.WriteRune(r)
			} else {
				b.WriteRune('\\')
				b.WriteRune(r)
			}
			escaped = false
			continue
		}
		if r == '\\' {
			escaped = true
			continue
		}
		b.WriteRune(r)
	}
	return html.UnescapeString(b.String())
}

func hasUnescapedTrailingPipe(line string) bool {
	line = strings.TrimSpace(line)
	if !strings.HasSuffix(line, "|") {
		return false
	}
	backslashes := 0
	for i := len(line) - 2; i >= 0 && line[i] == '\\'; i-- {
		backslashes++
	}
	return backslashes%2 == 0
}

func writeMarkdown(t table) ([]byte, error) {
	if len(t.Header) == 0 && len(t.Rows) == 0 {
		return nil, nil
	}
	if len(t.Header) == 0 {
		return nil, fmt.Errorf("markdown table has rows but no header")
	}
	var b strings.Builder
	writeRow := func(row []string) {
		b.WriteString("| ")
		for i, value := range row {
			if i > 0 {
				b.WriteString(" | ")
			}
			b.WriteString(escapeMarkdownCell(value))
		}
		b.WriteString(" |\n")
	}
	writeRow(t.Header)
	b.WriteString("|")
	for range t.Header {
		b.WriteString(" --- |")
	}
	b.WriteByte('\n')
	for _, row := range t.Rows {
		writeRow(padRow(row, len(t.Header)))
	}
	return []byte(b.String()), nil
}

func escapeMarkdownCell(value string) string {
	leadingEnd := len(value) - len(strings.TrimLeft(value, " \t"))
	trailingStart := len(strings.TrimRight(value, " \t"))
	var b strings.Builder
	for i, r := range value {
		switch r {
		case '\\':
			b.WriteString("\\\\")
		case '|':
			b.WriteString("\\|")
		case '\r':
			b.WriteString("&#13;")
		case '\n':
			b.WriteString("&#10;")
		case ' ':
			if i < leadingEnd || i >= trailingStart {
				b.WriteString("&#32;")
			} else {
				b.WriteByte(' ')
			}
		case '\t':
			if i < leadingEnd || i >= trailingStart {
				b.WriteString("&#9;")
			} else {
				b.WriteByte('\t')
			}
		default:
			switch r {
			case '&':
				b.WriteString("&amp;")
			case '<':
				b.WriteString("&lt;")
			case '>':
				b.WriteString("&gt;")
			default:
				b.WriteRune(r)
			}
		}
	}
	return b.String()
}

func writeDelimited(t table, comma rune) ([]byte, error) {
	var b bytes.Buffer
	w := csv.NewWriter(&b)
	w.Comma = comma
	writeRecord := func(row []string) error {
		if len(row) == 1 && row[0] == "" {
			// encoding/csv writes this record as a bare empty line, which its
			// Reader intentionally skips. Quote the field so the row survives a
			// table-to-table round trip.
			w.Flush()
			if err := w.Error(); err != nil {
				return err
			}
			b.WriteString("\"\"\n")
			return nil
		}
		return w.Write(row)
	}
	if err := writeRecord(t.Header); err != nil {
		return nil, err
	}
	for _, row := range t.Rows {
		if err := writeRecord(row); err != nil {
			return nil, err
		}
	}
	w.Flush()
	return b.Bytes(), w.Error()
}

func writeXLSX(t table, sheet string) ([]byte, error) {
	f := excelize.NewFile()
	defer f.Close()
	defaultSheet := f.GetSheetName(0)
	if sheet == "" {
		sheet = "Sheet1"
	}
	if defaultSheet != sheet {
		if err := f.SetSheetName(defaultSheet, sheet); err != nil {
			return nil, err
		}
	}
	rows := append([][]string{t.Header}, t.Rows...)
	for rowIndex, row := range rows {
		for columnIndex, value := range row {
			cell, err := excelize.CoordinatesToCellName(columnIndex+1, rowIndex+1)
			if err != nil {
				return nil, err
			}
			if err := f.SetCellValue(sheet, cell, value); err != nil {
				return nil, err
			}
		}
	}
	if err := styleHeader(f, sheet, len(t.Header)); err != nil {
		return nil, err
	}
	var b bytes.Buffer
	if err := f.Write(&b); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

// styleHeader applies editing conveniences without trying to preserve or
// emulate arbitrary workbook formatting.
func styleHeader(f *excelize.File, sheet string, width int) error {
	if width == 0 {
		return nil
	}
	end, err := excelize.CoordinatesToCellName(width, 1)
	if err != nil {
		return err
	}
	style, err := f.NewStyle(&excelize.Style{Font: &excelize.Font{Bold: true}})
	if err != nil {
		return err
	}
	if err := f.SetCellStyle(sheet, "A1", end, style); err != nil {
		return err
	}
	return f.SetPanes(sheet, &excelize.Panes{
		Freeze: true, YSplit: 1, TopLeftCell: "A2", ActivePane: "bottomLeft",
	})
}

func valueToTable(value any) table {
	records := recordsFromValue(value)
	flatRows := make([]map[string]string, 0, len(records))
	seen := make(map[string]bool)
	var header []string
	for _, record := range records {
		flat := make(map[string]string)
		flatten("", record, flat)
		for key := range flat {
			if !seen[key] {
				seen[key] = true
				header = append(header, key)
			}
		}
		flatRows = append(flatRows, flat)
	}
	// Stable column ordering makes generated files and their diffs predictable.
	sort.Strings(header)
	rows := make([][]string, 0, len(flatRows))
	for _, flat := range flatRows {
		row := make([]string, len(header))
		for i, key := range header {
			row[i] = flat[key]
		}
		rows = append(rows, row)
	}
	return table{Header: header, Rows: rows}
}

func recordsFromValue(value any) []any {
	switch v := value.(type) {
	case []any:
		return v
	case map[string]any:
		// A sole wrapper key lets object-shaped formats such as TOML carry row
		// arrays. If siblings exist, unwrapping would silently discard them, so
		// preserve the complete object as one table record instead.
		if len(v) == 1 {
			for _, key := range []string{"rows", "records", "items"} {
				if rows, ok := v[key].([]any); ok {
					return rows
				}
			}
		}
		return []any{v}
	default:
		return []any{v}
	}
}

// flatten converts nested values to column paths such as user.name and
// items[0].sku. Empty maps and arrays use explicit textual markers; these are
// visible to editors but currently return as strings when converted back.
func flatten(prefix string, value any, out map[string]string) {
	switch v := value.(type) {
	case map[string]any:
		if len(v) == 0 && prefix != "" {
			out[prefix] = "{}"
			return
		}
		keys := make([]string, 0, len(v))
		for key := range v {
			keys = append(keys, key)
		}
		sort.Strings(keys)
		for _, key := range keys {
			childPrefix := appendPathKey(prefix, key)
			flatten(childPrefix, v[key], out)
		}
	case []any:
		if len(v) == 0 && prefix != "" {
			out[prefix] = "[]"
			return
		}
		for i, child := range v {
			flatten(fmt.Sprintf("%s[%d]", prefix, i), child, out)
		}
	default:
		if prefix == "" {
			prefix = "value"
		}
		out[prefix] = scalarString(v)
	}
}

func scalarString(value any) string {
	switch v := value.(type) {
	case nil:
		return ""
	case json.Number:
		return v.String()
	case string:
		return v
	case bool:
		return strconv.FormatBool(v)
	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
		return fmt.Sprint(v)
	default:
		data, err := json.Marshal(v)
		if err != nil {
			return fmt.Sprint(v)
		}
		return string(data)
	}
}