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, commandmeta.EffectiveFlags(context.command.Flags, context.command.CoreFlags)...)
	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, commandmeta.ShortFlagToken:
			for offset := 1; ; {
				fact := commandmeta.ConsumeFlag(lookup, token, offset)
				if !fact.Found {
					state.invalid = true
					return state
				}
				state.supplied[fact.Flag.Name] = true
				if fact.NeedsNext {
					if index+1 == len(tokens) {
						flag := fact.Flag
						state.pending = &flag
						return state
					}
					index++
				}
				if fact.NextOffset == 0 {
					break
				}
				offset = fact.NextOffset
			}
		}
	}
	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, "--") {
		fact := commandmeta.ConsumeFlag(lookup, current, 1)
		if fact.Attached {
			if fact.Found {
				return completionTarget{kind: completionTargetValue, hint: fact.Flag.Completion, prefix: fact.Spelling + "=", valuePrefix: fact.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); {
		fact := commandmeta.ConsumeFlag(lookup, current, offset)
		if !fact.Found {
			return completionTarget{}
		}
		used[fact.Flag.Name] = true
		if fact.Flag.Boolean && fact.Attached {
			return completionTarget{kind: completionTargetValue, hint: fact.Flag.Completion, prefix: current[:offset+2], valuePrefix: fact.Value}
		}
		if !fact.Flag.Boolean {
			prefix := current[:offset+1]
			if strings.HasPrefix(current[offset+1:], "=") {
				prefix += "="
			}
			return completionTarget{kind: completionTargetValue, hint: fact.Flag.Completion, prefix: prefix, valuePrefix: fact.Value}
		}
		if fact.NextOffset == 0 {
			break
		}
		offset = fact.NextOffset
	}
	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
}