1
package readline
2

3
import (
4
	"bufio"
5
	"bytes"
6
	"io"
7
	"strconv"
8
	"strings"
9
	"sync"
10
)
11

12
type runeBufferBck struct {
13
	buf []rune
14
	idx int
15
}
16

17
type RuneBuffer struct {
18
	buf    []rune
19
	idx    int
20
	prompt []rune
21
	w      io.Writer
22

23
	hadClean    bool
24
	interactive bool
25
	cfg         *Config
26

27
	width int
28

29
	bck *runeBufferBck
30

31
	offset string
32

33
	lastKill []rune
34

35
	sync.Mutex
36
}
37

38
func (r *RuneBuffer) pushKill(text []rune) {
39
	r.lastKill = append([]rune{}, text...)
40
}
41

42
func (r *RuneBuffer) OnWidthChange(newWidth int) {
43
	r.Lock()
44
	r.width = newWidth
45
	r.Unlock()
46
}
47

48
func (r *RuneBuffer) Backup() {
49
	r.Lock()
50
	r.bck = &runeBufferBck{r.buf, r.idx}
51
	r.Unlock()
52
}
53

54
func (r *RuneBuffer) Restore() {
55
	r.Refresh(func() {
56
		if r.bck == nil {
57
			return
58
		}
59
		r.buf = r.bck.buf
60
		r.idx = r.bck.idx
61
	})
62
}
63

64
func NewRuneBuffer(w io.Writer, prompt string, cfg *Config, width int) *RuneBuffer {
65
	rb := &RuneBuffer{
66
		w:           w,
67
		interactive: cfg.useInteractive(),
68
		cfg:         cfg,
69
		width:       width,
70
	}
71
	rb.SetPrompt(prompt)
72
	return rb
73
}
74

75
func (r *RuneBuffer) SetConfig(cfg *Config) {
76
	r.Lock()
77
	r.cfg = cfg
78
	r.interactive = cfg.useInteractive()
79
	r.Unlock()
80
}
81

82
func (r *RuneBuffer) SetMask(m rune) {
83
	r.Lock()
84
	r.cfg.MaskRune = m
85
	r.Unlock()
86
}
87

88
func (r *RuneBuffer) CurrentWidth(x int) int {
89
	r.Lock()
90
	defer r.Unlock()
91
	return runes.WidthAll(r.buf[:x])
92
}
93

94
func (r *RuneBuffer) PromptLen() int {
95
	r.Lock()
96
	width := r.promptLen()
97
	r.Unlock()
98
	return width
99
}
100

101
func (r *RuneBuffer) promptLen() int {
102
	return runes.WidthAll(runes.ColorFilter(r.prompt))
103
}
104

105
func (r *RuneBuffer) RuneSlice(i int) []rune {
106
	r.Lock()
107
	defer r.Unlock()
108

109
	if i > 0 {
110
		rs := make([]rune, i)
111
		copy(rs, r.buf[r.idx:r.idx+i])
112
		return rs
113
	}
114
	rs := make([]rune, -i)
115
	copy(rs, r.buf[r.idx+i:r.idx])
116
	return rs
117
}
118

119
func (r *RuneBuffer) Runes() []rune {
120
	r.Lock()
121
	newr := make([]rune, len(r.buf))
122
	copy(newr, r.buf)
123
	r.Unlock()
124
	return newr
125
}
126

127
func (r *RuneBuffer) Pos() int {
128
	r.Lock()
129
	defer r.Unlock()
130
	return r.idx
131
}
132

133
func (r *RuneBuffer) Len() int {
134
	r.Lock()
135
	defer r.Unlock()
136
	return len(r.buf)
137
}
138

139
func (r *RuneBuffer) MoveToLineStart() {
140
	r.Refresh(func() {
141
		if r.idx == 0 {
142
			return
143
		}
144
		r.idx = 0
145
	})
146
}
147

148
func (r *RuneBuffer) MoveBackward() {
149
	r.Refresh(func() {
150
		if r.idx == 0 {
151
			return
152
		}
153
		r.idx--
154
	})
155
}
156

157
func (r *RuneBuffer) WriteString(s string) {
158
	r.WriteRunes([]rune(s))
159
}
160

161
func (r *RuneBuffer) WriteRune(s rune) {
162
	r.WriteRunes([]rune{s})
163
}
164

165
func (r *RuneBuffer) WriteRunes(s []rune) {
166
	r.Refresh(func() {
167
		tail := append(s, r.buf[r.idx:]...)
168
		r.buf = append(r.buf[:r.idx], tail...)
169
		r.idx += len(s)
170
	})
171
}
172

173
func (r *RuneBuffer) MoveForward() {
174
	r.Refresh(func() {
175
		if r.idx == len(r.buf) {
176
			return
177
		}
178
		r.idx++
179
	})
180
}
181

182
func (r *RuneBuffer) IsCursorInEnd() bool {
183
	r.Lock()
184
	defer r.Unlock()
185
	return r.idx == len(r.buf)
186
}
187

188
func (r *RuneBuffer) Replace(ch rune) {
189
	r.Refresh(func() {
190
		r.buf[r.idx] = ch
191
	})
192
}
193

194
func (r *RuneBuffer) Erase() {
195
	r.Refresh(func() {
196
		r.idx = 0
197
		r.pushKill(r.buf[:])
198
		r.buf = r.buf[:0]
199
	})
200
}
201

202
func (r *RuneBuffer) Delete() (success bool) {
203
	r.Refresh(func() {
204
		if r.idx == len(r.buf) {
205
			return
206
		}
207
		r.pushKill(r.buf[r.idx : r.idx+1])
208
		r.buf = append(r.buf[:r.idx], r.buf[r.idx+1:]...)
209
		success = true
210
	})
211
	return
212
}
213

214
func (r *RuneBuffer) DeleteWord() {
215
	if r.idx == len(r.buf) {
216
		return
217
	}
218
	init := r.idx
219
	for init < len(r.buf) && IsWordBreak(r.buf[init]) {
220
		init++
221
	}
222
	for i := init + 1; i < len(r.buf); i++ {
223
		if !IsWordBreak(r.buf[i]) && IsWordBreak(r.buf[i-1]) {
224
			r.pushKill(r.buf[r.idx : i-1])
225
			r.Refresh(func() {
226
				r.buf = append(r.buf[:r.idx], r.buf[i-1:]...)
227
			})
228
			return
229
		}
230
	}
231
	r.Kill()
232
}
233

234
func (r *RuneBuffer) MoveToPrevWord() (success bool) {
235
	r.Refresh(func() {
236
		if r.idx == 0 {
237
			return
238
		}
239

240
		for i := r.idx - 1; i > 0; i-- {
241
			if !IsWordBreak(r.buf[i]) && IsWordBreak(r.buf[i-1]) {
242
				r.idx = i
243
				success = true
244
				return
245
			}
246
		}
247
		r.idx = 0
248
		success = true
249
	})
250
	return
251
}
252

253
func (r *RuneBuffer) KillFront() {
254
	r.Refresh(func() {
255
		if r.idx == 0 {
256
			return
257
		}
258

259
		length := len(r.buf) - r.idx
260
		r.pushKill(r.buf[:r.idx])
261
		copy(r.buf[:length], r.buf[r.idx:])
262
		r.idx = 0
263
		r.buf = r.buf[:length]
264
	})
265
}
266

267
func (r *RuneBuffer) Kill() {
268
	r.Refresh(func() {
269
		r.pushKill(r.buf[r.idx:])
270
		r.buf = r.buf[:r.idx]
271
	})
272
}
273

274
func (r *RuneBuffer) Transpose() {
275
	r.Refresh(func() {
276
		if len(r.buf) == 1 {
277
			r.idx++
278
		}
279

280
		if len(r.buf) < 2 {
281
			return
282
		}
283

284
		if r.idx == 0 {
285
			r.idx = 1
286
		} else if r.idx >= len(r.buf) {
287
			r.idx = len(r.buf) - 1
288
		}
289
		r.buf[r.idx], r.buf[r.idx-1] = r.buf[r.idx-1], r.buf[r.idx]
290
		r.idx++
291
	})
292
}
293

294
func (r *RuneBuffer) MoveToNextWord() {
295
	r.Refresh(func() {
296
		for i := r.idx + 1; i < len(r.buf); i++ {
297
			if !IsWordBreak(r.buf[i]) && IsWordBreak(r.buf[i-1]) {
298
				r.idx = i
299
				return
300
			}
301
		}
302

303
		r.idx = len(r.buf)
304
	})
305
}
306

307
func (r *RuneBuffer) MoveToEndWord() {
308
	r.Refresh(func() {
309
		// already at the end, so do nothing
310
		if r.idx == len(r.buf) {
311
			return
312
		}
313
		// if we are at the end of a word already, go to next
314
		if !IsWordBreak(r.buf[r.idx]) && IsWordBreak(r.buf[r.idx+1]) {
315
			r.idx++
316
		}
317

318
		// keep going until at the end of a word
319
		for i := r.idx + 1; i < len(r.buf); i++ {
320
			if IsWordBreak(r.buf[i]) && !IsWordBreak(r.buf[i-1]) {
321
				r.idx = i - 1
322
				return
323
			}
324
		}
325
		r.idx = len(r.buf)
326
	})
327
}
328

329
func (r *RuneBuffer) BackEscapeWord() {
330
	r.Refresh(func() {
331
		if r.idx == 0 {
332
			return
333
		}
334
		for i := r.idx - 1; i > 0; i-- {
335
			if !IsWordBreak(r.buf[i]) && IsWordBreak(r.buf[i-1]) {
336
				r.pushKill(r.buf[i:r.idx])
337
				r.buf = append(r.buf[:i], r.buf[r.idx:]...)
338
				r.idx = i
339
				return
340
			}
341
		}
342

343
		r.buf = r.buf[:0]
344
		r.idx = 0
345
	})
346
}
347

348
func (r *RuneBuffer) Yank() {
349
	if len(r.lastKill) == 0 {
350
		return
351
	}
352
	r.Refresh(func() {
353
		buf := make([]rune, 0, len(r.buf)+len(r.lastKill))
354
		buf = append(buf, r.buf[:r.idx]...)
355
		buf = append(buf, r.lastKill...)
356
		buf = append(buf, r.buf[r.idx:]...)
357
		r.buf = buf
358
		r.idx += len(r.lastKill)
359
	})
360
}
361

362
func (r *RuneBuffer) Backspace() {
363
	r.Refresh(func() {
364
		if r.idx == 0 {
365
			return
366
		}
367

368
		r.idx--
369
		r.buf = append(r.buf[:r.idx], r.buf[r.idx+1:]...)
370
	})
371
}
372

373
func (r *RuneBuffer) MoveToLineEnd() {
374
	r.Refresh(func() {
375
		if r.idx == len(r.buf) {
376
			return
377
		}
378

379
		r.idx = len(r.buf)
380
	})
381
}
382

383
func (r *RuneBuffer) LineCount(width int) int {
384
	if width == -1 {
385
		width = r.width
386
	}
387
	return LineCount(width,
388
		runes.WidthAll(r.buf)+r.PromptLen())
389
}
390

391
func (r *RuneBuffer) MoveTo(ch rune, prevChar, reverse bool) (success bool) {
392
	r.Refresh(func() {
393
		if reverse {
394
			for i := r.idx - 1; i >= 0; i-- {
395
				if r.buf[i] == ch {
396
					r.idx = i
397
					if prevChar {
398
						r.idx++
399
					}
400
					success = true
401
					return
402
				}
403
			}
404
			return
405
		}
406
		for i := r.idx + 1; i < len(r.buf); i++ {
407
			if r.buf[i] == ch {
408
				r.idx = i
409
				if prevChar {
410
					r.idx--
411
				}
412
				success = true
413
				return
414
			}
415
		}
416
	})
417
	return
418
}
419

420
func (r *RuneBuffer) isInLineEdge() bool {
421
	if isWindows {
422
		return false
423
	}
424
	sp := r.getSplitByLine(r.buf)
425
	return len(sp[len(sp)-1]) == 0
426
}
427

428
func (r *RuneBuffer) getSplitByLine(rs []rune) []string {
429
	return SplitByLine(r.promptLen(), r.width, rs)
430
}
431

432
func (r *RuneBuffer) IdxLine(width int) int {
433
	r.Lock()
434
	defer r.Unlock()
435
	return r.idxLine(width)
436
}
437

438
func (r *RuneBuffer) idxLine(width int) int {
439
	if width == 0 {
440
		return 0
441
	}
442
	sp := r.getSplitByLine(r.buf[:r.idx])
443
	return len(sp) - 1
444
}
445

446
func (r *RuneBuffer) CursorLineCount() int {
447
	return r.LineCount(r.width) - r.IdxLine(r.width)
448
}
449

450
func (r *RuneBuffer) Refresh(f func()) {
451
	r.Lock()
452
	defer r.Unlock()
453

454
	if !r.interactive {
455
		if f != nil {
456
			f()
457
		}
458
		return
459
	}
460

461
	r.clean()
462
	if f != nil {
463
		f()
464
	}
465
	r.print()
466
}
467

468
func (r *RuneBuffer) SetOffset(offset string) {
469
	r.Lock()
470
	r.offset = offset
471
	r.Unlock()
472
}
473

474
func (r *RuneBuffer) print() {
475
	r.w.Write(r.output())
476
	r.hadClean = false
477
}
478

479
func (r *RuneBuffer) output() []byte {
480
	buf := bytes.NewBuffer(nil)
481
	buf.WriteString(string(r.prompt))
482
	if r.cfg.EnableMask && len(r.buf) > 0 {
483
		buf.Write([]byte(strings.Repeat(string(r.cfg.MaskRune), len(r.buf)-1)))
484
		if r.buf[len(r.buf)-1] == '\n' {
485
			buf.Write([]byte{'\n'})
486
		} else {
487
			buf.Write([]byte(string(r.cfg.MaskRune)))
488
		}
489
		if len(r.buf) > r.idx {
490
			buf.Write(r.getBackspaceSequence())
491
		}
492

493
	} else {
494
		for _, e := range r.cfg.Painter.Paint(r.buf, r.idx) {
495
			if e == '\t' {
496
				buf.WriteString(strings.Repeat(" ", TabWidth))
497
			} else {
498
				buf.WriteRune(e)
499
			}
500
		}
501
		if r.isInLineEdge() {
502
			buf.Write([]byte(" \b"))
503
		}
504
	}
505
	// cursor position
506
	if len(r.buf) > r.idx {
507
		buf.Write(r.getBackspaceSequence())
508
	}
509
	return buf.Bytes()
510
}
511

512
func (r *RuneBuffer) getBackspaceSequence() []byte {
513
	var sep = map[int]bool{}
514

515
	var i int
516
	for {
517
		if i >= runes.WidthAll(r.buf) {
518
			break
519
		}
520

521
		if i == 0 {
522
			i -= r.promptLen()
523
		}
524
		i += r.width
525

526
		sep[i] = true
527
	}
528
	var buf []byte
529
	for i := len(r.buf); i > r.idx; i-- {
530
		// move input to the left of one
531
		buf = append(buf, '\b')
532
		if sep[i] {
533
			// up one line, go to the start of the line and move cursor right to the end (r.width)
534
			buf = append(buf, "\033[A\r"+"\033["+strconv.Itoa(r.width)+"C"...)
535
		}
536
	}
537

538
	return buf
539

540
}
541

542
func (r *RuneBuffer) Reset() []rune {
543
	ret := runes.Copy(r.buf)
544
	r.buf = r.buf[:0]
545
	r.idx = 0
546
	return ret
547
}
548

549
func (r *RuneBuffer) calWidth(m int) int {
550
	if m > 0 {
551
		return runes.WidthAll(r.buf[r.idx : r.idx+m])
552
	}
553
	return runes.WidthAll(r.buf[r.idx+m : r.idx])
554
}
555

556
func (r *RuneBuffer) SetStyle(start, end int, style string) {
557
	if end < start {
558
		panic("end < start")
559
	}
560

561
	// goto start
562
	move := start - r.idx
563
	if move > 0 {
564
		r.w.Write([]byte(string(r.buf[r.idx : r.idx+move])))
565
	} else {
566
		r.w.Write(bytes.Repeat([]byte("\b"), r.calWidth(move)))
567
	}
568
	r.w.Write([]byte("\033[" + style + "m"))
569
	r.w.Write([]byte(string(r.buf[start:end])))
570
	r.w.Write([]byte("\033[0m"))
571
	// TODO: move back
572
}
573

574
func (r *RuneBuffer) SetWithIdx(idx int, buf []rune) {
575
	r.Refresh(func() {
576
		r.buf = buf
577
		r.idx = idx
578
	})
579
}
580

581
func (r *RuneBuffer) Set(buf []rune) {
582
	r.SetWithIdx(len(buf), buf)
583
}
584

585
func (r *RuneBuffer) SetPrompt(prompt string) {
586
	r.Lock()
587
	r.prompt = []rune(prompt)
588
	r.Unlock()
589
}
590

591
func (r *RuneBuffer) cleanOutput(w io.Writer, idxLine int) {
592
	buf := bufio.NewWriter(w)
593

594
	if r.width == 0 {
595
		buf.WriteString(strings.Repeat("\r\b", len(r.buf)+r.promptLen()))
596
		buf.Write([]byte("\033[J"))
597
	} else {
598
		buf.Write([]byte("\033[J")) // just like ^k :)
599
		if idxLine == 0 {
600
			buf.WriteString("\033[2K")
601
			buf.WriteString("\r")
602
		} else {
603
			for i := 0; i < idxLine; i++ {
604
				io.WriteString(buf, "\033[2K\r\033[A")
605
			}
606
			io.WriteString(buf, "\033[2K\r")
607
		}
608
	}
609
	buf.Flush()
610
	return
611
}
612

613
func (r *RuneBuffer) Clean() {
614
	r.Lock()
615
	r.clean()
616
	r.Unlock()
617
}
618

619
func (r *RuneBuffer) clean() {
620
	r.cleanWithIdxLine(r.idxLine(r.width))
621
}
622

623
func (r *RuneBuffer) cleanWithIdxLine(idxLine int) {
624
	if r.hadClean || !r.interactive {
625
		return
626
	}
627
	r.hadClean = true
628
	r.cleanOutput(r.w, idxLine)
629
}
630

631