1package tview
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"strconv"
8	"strings"
9)
10
11// The states of the ANSI escape code parser.
12const (
13	ansiText = iota
14	ansiEscape
15	ansiSubstring
16	ansiControlSequence
17)
18
19// ansi is a io.Writer which translates ANSI escape codes into tview color
20// tags.
21type ansi struct {
22	io.Writer
23
24	// Reusable buffers.
25	buffer                        *bytes.Buffer // The entire output text of one Write().
26	csiParameter, csiIntermediate *bytes.Buffer // Partial CSI strings.
27	attributes                    string        // The buffer's current text attributes (a tview attribute string).
28
29	// The current state of the parser. One of the ansi constants.
30	state int
31}
32
33// ANSIWriter returns an io.Writer which translates any ANSI escape codes
34// written to it into tview color tags. Other escape codes don't have an effect
35// and are simply removed. The translated text is written to the provided
36// writer.
37func ANSIWriter(writer io.Writer) io.Writer {
38	return &ansi{
39		Writer:          writer,
40		buffer:          new(bytes.Buffer),
41		csiParameter:    new(bytes.Buffer),
42		csiIntermediate: new(bytes.Buffer),
43		state:           ansiText,
44	}
45}
46
47// Write parses the given text as a string of runes, translates ANSI escape
48// codes to color tags and writes them to the output writer.
49func (a *ansi) Write(text []byte) (int, error) {
50	defer func() {
51		a.buffer.Reset()
52	}()
53
54	for _, r := range string(text) {
55		switch a.state {
56
57		// We just entered an escape sequence.
58		case ansiEscape:
59			switch r {
60			case '[': // Control Sequence Introducer.
61				a.csiParameter.Reset()
62				a.csiIntermediate.Reset()
63				a.state = ansiControlSequence
64			case 'c': // Reset.
65				fmt.Fprint(a.buffer, "[-:-:-]")
66				a.state = ansiText
67			case 'P', ']', 'X', '^', '_': // Substrings and commands.
68				a.state = ansiSubstring
69			default: // Ignore.
70				a.state = ansiText
71			}
72
73		// CSI Sequences.
74		case ansiControlSequence:
75			switch {
76			case r >= 0x30 && r <= 0x3f: // Parameter bytes.
77				if _, err := a.csiParameter.WriteRune(r); err != nil {
78					return 0, err
79				}
80			case r >= 0x20 && r <= 0x2f: // Intermediate bytes.
81				if _, err := a.csiIntermediate.WriteRune(r); err != nil {
82					return 0, err
83				}
84			case r >= 0x40 && r <= 0x7e: // Final byte.
85				switch r {
86				case 'E': // Next line.
87					count, _ := strconv.Atoi(a.csiParameter.String())
88					if count == 0 {
89						count = 1
90					}
91					fmt.Fprint(a.buffer, strings.Repeat("\n", count))
92				case 'm': // Select Graphic Rendition.
93					var background, foreground string
94					params := a.csiParameter.String()
95					fields := strings.Split(params, ";")
96					if len(params) == 0 || len(fields) == 1 && fields[0] == "0" {
97						// Reset.
98						a.attributes = ""
99						if _, err := a.buffer.WriteString("[-:-:-]"); err != nil {
100							return 0, err
101						}
102						break
103					}
104					lookupColor := func(colorNumber int) string {
105						if colorNumber < 0 || colorNumber > 15 {
106							return "black"
107						}
108						return []string{
109							"black",
110							"maroon",
111							"green",
112							"olive",
113							"navy",
114							"purple",
115							"teal",
116							"silver",
117							"gray",
118							"red",
119							"lime",
120							"yellow",
121							"blue",
122							"fuchsia",
123							"aqua",
124							"white",
125						}[colorNumber]
126					}
127				FieldLoop:
128					for index, field := range fields {
129						switch field {
130						case "1", "01":
131							if strings.IndexRune(a.attributes, 'b') < 0 {
132								a.attributes += "b"
133							}
134						case "2", "02":
135							if strings.IndexRune(a.attributes, 'd') < 0 {
136								a.attributes += "d"
137							}
138						case "4", "04":
139							if strings.IndexRune(a.attributes, 'u') < 0 {
140								a.attributes += "u"
141							}
142						case "5", "05":
143							if strings.IndexRune(a.attributes, 'l') < 0 {
144								a.attributes += "l"
145							}
146						case "22":
147							if i := strings.IndexRune(a.attributes, 'b'); i >= 0 {
148								a.attributes = a.attributes[:i] + a.attributes[i+1:]
149							}
150							if i := strings.IndexRune(a.attributes, 'd'); i >= 0 {
151								a.attributes = a.attributes[:i] + a.attributes[i+1:]
152							}
153						case "24":
154							if i := strings.IndexRune(a.attributes, 'u'); i >= 0 {
155								a.attributes = a.attributes[:i] + a.attributes[i+1:]
156							}
157						case "25":
158							if i := strings.IndexRune(a.attributes, 'l'); i >= 0 {
159								a.attributes = a.attributes[:i] + a.attributes[i+1:]
160							}
161						case "30", "31", "32", "33", "34", "35", "36", "37":
162							colorNumber, _ := strconv.Atoi(field)
163							foreground = lookupColor(colorNumber - 30)
164						case "39":
165							foreground = "-"
166						case "40", "41", "42", "43", "44", "45", "46", "47":
167							colorNumber, _ := strconv.Atoi(field)
168							background = lookupColor(colorNumber - 40)
169						case "49":
170							background = "-"
171						case "90", "91", "92", "93", "94", "95", "96", "97":
172							colorNumber, _ := strconv.Atoi(field)
173							foreground = lookupColor(colorNumber - 82)
174						case "100", "101", "102", "103", "104", "105", "106", "107":
175							colorNumber, _ := strconv.Atoi(field)
176							background = lookupColor(colorNumber - 92)
177						case "38", "48":
178							var color string
179							if len(fields) > index+1 {
180								if fields[index+1] == "5" && len(fields) > index+2 { // 8-bit colors.
181									colorNumber, _ := strconv.Atoi(fields[index+2])
182									if colorNumber <= 15 {
183										color = lookupColor(colorNumber)
184									} else if colorNumber <= 231 {
185										red := (colorNumber - 16) / 36
186										green := ((colorNumber - 16) / 6) % 6
187										blue := (colorNumber - 16) % 6
188										color = fmt.Sprintf("#%02x%02x%02x", 255*red/5, 255*green/5, 255*blue/5)
189									} else if colorNumber <= 255 {
190										grey := 255 * (colorNumber - 232) / 23
191										color = fmt.Sprintf("#%02x%02x%02x", grey, grey, grey)
192									}
193								} else if fields[index+1] == "2" && len(fields) > index+4 { // 24-bit colors.
194									red, _ := strconv.Atoi(fields[index+2])
195									green, _ := strconv.Atoi(fields[index+3])
196									blue, _ := strconv.Atoi(fields[index+4])
197									color = fmt.Sprintf("#%02x%02x%02x", red, green, blue)
198								}
199							}
200							if len(color) > 0 {
201								if field == "38" {
202									foreground = color
203								} else {
204									background = color
205								}
206							}
207							break FieldLoop
208						}
209					}
210					var colon string
211					if len(a.attributes) > 0 {
212						colon = ":"
213					}
214					if len(foreground) > 0 || len(background) > 0 || len(a.attributes) > 0 {
215						fmt.Fprintf(a.buffer, "[%s:%s%s%s]", foreground, background, colon, a.attributes)
216					}
217				}
218				a.state = ansiText
219			default: // Undefined byte.
220				a.state = ansiText // Abort CSI.
221			}
222
223			// We just entered a substring/command sequence.
224		case ansiSubstring:
225			if r == 27 { // Most likely the end of the substring.
226				a.state = ansiEscape
227			} // Ignore all other characters.
228
229			// "ansiText" and all others.
230		default:
231			if r == 27 {
232				// This is the start of an escape sequence.
233				a.state = ansiEscape
234			} else {
235				// Just a regular rune. Send to buffer.
236				if _, err := a.buffer.WriteRune(r); err != nil {
237					return 0, err
238				}
239			}
240		}
241	}
242
243	// Write buffer to target writer.
244	n, err := a.buffer.WriteTo(a.Writer)
245	if err != nil {
246		return int(n), err
247	}
248	return len(text), nil
249}
250
251// TranslateANSI replaces ANSI escape sequences found in the provided string
252// with tview's color tags and returns the resulting string.
253func TranslateANSI(text string) string {
254	var buffer bytes.Buffer
255	writer := ANSIWriter(&buffer)
256	writer.Write([]byte(text))
257	return buffer.String()
258}
259