1package cellbuf
 2
 3import (
 4	"io"
 5
 6	"github.com/charmbracelet/x/ansi"
 7)
 8
 9// PenWriter is a writer that writes to a buffer and keeps track of the current
10// pen style and link state for the purpose of wrapping with newlines.
11type PenWriter struct {
12	w     io.Writer
13	p     *ansi.Parser
14	style Style
15	link  Link
16}
17
18// NewPenWriter returns a new PenWriter.
19func NewPenWriter(w io.Writer) *PenWriter {
20	pw := &PenWriter{w: w}
21	pw.p = ansi.NewParser()
22	pw.p.SetParamsSize(32)            // 32 parameters
23	pw.p.SetDataSize(4 * 1024 * 1024) // 4MB of data buffer
24	handleCsi := func(cmd ansi.Cmd, params ansi.Params) {
25		if cmd == 'm' {
26			ReadStyle(params, &pw.style)
27		}
28	}
29	handleOsc := func(cmd int, data []byte) {
30		if cmd == 8 {
31			ReadLink(data, &pw.link)
32		}
33	}
34	pw.p.SetHandler(ansi.Handler{
35		HandleCsi: handleCsi,
36		HandleOsc: handleOsc,
37	})
38	return pw
39}
40
41// Style returns the current pen style.
42func (w *PenWriter) Style() Style {
43	return w.style
44}
45
46// Link returns the current pen link.
47func (w *PenWriter) Link() Link {
48	return w.link
49}
50
51// Write writes to the buffer.
52func (w *PenWriter) Write(p []byte) (int, error) {
53	for i := 0; i < len(p); i++ {
54		b := p[i]
55		w.p.Advance(b)
56		if b == '\n' {
57			if !w.style.Empty() {
58				_, _ = w.w.Write([]byte(ansi.ResetStyle))
59			}
60			if !w.link.Empty() {
61				_, _ = w.w.Write([]byte(ansi.ResetHyperlink()))
62			}
63		}
64
65		_, _ = w.w.Write([]byte{b})
66		if b == '\n' {
67			if !w.link.Empty() {
68				_, _ = w.w.Write([]byte(ansi.SetHyperlink(w.link.URL, w.link.Params)))
69			}
70			if !w.style.Empty() {
71				_, _ = w.w.Write([]byte(w.style.Sequence()))
72			}
73		}
74	}
75
76	return len(p), nil
77}
78
79// Close closes the writer and resets the style and link if necessary.
80func (w *PenWriter) Close() error {
81	if !w.style.Empty() {
82		_, _ = w.w.Write([]byte(ansi.ResetStyle))
83	}
84	if !w.link.Empty() {
85		_, _ = w.w.Write([]byte(ansi.ResetHyperlink()))
86	}
87	return nil
88}