interpres/document.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package interpres

import (
	"maps"
	"slices"
)

// A Document is a parsed TOML document: the values, plus what the map shape
// cannot carry, which is the order the keys were written in, whether a table
// was written inline or under a header, and the comments.
//
// The values are the tree ParseMap returns, shared rather than copied, so a
// value read from a Document and from that map is the same value. A comment
// belongs to the statement it precedes: the lines above a key belong to the
// key, the lines above a header belong to the header's table, and a comment
// block after the last statement belongs to the Document.
//
// Comments inside array and inline-table values are not carried yet; the
// parser skips them as it always has.
type Document struct {
	root   *Table
	footer []string
}

// Root returns the document's root table. A nil document has no root.
func (d *Document) Root() *Table {
	if d == nil {
		return nil
	}
	return d.root
}

// Map returns the value tree, the shape ParseMap gives. It is the tree the
// document was parsed into, not a copy. A nil document or one with no root
// holds no values.
func (d *Document) Map() map[string]any {
	if d == nil || d.root == nil {
		return nil
	}
	return d.root.values
}

// Footer returns the comment lines that follow the last statement, and every
// line of a document that holds no statement at all.
func (d *Document) Footer() []string {
	if d == nil {
		return nil
	}
	return d.footer
}

// SetFooter replaces those lines.
func (d *Document) SetFooter(lines []string) {
	if d == nil {
		return
	}
	d.footer = lines
}

// The document-level convenience forms of the Table edit API; they act on
// the root table.

// Get returns the root table's entry for key, and whether the document has
// one. See Table.Get.
func (d *Document) Get(key string) (*Entry, bool) { return d.Root().Get(key) }

// GetString returns the string the key holds, and whether it holds one.
func (d *Document) GetString(key string) (string, bool) { return d.Root().GetString(key) }

// GetInt returns the integer the key holds, and whether it holds one.
func (d *Document) GetInt(key string) (int64, bool) { return d.Root().GetInt(key) }

// GetFloat returns the float the key holds, and whether it holds one.
func (d *Document) GetFloat(key string) (float64, bool) { return d.Root().GetFloat(key) }

// GetBool returns the boolean the key holds, and whether it holds one.
func (d *Document) GetBool(key string) (bool, bool) { return d.Root().GetBool(key) }

// GetArray returns the value array the key holds, and whether it holds one.
func (d *Document) GetArray(key string) ([]any, bool) { return d.Root().GetArray(key) }

// GetTable returns the node of the table the key holds, and whether it holds
// one.
func (d *Document) GetTable(key string) (*Table, bool) { return d.Root().GetTable(key) }

// Set stores value under the key in the root table. See Table.Set.
func (d *Document) Set(key string, value any) { d.Root().Set(key, value) }

// Delete removes the key from the root table. See Table.Delete.
func (d *Document) Delete(key string) { d.Root().Delete(key) }

// A Table is one TOML table: its values, its keys in written order, and the
// comments around the header or the key that introduced it.
type Table struct {
	values  map[string]any
	entries []*Entry
	index   map[string]*Entry

	// inline records that the table was written as an inline table, `{…}`,
	// rather than under a header or as a dotted key.
	inline bool

	// dotted records that a dotted key introduced the table, `a.b = 1`
	// building the a around the leaf: the write side gives such a table back
	// as dotted key lines, the form that holds the position of a line.
	dotted bool

	// comments are the lines above the table's header, trailing is the comment
	// on the header's own line. Both are empty for a table a dotted key
	// introduced, which has no line of its own.
	comments []string
	trailing string
}

func newTable(values map[string]any) *Table {
	return &Table{values: values, index: map[string]*Entry{}}
}

// Keys returns the table's keys in the order they were written. A nil table
// holds none, the answer a document without a root gives through Root.
func (t *Table) Keys() []string {
	if t == nil {
		return nil
	}
	keys := make([]string, len(t.entries))
	for i, e := range t.entries {
		keys[i] = e.key
	}
	return keys
}

// Values returns the table's values, which is the map the value tree holds for
// it.
func (t *Table) Values() map[string]any {
	if t == nil {
		return nil
	}
	return t.values
}

// Entries returns the table's entries in written order.
func (t *Table) Entries() []*Entry {
	if t == nil {
		return nil
	}
	return t.entries
}

// Get returns the entry for key, and whether the table has one.
func (t *Table) Get(key string) (*Entry, bool) {
	if t == nil {
		return nil, false
	}
	e, ok := t.index[key]
	return e, ok
}

// Inline reports whether the table was written as an inline table, `{…}`,
// rather than under a header or introduced by a dotted key.
func (t *Table) Inline() bool { return t != nil && t.inline }

// Comments returns the comment lines above the table's header, or above the
// key that introduced it. Lines carry no leading '#' and no surrounding space.
func (t *Table) Comments() []string {
	if t == nil {
		return nil
	}
	return t.comments
}

// SetComments replaces those lines. Each line is written back with a "# " in
// front of it, so a line should not carry one.
func (t *Table) SetComments(lines []string) {
	if t == nil {
		return
	}
	t.comments = lines
}

// Trailing returns the comment on the header's own line, without the '#'.
func (t *Table) Trailing() string {
	if t == nil {
		return ""
	}
	return t.trailing
}

// SetTrailing replaces that comment.
func (t *Table) SetTrailing(line string) {
	if t == nil {
		return
	}
	t.trailing = line
}

// addValue records a key of the table, in written order. The caller gives
// the entry a table node or element nodes when the value has that shape; a
// map value left without a node writes as an inline table.
func (t *Table) addValue(key string, val any, inline bool) *Entry {
	e := &Entry{table: t, key: key, inline: inline}
	t.entries = append(t.entries, e)
	t.index[key] = e
	return e
}

// addTable records a key whose value is a table introduced by a header or a
// dotted key, and returns the table's node.
func (t *Table) addTable(key string, values map[string]any) *Table {
	if e, ok := t.index[key]; ok {
		// The key was seen before, as the leaf of an earlier dotted key.
		if e.child == nil {
			e.child = newTable(values)
		}
		return e.child
	}
	e := t.addValue(key, values, false)
	e.child = newTable(values)
	return e.child
}

// addElement records one element of an array of tables, and returns its node.
func (t *Table) addElement(key string, values map[string]any) *Table {
	e, ok := t.index[key]
	if !ok {
		e = t.addValue(key, nil, false)
		e.elements = []*Table{}
	}
	el := newTable(values)
	e.elements = append(e.elements, el)
	return el
}

// child returns the node of a table-valued key, or nil.
func (t *Table) child(key string) *Table {
	if t == nil {
		return nil
	}
	if e, ok := t.index[key]; ok {
		return e.child
	}
	return nil
}

// lastElement returns the node of the newest element of an array of tables.
func (t *Table) lastElement(key string) *Table {
	if e, ok := t.index[key]; ok && len(e.elements) > 0 {
		return e.elements[len(e.elements)-1]
	}
	return nil
}

// An Entry is one key of a table: the value and the comments around the key.
type Entry struct {
	table    *Table
	key      string
	inline   bool
	comments []string
	trailing string

	// child is the table the value is, and elements are the tables of an array
	// of tables; one of them is set only when the value has that shape.
	child    *Table
	elements []*Table
}

// Key returns the key as it was written.
func (e *Entry) Key() string { return e.key }

// Value returns the value the key holds. It is read from the table's map, so
// it stays current if that map is changed.
func (e *Entry) Value() any { return e.table.values[e.key] }

// Inline reports whether the value was written as an inline table, `{…}`.
func (e *Entry) Inline() bool { return e.inline }

// Table returns the table the value is, or nil when it is not a table.
func (e *Entry) Table() *Table { return e.child }

// Elements returns the tables of an array of tables, or nil when the value is
// not one.
func (e *Entry) Elements() []*Table { return e.elements }

// Comments returns the comment lines above the key. Lines carry no leading '#'
// and no surrounding space.
func (e *Entry) Comments() []string { return e.comments }

// SetComments replaces those lines. Each line is written back with a "# " in
// front of it, so a line should not carry one.
func (e *Entry) SetComments(lines []string) { e.comments = lines }

// Trailing returns the comment on the key's own line, without the '#'.
func (e *Entry) Trailing() string { return e.trailing }

// SetTrailing replaces that comment.
func (e *Entry) SetTrailing(line string) { e.trailing = line }

// GetString returns the string the key holds, and whether it holds one.
func (t *Table) GetString(key string) (string, bool) {
	if t == nil {
		return "", false
	}
	v, ok := t.values[key]
	s, ok := v.(string)
	return s, ok
}

// GetInt returns the integer the key holds, and whether it holds one.
func (t *Table) GetInt(key string) (int64, bool) {
	if t == nil {
		return 0, false
	}
	v, ok := t.values[key]
	i, ok := v.(int64)
	return i, ok
}

// GetFloat returns the float the key holds, and whether it holds one.
func (t *Table) GetFloat(key string) (float64, bool) {
	if t == nil {
		return 0, false
	}
	v, ok := t.values[key]
	f, ok := v.(float64)
	return f, ok
}

// GetBool returns the boolean the key holds, and whether it holds one.
func (t *Table) GetBool(key string) (bool, bool) {
	if t == nil {
		return false, false
	}
	v, ok := t.values[key]
	b, ok := v.(bool)
	return b, ok
}

// GetArray returns the value array the key holds, and whether it holds one.
func (t *Table) GetArray(key string) ([]any, bool) {
	if t == nil {
		return nil, false
	}
	v, ok := t.values[key]
	a, ok := v.([]any)
	return a, ok
}

// GetTable returns the node of the table the key holds, and whether it holds
// one, whichever way the document wrote the table.
func (t *Table) GetTable(key string) (*Table, bool) {
	c := t.child(key)
	return c, c != nil
}

// Set stores value under key. A key the table already has keeps its position
// and its comments; a new one joins the end. A value of map[string]any
// becomes a table node of its own, written under a header like any other
// table, and replaces the node the key held, which belonged to the value the
// key held; a Go map carries no order, so its keys take sorted order. A value
// of []map[string]any becomes an array-of-tables node.
func (t *Table) Set(key string, value any) {
	if t == nil {
		return
	}
	e, ok := t.index[key]
	if !ok {
		t.values[key] = value
		e = t.addValue(key, value, false)
		if m, isMap := value.(map[string]any); isMap {
			e.child = newOrderedTable(m)
		}
		if items, isArray := value.([]map[string]any); isArray {
			e.elements = make([]*Table, len(items))
			for i, item := range items {
				e.elements[i] = newOrderedTable(item)
			}
		}
		return
	}
	t.values[key] = value
	switch v := value.(type) {
	case map[string]any:
		// The node is rebuilt rather than patched: the entries and the index
		// belong to the table the key held, and writing the new value
		// through them would leave the old table's keys in the output.
		e.child = newOrderedTable(v)
		e.elements = nil
	case []map[string]any:
		e.child = nil
		e.elements = make([]*Table, len(v))
		for i, item := range v {
			e.elements[i] = newOrderedTable(item)
		}
	default:
		e.child = nil
		e.elements = nil
	}
}

// newOrderedTable builds a table node for a value the caller set, its keys
// entered in sorted order, the order Marshal writes maps in.
func newOrderedTable(m map[string]any) *Table {
	return orderedTable(m, 0)
}

// orderedTable is newOrderedTable's recursion. The depth bound is the value
// encoder's: a cyclic map stopped here is written by the value writer, which
// reports it instead of running the stack out.
func orderedTable(m map[string]any, depth int) *Table {
	t := newTable(m)
	for _, k := range slices.Sorted(maps.Keys(m)) {
		v := m[k]
		e := t.addValue(k, v, false)
		if depth >= maxEncodeDepth {
			continue
		}
		switch val := v.(type) {
		case map[string]any:
			e.child = orderedTable(val, depth+1)
		case []map[string]any:
			e.elements = make([]*Table, len(val))
			for i, item := range val {
				e.elements[i] = orderedTable(item, depth+1)
			}
		}
	}
	return t
}

// Delete removes key and everything it holds.
func (t *Table) Delete(key string) {
	if t == nil {
		return
	}
	if _, ok := t.values[key]; !ok {
		return
	}
	delete(t.values, key)
	delete(t.index, key)
	for i, e := range t.entries {
		if e.key == key {
			t.entries = append(t.entries[:i], t.entries[i+1:]...)
			break
		}
	}
}