commandkit/completion_flags.go

package commandkit

import (
	"strings"

	commandmeta "go.maximumdirect.net/commandkit/internal/commandmeta"
)

type completionTargetKind uint8

const (
	completionTargetNone completionTargetKind = iota
	completionTargetCommands
	completionTargetFlags
	completionTargetValue
	completionTargetOperand
)

type completionTarget struct {
	kind        completionTargetKind
	choices     []commandmeta.CompletionChoice
	hint        commandmeta.CompletionHint
	prefix      string
	valuePrefix string
}

type completionFlagState struct {
	supplied map[string]bool
	operands []string
	pending  *commandmeta.Flag
	ended    bool
	invalid  bool
}

func completionFlags(context completionContext) []commandmeta.Flag {
	flags := make([]commandmeta.Flag, 0, len(context.command.Flags)+len(context.command.CoreFlags)+1)
	flags = append(flags, commandmeta.HelpFlag())
	flags = append(flags, context.command.CoreFlags...)
	flags = append(flags, context.command.Flags...)
	return flags
}

func walkCompletionFlags(tokens []string, lookup commandmeta.FlagLookup) completionFlagState {
	state := completionFlagState{supplied: make(map[string]bool)}
	for index := 0; index < len(tokens); index++ {
		token := tokens[index]
		if state.ended {
			state.operands = append(state.operands, token)
			continue
		}
		switch commandmeta.ClassifyFlagToken(token) {
		case commandmeta.TerminatorToken:
			state.ended = true
		case commandmeta.OperandToken:
			state.operands = append(state.operands, token)
		case commandmeta.LongFlagToken:
			spelling, _, attached := commandmeta.SplitLongFlag(token)
			flag, ok := lookup.Find(spelling)
			if !ok {
				state.invalid = true
				return state
			}
			state.supplied[flag.Name] = true
			if !flag.Boolean && !attached {
				if index+1 == len(tokens) {
					state.pending = &flag
					return state
				}
				index++
			}
		case commandmeta.ShortFlagToken:
			for offset := 1; offset < len(token); offset++ {
				flag, ok := lookup.Find("-" + token[offset:offset+1])
				if !ok {
					state.invalid = true
					return state
				}
				state.supplied[flag.Name] = true
				suffix := token[offset+1:]
				if flag.Boolean {
					if strings.HasPrefix(suffix, "=") {
						break
					}
					continue
				}
				if suffix == "" {
					if index+1 == len(tokens) {
						state.pending = &flag
						return state
					}
					index++
				}
				break
			}
		}
	}
	return state
}

func completionTargetFor(snapshot *routingSnapshot, context completionContext) completionTarget {
	if context.command == nil {
		return completionTarget{}
	}
	flags := completionFlags(context)
	lookup := commandmeta.NewFlagLookup(flags, context.aliases)
	state := walkCompletionFlags(context.tail, lookup)
	if state.invalid {
		return completionTarget{}
	}
	if state.pending != nil {
		return completionTarget{kind: completionTargetValue, hint: state.pending.Completion, valuePrefix: context.current}
	}
	if !state.ended && context.commandPosition {
		if children := completionChildren(snapshot, context); len(children) > 0 {
			return completionTarget{kind: completionTargetCommands, choices: children}
		}
	}
	if !state.ended && strings.HasPrefix(context.current, "-") {
		return completionOptionTarget(context.current, flags, lookup, state.supplied)
	}
	if hint, ok := completionOperandHint(context.command.Arguments, len(state.operands)); ok {
		return completionTarget{kind: completionTargetOperand, hint: hint, valuePrefix: context.current}
	}
	if !state.ended && context.current == "" {
		return completionTarget{kind: completionTargetFlags, choices: completionFlagNames(flags, state.supplied, "")}
	}
	return completionTarget{}
}

func completionOptionTarget(current string, flags []commandmeta.Flag, lookup commandmeta.FlagLookup, supplied map[string]bool) completionTarget {
	if current == "-" {
		return completionTarget{kind: completionTargetFlags, choices: completionFlagNames(flags, supplied, current)}
	}
	if strings.HasPrefix(current, "--") {
		spelling, value, attached := commandmeta.SplitLongFlag(current)
		if attached {
			if flag, ok := lookup.Find(spelling); ok {
				return completionTarget{kind: completionTargetValue, hint: flag.Completion, prefix: spelling + "=", valuePrefix: value}
			}
			return completionTarget{}
		}
		return completionTarget{kind: completionTargetFlags, choices: completionFlagNames(flags, supplied, current)}
	}
	used := make(map[string]bool, len(supplied))
	for name, value := range supplied {
		used[name] = value
	}
	for offset := 1; offset < len(current); offset++ {
		flag, ok := lookup.Find("-" + current[offset:offset+1])
		if !ok {
			return completionTarget{}
		}
		used[flag.Name] = true
		suffix := current[offset+1:]
		if flag.Boolean {
			if strings.HasPrefix(suffix, "=") {
				return completionTarget{kind: completionTargetValue, hint: flag.Completion, prefix: current[:offset+2], valuePrefix: suffix[1:]}
			}
			continue
		}
		prefix := current[:offset+1]
		if strings.HasPrefix(suffix, "=") {
			prefix += "="
			suffix = suffix[1:]
		}
		return completionTarget{kind: completionTargetValue, hint: flag.Completion, prefix: prefix, valuePrefix: suffix}
	}
	var choices []commandmeta.CompletionChoice
	for _, flag := range flags {
		if flag.Short != "" && completionFlagEligible(flag, used) {
			choices = append(choices, commandmeta.CompletionChoice{Value: current + flag.Short, Description: flag.Usage})
		}
	}
	return completionTarget{kind: completionTargetFlags, choices: choices}
}

func completionFlagNames(flags []commandmeta.Flag, supplied map[string]bool, prefix string) []commandmeta.CompletionChoice {
	var choices []commandmeta.CompletionChoice
	for _, flag := range flags {
		if !completionFlagEligible(flag, supplied) {
			continue
		}
		for _, spelling := range []string{"--" + flag.Name, shortCompletionSpelling(flag.Short)} {
			if spelling != "" && strings.HasPrefix(spelling, prefix) {
				choices = append(choices, commandmeta.CompletionChoice{Value: spelling, Description: flag.Usage})
			}
		}
	}
	return choices
}

func completionFlagEligible(flag commandmeta.Flag, supplied map[string]bool) bool {
	return !supplied[flag.Name] || flag.Repeated || commandmeta.IsReportingFlag(flag.Name)
}

func shortCompletionSpelling(short string) string {
	if short == "" {
		return ""
	}
	return "-" + short
}

func completionOperandHint(arguments []commandmeta.Argument, completed int) (commandmeta.CompletionHint, bool) {
	total := completed + 1
	counts, err := allocateOperandCounts(total, arguments, false)
	if err != nil {
		return commandmeta.CompletionHint{}, false
	}
	position := 0
	for index, count := range counts {
		position += count
		if position >= total {
			return arguments[index].Completion, true
		}
	}
	return commandmeta.CompletionHint{}, false
}