1// Copyright (c) 2021, Peter Ohler, All rights reserved.
2
3package pretty
4
5import (
6	"fmt"
7	"io"
8	"math"
9
10	"github.com/ohler55/ojg"
11)
12
13const (
14	nullStr  = "null"
15	trueStr  = "true"
16	falseStr = "false"
17	spaces   = "\n                                                                " +
18		"                                                                "
19)
20
21// Writer writes data in either JSON or SEN format using setting to determine
22// the output.
23type Writer struct {
24	ojg.Options
25
26	// Width is the suggested maximum width. In some cases it may not be
27	// possible to stay withing the specified width.
28	Width int
29
30	// MaxDepth is the maximum depth of an element on a single line.
31	MaxDepth int
32
33	// Align if true attempts to align elements of children in list.
34	Align bool
35
36	// SEN format if true otherwise JSON encoding.
37	SEN bool
38
39	buf []byte
40	w   io.Writer
41}
42
43// Encode data. Any panics during encoding will cause an empty return but will
44// not fail.
45func (w *Writer) Encode(data interface{}) []byte {
46	b, _ := w.encode(data)
47
48	return b
49}
50
51// Marshal data. The same as Encode but a panics during encoding will result
52// in an error return.
53func (w *Writer) Marshal(data interface{}) ([]byte, error) {
54	return w.encode(data)
55}
56
57// Write encoded data to the op.Writer.
58func (w *Writer) Write(wr io.Writer, data interface{}) (err error) {
59	w.w = wr
60	_, err = w.encode(data)
61
62	return
63}
64func (w *Writer) config(args []interface{}) {
65	for _, arg := range args {
66		switch ta := arg.(type) {
67		case int:
68			w.Width = ta
69		case float64:
70			if 0.0 < ta {
71				if ta < 1.0 {
72					w.MaxDepth = int(math.Round(ta * 10.0))
73				} else {
74					w.Width = int(ta)
75					w.MaxDepth = int(math.Round((ta - float64(w.Width)) * 10.0))
76				}
77				if w.MaxDepth == 0 { // use the default
78					w.MaxDepth = 2
79				}
80			}
81		case bool:
82			w.Align = ta
83		case *ojg.Options:
84			sw := w.w
85			w.Options = *ta
86			w.w = sw
87		}
88	}
89}
90
91func (w *Writer) encode(data interface{}) (out []byte, err error) {
92	if w.InitSize == 0 {
93		w.InitSize = 256
94	}
95	if len(spaces)-1 < w.Width {
96		w.Width = len(spaces) - 1
97	}
98	if w.WriteLimit == 0 {
99		w.WriteLimit = 1024
100	}
101	if cap(w.buf) < w.InitSize {
102		w.buf = make([]byte, 0, w.InitSize)
103	} else {
104		w.buf = w.buf[:0]
105	}
106	defer func() {
107		if r := recover(); r != nil {
108			if err, _ = r.(error); err == nil {
109				err = fmt.Errorf("%v", r)
110				out = []byte{}
111				if w.Color && w.w != nil {
112					_, err = w.w.Write([]byte(w.NoColor))
113				}
114			}
115		}
116	}()
117	tree := w.build(data)
118	w.buf = w.buf[:0]
119	w.Indent = 2
120	if w.Width*3/8 < tree.depth {
121		w.Indent = 1
122	}
123	w.fill(tree, 0, false)
124	if w.w != nil && 0 < len(w.buf) {
125		_, err = w.w.Write(w.buf)
126		w.buf = w.buf[:0]
127	}
128	out = w.buf
129
130	return
131}
132
133func (w *Writer) fill(n *node, depth int, flat bool) {
134	start := depth * w.Indent
135	switch n.kind {
136	case strNode, numNode:
137		w.buf = append(w.buf, n.buf...)
138	case arrayNode:
139		var comma []byte
140		if w.Color {
141			if !w.SEN {
142				comma = append(comma, w.SyntaxColor...)
143				comma = append(comma, ',')
144				comma = append(comma, w.NoColor...)
145			}
146			w.buf = append(w.buf, w.SyntaxColor...)
147			w.buf = append(w.buf, '[')
148			w.buf = append(w.buf, w.NoColor...)
149		} else {
150			if !w.SEN {
151				comma = append(comma, ',')
152			}
153			w.buf = append(w.buf, '[')
154		}
155		if !flat && start+n.size < w.Width && n.depth < w.MaxDepth {
156			flat = true
157		}
158		d2 := depth + 1
159		var cs []byte
160		var is []byte
161
162		if flat {
163			cs = []byte{' '}
164		} else {
165			x := d2*w.Indent + 1
166			if len(spaces) < x {
167				flat = true
168			} else {
169				cs = []byte(spaces[0:x])
170				x = depth*w.Indent + 1
171				is = []byte(spaces[0:x])
172			}
173		}
174		if !w.Align || w.MaxDepth < n.depth || len(n.members) < 2 || w.checkAlign(n, start, comma, cs) {
175			for i, m := range n.members {
176				if 0 < i {
177					w.buf = append(w.buf, comma...)
178					w.buf = append(w.buf, []byte(cs)...)
179				} else if !flat {
180					w.buf = append(w.buf, []byte(cs)...)
181				}
182				w.fill(m, d2, flat)
183			}
184		}
185		w.buf = append(w.buf, []byte(is)...)
186		if w.Color {
187			w.buf = append(w.buf, w.SyntaxColor...)
188			w.buf = append(w.buf, ']')
189			w.buf = append(w.buf, w.NoColor...)
190		} else {
191			w.buf = append(w.buf, ']')
192		}
193	case mapNode:
194		var comma []byte
195		if w.Color {
196			if !w.SEN {
197				comma = append(comma, w.SyntaxColor...)
198				comma = append(comma, ',')
199				comma = append(comma, w.NoColor...)
200			}
201			w.buf = append(w.buf, w.SyntaxColor...)
202			w.buf = append(w.buf, '{')
203			w.buf = append(w.buf, w.NoColor...)
204		} else {
205			if !w.SEN {
206				comma = append(comma, ',')
207			}
208			w.buf = append(w.buf, '{')
209		}
210		d2 := depth + 1
211		var cs []byte
212		var is []byte
213
214		if !flat && start+n.size < w.Width && n.depth < w.MaxDepth {
215			flat = true
216		}
217		if flat {
218			cs = []byte{' '}
219		} else {
220			x := d2*w.Indent + 1
221			if len(spaces) < x {
222				flat = true
223			} else {
224				cs = []byte(spaces[0:x])
225				x = depth*w.Indent + 1
226				is = []byte(spaces[0:x])
227			}
228		}
229		keyWidth := 1
230		if w.Align {
231			for _, m := range n.members {
232				if keyWidth < len(m.key) {
233					keyWidth = len(m.key)
234				}
235			}
236		}
237		for i, m := range n.members {
238			if 0 < i {
239				w.buf = append(w.buf, comma...)
240				w.buf = append(w.buf, []byte(cs)...)
241			} else if !flat {
242				w.buf = append(w.buf, []byte(cs)...)
243			}
244			w.buf = append(w.buf, m.key...)
245			if w.Color {
246				w.buf = append(w.buf, w.SyntaxColor...)
247				w.buf = append(w.buf, ':')
248				w.buf = append(w.buf, w.NoColor...)
249				w.buf = append(w.buf, ' ')
250			} else {
251				w.buf = append(w.buf, ": "...)
252			}
253			for i := keyWidth - len(m.key); 0 < i; i-- {
254				w.buf = append(w.buf, ' ')
255			}
256			w.fill(m, d2, flat)
257		}
258		w.buf = append(w.buf, []byte(is)...)
259		if w.Color {
260			w.buf = append(w.buf, w.SyntaxColor...)
261			w.buf = append(w.buf, '}')
262			w.buf = append(w.buf, w.NoColor...)
263		} else {
264			w.buf = append(w.buf, '}')
265		}
266	}
267	if w.w != nil && w.WriteLimit < len(w.buf) {
268		if _, err := w.w.Write(w.buf); err != nil {
269			panic(err)
270		}
271		w.buf = w.buf[:0]
272	}
273}
274
275// Return true if not filled.
276func (w *Writer) checkAlign(n *node, start int, comma, cs []byte) bool {
277	c := n.genTables(w.SEN)
278	if c == nil || w.Width < start+c.size {
279		return true
280	}
281	for i, m := range n.members {
282		if 0 < i {
283			w.buf = append(w.buf, comma...)
284		}
285		w.buf = append(w.buf, []byte(cs)...)
286		switch m.kind {
287		case arrayNode:
288			w.alignArray(m, c, comma, cs)
289		case mapNode:
290			w.alignMap(m, c, comma, cs)
291		}
292	}
293	return false
294}
295
296func (w *Writer) alignArray(n *node, t *table, comma, cs []byte) {
297	if w.Color {
298		w.buf = append(w.buf, w.SyntaxColor...)
299		w.buf = append(w.buf, '[')
300		w.buf = append(w.buf, w.NoColor...)
301	} else {
302		w.buf = append(w.buf, '[')
303	}
304	for k, col := range t.columns {
305		if len(n.members) <= k {
306			break
307		}
308		if 0 < k {
309			w.buf = append(w.buf, comma...)
310			w.buf = append(w.buf, ' ')
311		}
312		m := n.members[k]
313		cw := col.size
314		switch m.kind {
315		case strNode:
316			w.buf = append(w.buf, m.buf...)
317			if m.size < cw {
318				w.buf = append(w.buf, spaces[1:cw-m.size+1]...)
319			}
320		case numNode:
321			if m.size < cw {
322				w.buf = append(w.buf, spaces[1:cw-m.size+1]...)
323			}
324			w.buf = append(w.buf, m.buf...)
325		case arrayNode:
326			w.alignArray(m, col, comma, []byte{' '})
327		case mapNode:
328			w.alignMap(m, col, comma, []byte{' '})
329		}
330	}
331	if w.Color {
332		w.buf = append(w.buf, w.SyntaxColor...)
333		w.buf = append(w.buf, ']')
334		w.buf = append(w.buf, w.NoColor...)
335	} else {
336		w.buf = append(w.buf, ']')
337	}
338}
339
340func (w *Writer) alignMap(n *node, t *table, comma, cs []byte) {
341	if w.Color {
342		w.buf = append(w.buf, w.SyntaxColor...)
343		w.buf = append(w.buf, '{')
344		w.buf = append(w.buf, w.NoColor...)
345	} else {
346		w.buf = append(w.buf, '{')
347	}
348	prevExist := false
349	for i, col := range t.columns {
350		k, _ := col.key.(string)
351		var m *node
352		for _, mm := range n.members {
353			if string(mm.key) == k {
354				m = mm
355				break
356			}
357		}
358		if prevExist {
359			w.buf = append(w.buf, comma...)
360			w.buf = append(w.buf, ' ')
361		}
362		if m == nil {
363			prevExist = false
364			pad := len(k) + 2 + col.size
365			if i < len(t.columns)-1 {
366				if w.SEN {
367					pad++
368				} else {
369					pad += 2
370				}
371			}
372			w.buf = append(w.buf, spaces[1:pad+1]...)
373		} else {
374			prevExist = true
375			w.buf = append(w.buf, k...)
376			w.buf = append(w.buf, ':')
377			w.buf = append(w.buf, ' ')
378			cw := col.size
379			switch m.kind {
380			case strNode:
381				w.buf = append(w.buf, m.buf...)
382				if m.size < cw {
383					w.buf = append(w.buf, spaces[1:cw-m.size+1]...)
384				}
385			case numNode:
386				if m.size < cw {
387					w.buf = append(w.buf, spaces[1:cw-m.size+1]...)
388				}
389				w.buf = append(w.buf, m.buf...)
390			case arrayNode:
391				w.alignArray(m, col, comma, []byte{' '})
392			case mapNode:
393				w.alignMap(m, col, comma, []byte{' '})
394			}
395		}
396	}
397	if w.Color {
398		w.buf = append(w.buf, w.SyntaxColor...)
399		w.buf = append(w.buf, '}')
400		w.buf = append(w.buf, w.NoColor...)
401	} else {
402		w.buf = append(w.buf, '}')
403	}
404}
405