1package ansi
 2
 3import (
 4	"bytes"
 5	"io"
 6	"unicode/utf8"
 7)
 8
 9type Writer struct {
10	Forward io.Writer
11
12	ansi       bool
13	ansiseq    bytes.Buffer
14	lastseq    bytes.Buffer
15	seqchanged bool
16	runeBuf    []byte
17}
18
19// Write is used to write content to the ANSI buffer.
20func (w *Writer) Write(b []byte) (int, error) {
21	for _, c := range string(b) {
22		if c == Marker {
23			// ANSI escape sequence
24			w.ansi = true
25			w.seqchanged = true
26			_, _ = w.ansiseq.WriteRune(c)
27		} else if w.ansi {
28			_, _ = w.ansiseq.WriteRune(c)
29			if IsTerminator(c) {
30				// ANSI sequence terminated
31				w.ansi = false
32
33				if bytes.HasSuffix(w.ansiseq.Bytes(), []byte("[0m")) {
34					// reset sequence
35					w.lastseq.Reset()
36					w.seqchanged = false
37				} else if c == 'm' {
38					// color code
39					_, _ = w.lastseq.Write(w.ansiseq.Bytes())
40				}
41
42				_, _ = w.ansiseq.WriteTo(w.Forward)
43			}
44		} else {
45			_, err := w.writeRune(c)
46			if err != nil {
47				return 0, err
48			}
49		}
50	}
51
52	return len(b), nil
53}
54
55func (w *Writer) writeRune(r rune) (int, error) {
56	if w.runeBuf == nil {
57		w.runeBuf = make([]byte, utf8.UTFMax)
58	}
59	n := utf8.EncodeRune(w.runeBuf, r)
60	return w.Forward.Write(w.runeBuf[:n])
61}
62
63func (w *Writer) LastSequence() string {
64	return w.lastseq.String()
65}
66
67func (w *Writer) ResetAnsi() {
68	if !w.seqchanged {
69		return
70	}
71	_, _ = w.Forward.Write([]byte("\x1b[0m"))
72}
73
74func (w *Writer) RestoreAnsi() {
75	_, _ = w.Forward.Write(w.lastseq.Bytes())
76}