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.
func (d *Document) Map() map[string]any { 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 { return d.footer }

// SetFooter replaces those lines.
func (d *Document) SetFooter(lines []string) { 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

	// 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.
func (t *Table) Keys() []string {
	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 { return t.values }

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

// Get returns the entry for key, and whether the table has one.
func (t *Table) Get(key string) (*Entry, bool) {
	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.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 { 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) { t.comments = lines }

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

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

// addValue records a key of the table, in written order.
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
	if node, ok := val.(map[string]any); ok {
		e.child = newTable(node)
	}
	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 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) {
	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) {
	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) {
	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) {
	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) {
	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; 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) {
	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:
		if e.child == nil {
			e.child = newOrderedTable(v)
		} else {
			e.child.values = 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 as entries in sorted order, the order Marshal writes maps in.
func newOrderedTable(m map[string]any) *Table {
	t := newTable(m)
	for _, k := range slices.Sorted(maps.Keys(m)) {
		v := m[k]
		_, isMap := v.(map[string]any)
		e := t.addValue(k, v, false)
		if isMap {
			e.child = newOrderedTable(v.(map[string]any))
		}
	}
	return t
}

// Delete removes key and everything it holds.
func (t *Table) Delete(key string) {
	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
		}
	}
}