stario/io.go

package stario

import (
	"bufio"
	"fmt"
	"golang.org/x/term"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
)

type InputMsg struct {
	msg             string
	err             error
	skipSliceSigErr bool
}

type rawInputSignalMode uint8

const (
	rawInputSignalIgnore rawInputSignalMode = iota
	rawInputSignalExit
	rawInputSignalReturnError
)

type rawTerminalSession struct {
	fd           int
	state        *term.State
	reader       *bufio.Reader
	input        io.Closer
	redrawHint   string
	printNewline bool
	mu           sync.Mutex
}

var rawTerminalSessionFactory = newRawTerminalSession
var inputSignalHandler = signal

// Passwd reads one password-style line in raw mode.
//
// When the user presses an input signal such as Ctrl+C, this compatibility
// entry exits the current flow and returns an empty message with a nil error.
func Passwd(hint string, defaultVal string) InputMsg {
	return passwd(hint, defaultVal, "", rawInputSignalExit)
}

// PasswdWithMask is like Passwd but echoes the provided mask string.
func PasswdWithMask(hint string, defaultVal string, mask string) InputMsg {
	return passwd(hint, defaultVal, mask, rawInputSignalExit)
}

// PasswdResponseSignal is like Passwd but preserves input-signal errors for
// callers that need to distinguish Ctrl+C / Ctrl+Z style exits.
func PasswdResponseSignal(hint string, defaultVal string) InputMsg {
	return passwd(hint, defaultVal, "", rawInputSignalReturnError)
}

// PasswdResponseSignalWithMask is like PasswdResponseSignal but echoes the
// provided mask string.
func PasswdResponseSignalWithMask(hint string, defaultVal string, mask string) InputMsg {
	return passwd(hint, defaultVal, mask, rawInputSignalReturnError)
}

// MessageBoxRaw reads one line in raw mode without treating control keys as
// exit signals.
func MessageBoxRaw(hint string, defaultVal string) InputMsg {
	return messageBox(hint, defaultVal)
}

func newRawTerminalSession(hint string, printNewline bool) (*rawTerminalSession, error) {
	if hint != "" {
		fmt.Print(hint)
	}
	input, err := openRawTerminalInput()
	if err != nil {
		return nil, err
	}
	fd := int(input.Fd())
	state, err := term.MakeRaw(fd)
	if err != nil {
		_ = input.Close()
		return nil, err
	}
	return &rawTerminalSession{
		fd:           fd,
		state:        state,
		reader:       bufio.NewReader(input),
		input:        input,
		redrawHint:   promptRedrawHint(hint),
		printNewline: printNewline,
	}, nil
}

func (session *rawTerminalSession) Close() {
	if session == nil {
		return
	}
	session.mu.Lock()
	defer session.mu.Unlock()
	if session.state != nil {
		_ = term.Restore(session.fd, session.state)
		session.state = nil
	}
	if session.printNewline {
		fmt.Println()
		session.printNewline = false
	}
	if session.input != nil {
		_ = session.input.Close()
		session.input = nil
	}
}

func (session *rawTerminalSession) Restore() error {
	if session == nil {
		return nil
	}
	session.mu.Lock()
	defer session.mu.Unlock()
	if session.state == nil {
		return nil
	}
	if err := term.Restore(session.fd, session.state); err != nil {
		return err
	}
	session.state = nil
	return nil
}

func (session *rawTerminalSession) Abort() {
	session.Close()
}

func promptRedrawHint(hint string) string {
	if strings.Index(hint, "\n") >= 0 {
		hint = hint[strings.LastIndex(hint, "\n")+1:]
	}
	return strings.TrimSpace(hint)
}

func finalizeInputValue(raw string, defaultVal string) string {
	raw = strings.TrimSpace(raw)
	if len(raw) == 0 {
		return defaultVal
	}
	return raw
}

func renderRawEcho(ioBuf []rune, mask string) string {
	if mask == "" {
		return string(ioBuf)
	}
	return strings.Repeat(mask, len(ioBuf))
}

func rawEchoRenderUnit(r rune, mask string, maskWidth int) (string, int, bool) {
	if mask != "" {
		if maskWidth <= 0 {
			return "", 0, false
		}
		return mask, maskWidth, true
	}
	width := runeDisplayWidth(r)
	if width <= 0 {
		return "", 0, false
	}
	return string(r), width, true
}

func redrawPromptLine(hint string, echo string, lastWidth int) int {
	nowWidth := stringDisplayWidth(hint) + stringDisplayWidth(echo)
	clearWidth := lastWidth
	if nowWidth > clearWidth {
		clearWidth = nowWidth
	}
	fmt.Print("\r")
	if clearWidth > 0 {
		fmt.Print(strings.Repeat(" ", clearWidth))
		fmt.Print("\r")
	}
	if hint != "" {
		fmt.Print(hint)
	}
	if echo != "" {
		fmt.Print(echo)
	}
	return nowWidth
}

func erasePromptTail(width int) {
	if width <= 0 {
		return
	}
	backtrack := strings.Repeat("\b", width)
	fmt.Print(backtrack)
	fmt.Print(strings.Repeat(" ", width))
	fmt.Print(backtrack)
}

func redrawPromptEcho(hint string, ioBuf []rune, mask string, lastWidth int) (int, int) {
	echo := renderRawEcho(ioBuf, mask)
	echoWidth := stringDisplayWidth(echo)
	return echoWidth, redrawPromptLine(hint, echo, lastWidth)
}

func stringDisplayWidth(text string) int {
	width := 0
	for _, r := range text {
		width += runeDisplayWidth(r)
	}
	return width
}

func runeDisplayWidth(r rune) int {
	switch {
	case r == '\t':
		return 4
	case r < 0x20 || (r >= 0x7f && r < 0xa0):
		return 0
	case isWideRune(r):
		return 2
	default:
		return 1
	}
}

func isWideRune(r rune) bool {
	return r >= 0x1100 && (r <= 0x115f ||
		r == 0x2329 || r == 0x232a ||
		(r >= 0x2e80 && r <= 0xa4cf && r != 0x303f) ||
		(r >= 0xac00 && r <= 0xd7a3) ||
		(r >= 0xf900 && r <= 0xfaff) ||
		(r >= 0xfe10 && r <= 0xfe19) ||
		(r >= 0xfe30 && r <= 0xfe6f) ||
		(r >= 0xff00 && r <= 0xff60) ||
		(r >= 0xffe0 && r <= 0xffe6) ||
		(r >= 0x1f300 && r <= 0x1f64f) ||
		(r >= 0x1f900 && r <= 0x1f9ff) ||
		(r >= 0x20000 && r <= 0x3fffd))
}

func signalInputResult(mode rawInputSignalMode, err error) InputMsg {
	switch mode {
	case rawInputSignalExit:
		return InputMsg{msg: "", err: nil}
	case rawInputSignalReturnError:
		return InputMsg{msg: "", err: err}
	default:
		return InputMsg{msg: "", err: nil}
	}
}

func rawLineInput(hint string, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
	session, err := rawTerminalSessionFactory(hint, true)
	if err != nil {
		return InputMsg{msg: "", err: err}
	}
	defer session.Close()
	return rawLineInputSession(session, defaultVal, mask, signalMode)
}

func rawLineInputSession(session *rawTerminalSession, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
	if session == nil || session.reader == nil {
		return InputMsg{msg: "", err: io.ErrClosedPipe}
	}
	ioBuf := make([]rune, 0, 16)
	promptWidth := stringDisplayWidth(session.redrawHint)
	maskWidth := stringDisplayWidth(mask)
	echoWidth := 0
	lastWidth := promptWidth
	for {
		b, _, err := session.reader.ReadRune()
		if err != nil {
			return InputMsg{msg: "", err: err}
		}
		if signalMode != rawInputSignalIgnore && isSignal(b) {
			session.Close()
			if signalMode == rawInputSignalExit {
				return signalInputResult(signalMode, nil)
			}
			return signalInputResult(signalMode, inputSignalHandler(b))
		}
		switch b {
		case 0x0d, 0x0a:
			return InputMsg{msg: finalizeInputValue(string(ioBuf), defaultVal), err: nil}
		case 0x08, 0x7F:
			if len(ioBuf) > 0 {
				removed := ioBuf[len(ioBuf)-1]
				ioBuf = ioBuf[:len(ioBuf)-1]
				if _, removedWidth, ok := rawEchoRenderUnit(removed, mask, maskWidth); ok {
					erasePromptTail(removedWidth)
					echoWidth -= removedWidth
					if echoWidth < 0 {
						echoWidth = 0
					}
					lastWidth = promptWidth + echoWidth
					continue
				}
			}
		default:
			ioBuf = append(ioBuf, b)
			if appendText, appendWidth, ok := rawEchoRenderUnit(b, mask, maskWidth); ok {
				fmt.Print(appendText)
				echoWidth += appendWidth
				lastWidth = promptWidth + echoWidth
				continue
			}
		}
		echoWidth, lastWidth = redrawPromptEcho(session.redrawHint, ioBuf, mask, lastWidth)
	}
}

func messageBox(hint string, defaultVal string) InputMsg {
	return rawLineInput(hint, defaultVal, "", rawInputSignalIgnore)
}

func isSignal(s rune) bool {
	switch s {
	case 0x03, 0x1a, 0x1c:
		return true
	default:
		return false
	}
}

func passwd(hint string, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
	return rawLineInput(hint, defaultVal, mask, signalMode)
}

// MessageBox reads one line in cooked mode and falls back to defaultVal when
// the trimmed input is empty.
func MessageBox(hint string, defaultVal string) InputMsg {
	if hint != "" {
		fmt.Print(hint)
	}
	inputReader := bufio.NewReader(os.Stdin)
	str, err := inputReader.ReadString('\n')
	if err != nil {
		return InputMsg{msg: str, err: err}
	}
	str = finalizeInputValue(str, defaultVal)
	return InputMsg{msg: str, err: err}
}

func (im *InputMsg) IgnoreSliceParseError(i bool) *InputMsg {
	im.skipSliceSigErr = i
	return im
}

func (im InputMsg) String() (string, error) {
	if im.err != nil {
		return "", im.err
	}
	return im.msg, nil
}

func (im InputMsg) MustString() string {
	res, _ := im.String()
	return res
}

func (im InputMsg) SliceString(sep string) ([]string, error) {
	if im.err != nil {
		return nil, im.err
	}
	if len(strings.TrimSpace(im.msg)) == 0 {
		return []string{}, nil
	}
	return strings.Split(im.msg, sep), nil
}

func (im InputMsg) MustSliceString(sep string) []string {
	res, _ := im.SliceString(sep)
	return res
}

func (im InputMsg) sliceFn(sep string, fn func(string) (interface{}, error)) ([]interface{}, error) {
	var res []interface{}
	data, err := im.SliceString(sep)
	if err != nil {
		return res, err
	}
	for _, v := range data {
		code, err := fn(strings.TrimSpace(v))
		if err != nil && !im.skipSliceSigErr {
			return nil, err
		} else if err == nil {
			res = append(res, code)
		}
	}
	return res, nil
}

func (im InputMsg) Int() (int, error) {
	if im.err != nil {
		return 0, im.err
	}
	return strconv.Atoi(im.msg)
}

func (im InputMsg) SliceInt(sep string) ([]int, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		return strconv.Atoi(v)
	})
	var res []int
	for _, v := range data {
		res = append(res, v.(int))
	}
	return res, err
}

func (im InputMsg) MustSliceInt(sep string) []int {
	res, _ := im.SliceInt(sep)
	return res
}

func (im InputMsg) MustInt() int {
	res, _ := im.Int()
	return res
}

func (im InputMsg) Int64() (int64, error) {
	if im.err != nil {
		return 0, im.err
	}
	return strconv.ParseInt(im.msg, 10, 64)
}

func (im InputMsg) MustInt64() int64 {
	res, _ := im.Int64()
	return res
}

func (im InputMsg) SliceInt64(sep string) ([]int64, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		return strconv.ParseInt(v, 10, 64)
	})
	var res []int64
	for _, v := range data {
		res = append(res, v.(int64))
	}
	return res, err
}

func (im InputMsg) MustSliceInt64(sep string) []int64 {
	res, _ := im.SliceInt64(sep)
	return res
}

func (im InputMsg) Uint64() (uint64, error) {
	if im.err != nil {
		return 0, im.err
	}
	return strconv.ParseUint(im.msg, 10, 64)
}

func (im InputMsg) MustUint64() uint64 {
	res, _ := im.Uint64()
	return res
}
func (im InputMsg) SliceUint64(sep string) ([]uint64, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		return strconv.ParseUint(v, 10, 64)
	})
	var res []uint64
	for _, v := range data {
		res = append(res, v.(uint64))
	}
	return res, err
}

func (im InputMsg) MustSliceUint64(sep string) []uint64 {
	res, _ := im.SliceUint64(sep)
	return res
}

func (im InputMsg) Bool() (bool, error) {
	if im.err != nil {
		return false, im.err
	}
	return strconv.ParseBool(im.msg)
}

func (im InputMsg) MustBool() bool {
	res, _ := im.Bool()
	return res
}

func (im InputMsg) SliceBool(sep string) ([]bool, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		return strconv.ParseBool(v)
	})
	var res []bool
	for _, v := range data {
		res = append(res, v.(bool))
	}
	return res, err
}

func (im InputMsg) MustSliceBool(sep string) []bool {
	res, _ := im.SliceBool(sep)
	return res
}

func (im InputMsg) Float64() (float64, error) {
	if im.err != nil {
		return 0, im.err
	}
	return strconv.ParseFloat(im.msg, 64)
}

func (im InputMsg) MustFloat64() float64 {
	res, _ := im.Float64()
	return res
}

func (im InputMsg) SliceFloat64(sep string) ([]float64, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		return strconv.ParseFloat(v, 64)
	})
	var res []float64
	for _, v := range data {
		res = append(res, v.(float64))
	}
	return res, err
}

func (im InputMsg) MustSliceFloat64(sep string) []float64 {
	res, _ := im.SliceFloat64(sep)
	return res
}

func (im InputMsg) Float32() (float32, error) {
	if im.err != nil {
		return 0, im.err
	}
	res, err := strconv.ParseFloat(im.msg, 32)
	return float32(res), err
}

func (im InputMsg) MustFloat32() float32 {
	res, _ := im.Float32()
	return res
}

func (im InputMsg) SliceFloat32(sep string) ([]float32, error) {
	data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
		f, err := strconv.ParseFloat(v, 32)
		return float32(f), err
	})
	var res []float32
	for _, v := range data {
		res = append(res, v.(float32))
	}
	return res, err
}

func (im InputMsg) MustSliceFloat32(sep string) []float32 {
	res, _ := im.SliceFloat32(sep)
	return res
}

func YesNo(hint string, defaults bool) bool {
	res, err := YesNoE(hint, defaults)
	if err != nil {
		return defaults
	}
	return res
}

func parseYesNoValue(raw string, defaults bool) (bool, bool) {
	raw = strings.TrimSpace(strings.ToUpper(raw))
	if raw == "" {
		return defaults, true
	}
	switch []rune(raw)[0] {
	case 'Y':
		return true, true
	case 'N':
		return false, true
	default:
		return false, false
	}
}

func YesNoE(hint string, defaults bool) (bool, error) {
	for {
		res, err := MessageBox(hint, "").String()
		if err != nil {
			return false, err
		}
		if answer, ok := parseYesNoValue(res, defaults); ok {
			return answer, nil
		}
	}
}

func buildTriggerPrefixTable(triggerRunes []rune) []int {
	if len(triggerRunes) == 0 {
		return nil
	}
	prefix := make([]int, len(triggerRunes))
	for i := 1; i < len(triggerRunes); i++ {
		j := prefix[i-1]
		for j > 0 && triggerRunes[i] != triggerRunes[j] {
			j = prefix[j-1]
		}
		if triggerRunes[i] == triggerRunes[j] {
			j++
		}
		prefix[i] = j
	}
	return prefix
}

func advanceTriggerIndex(triggerRunes []rune, prefix []int, current int, input rune) (int, bool) {
	if len(triggerRunes) == 0 {
		return 0, true
	}
	for current > 0 && input != triggerRunes[current] {
		current = prefix[current-1]
	}
	if input == triggerRunes[current] {
		current++
		if current == len(triggerRunes) {
			return current, true
		}
	}
	return current, false
}

// StopUntil keeps reading raw input until trigger is matched.
// When trigger == "", it returns after the first key press, which is used for
// "press any key to continue" style prompts.
func StopUntil(hint string, trigger string, repeat bool) error {
	triggerRunes := []rune(trigger)
	prefix := buildTriggerPrefixTable(triggerRunes)
	session, err := rawTerminalSessionFactory(hint, false)
	if err != nil {
		return err
	}
	defer session.Close()
	i := 0
	for {
		b, _, err := session.reader.ReadRune()
		if err != nil {
			return err
		}
		if trigger == "" {
			break
		}
		next, complete := advanceTriggerIndex(triggerRunes, prefix, i, b)
		if complete {
			break
		}
		if next > 0 {
			i = next
			continue
		}
		i = 0
		if hint != "" && repeat {
			fmt.Print("\r\n")
			fmt.Print(hint)
		}
	}
	return nil
}