commandkit/internal/cobra/completion.go

package cobra

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	commandmeta "gitea.maximumdirect.net/eric/CommandKit/internal/commandmeta"
	spfcobra "github.com/spf13/cobra"
)

// CompletionQueryToken is the hidden request understood by generated scripts.
const CompletionQueryToken = spfcobra.ShellCompRequestCmd

const completionPathPrefixMarker = "__commandkit_path_prefix__"

// CompletionReply translates neutral suggestions to the shell line protocol.
func CompletionReply(result commandmeta.CompletionResult) (string, error) {
	var directive spfcobra.ShellCompDirective
	switch result.Kind {
	case commandmeta.CompletionNone, commandmeta.CompletionChoices:
		directive = spfcobra.ShellCompDirectiveNoFileComp
	case commandmeta.CompletionFiles:
		directive = spfcobra.ShellCompDirectiveDefault
	case commandmeta.CompletionDirectories:
		directive = spfcobra.ShellCompDirectiveFilterDirs
	default:
		return "", fmt.Errorf("unsupported completion suggestion strategy")
	}
	// The first line is always metadata, even for a whole-word replacement.
	// Remaining lines before the directive are exclusively literal candidates.
	if strings.ContainsAny(result.Prefix, "\x00\r\n\t") {
		return "", fmt.Errorf("unrepresentable completion prefix")
	}
	var output strings.Builder
	output.WriteString(completionPathPrefixMarker)
	output.WriteString(result.Prefix)
	output.WriteByte('\n')
	if result.Kind == commandmeta.CompletionChoices {
		for _, choice := range result.Choices {
			if choice.Value == "" || strings.ContainsAny(choice.Value, "\x00\r\n\t") {
				continue
			}
			// Generated scripts restore the attached prefix when needed.
			output.WriteString(choice.Value)
			if choice.Description != "" {
				output.WriteByte('\t')
				output.WriteString(strings.NewReplacer("\x00", " ", "\r", " ", "\n", " ", "\t", " ").Replace(choice.Description))
			}
			output.WriteByte('\n')
		}
	}
	fmt.Fprintf(&output, ":%d\n", directive)
	return output.String(), nil
}

// CompletionErrorReply is the protocol response for a malformed query.
func CompletionErrorReply() string {
	return fmt.Sprintf(":%d\n", spfcobra.ShellCompDirectiveError)
}

// CompletionScript generates a shell script for the installed executable.
func CompletionScript(w io.Writer, executable, shell string) error {
	root := &spfcobra.Command{Use: executable}
	var generated bytes.Buffer
	var err error
	switch shell {
	case "bash":
		err = root.GenBashCompletionV2(&generated, true)
	case "zsh":
		err = root.GenZshCompletion(&generated)
	case "fish":
		err = root.GenFishCompletion(&generated, true)
	default:
		return fmt.Errorf("unsupported completion shell %q", shell)
	}
	if err != nil {
		return err
	}

	script := generated.String()
	if shell == "bash" {
		script, err = adjustBashCompletion(script, root.Name())
	} else if shell == "zsh" {
		script, err = adjustZshCompletion(script)
	} else if shell == "fish" {
		script, err = adjustFishCompletion(script, root.Name())
	}
	if err != nil {
		return err
	}
	n, err := io.WriteString(w, script)
	if err == nil && n != len(script) {
		return io.ErrShortWrite
	}
	return err
}

func adjustBashCompletion(script, name string) (string, error) {
	var err error
	script, err = replaceCompletionSection(script, "    local requestComp lastParam lastChar args", "    # Extract the directive", `    local -a requestWords=()
    local word char quote decoded next i backtick=$'\x60'
    # Bash's completion helpers supply shell spelling, not evaluated argv.
    # Remove quoting and escapes without expanding substitutions or globs.
    for word in "${words[@]}"; do
        quote=
        decoded=
        for ((i=0; i<${#word}; i++)); do
            char=${word:i:1}
            case "$quote:$char" in
                "':'"|'":"') quote= ;;
                ':"'|":'") quote=$char ;;
                :\\|'"':\\)
                    if ((i+1 < ${#word})); then
                        next=${word:i+1:1}
                        if [[ $quote == '"' && $next != '\' && $next != '"' && $next != '$' && $next != "$backtick" ]]; then
                            decoded+=$char
                            continue
                        fi
                        ((i+=1))
                        decoded+=${word:i:1}
                    else
                        decoded+=$char
                    fi ;;
                *) decoded+=$char ;;
            esac
        done
        requestWords+=("$decoded")
    done
    if [[ -z $cur && ${words[${#words[@]}-1]} != *= ]]; then
        # The Bash helpers may omit the empty word at a trailing space.
        if ((${#requestWords[@]} <= cword)); then requestWords+=(""); fi
    fi
    cur=${requestWords[${#requestWords[@]}-1]}
    out=$(COBRA_ACTIVE_HELP=0 "${requestWords[0]}" __complete "${requestWords[@]:1}" 2>/dev/null)

`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionText(script, `local out directive`, `local out directive completionPrefix completionOriginalCur`)
	if err != nil {
		return "", err
	}
	errorDebug := fmt.Sprintf(`__%s_debug "Received error from custom completion go code"`, name)
	script, err = replaceCompletionText(script, errorDebug, errorDebug+"\n        compopt +o default")
	if err != nil {
		return "", err
	}
	parsePrefix := `if [[ ${out} != ` + completionPathPrefixMarker + `* ]]; then
        compopt +o default
        return
    fi
    completionPrefix=${out%%$'\n'*}
    completionPrefix=${completionPrefix#` + completionPathPrefixMarker + `}
    completionOriginalCur=$cur
    if [[ $completionOriginalCur != "$completionPrefix"* ]]; then
        compopt +o default
        return
    fi
    if [[ $out == *$'\n'* ]]; then out=${out#*$'\n'}; else out=; fi
    cur=${completionOriginalCur#"$completionPrefix"}

    # Separate activeHelp from normal completions`
	script, err = replaceCompletionText(script, `# Separate activeHelp from normal completions`, parsePrefix)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionSection(script, "    local activeHelpMarker=", "\n}\n", `    while IFS='' read -r comp; do
        [[ -z $comp ]] || completions+=("$comp")
    done <<<"${out}"
`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionText(script,
		`if (((directive & shellCompDirectiveFilterFileExt) != 0)); then`,
		`if ((directive == 0)); then
        compopt +o default
        _filedir
    elif (((directive & shellCompDirectiveFilterFileExt) != 0)); then`)
	if err != nil {
		return "", err
	}
	old := fmt.Sprintf("__%[1]s_handle_special_char \"$cur\" :\n    __%[1]s_handle_special_char \"$cur\" =", name)
	replacement := fmt.Sprintf(`if [[ -n $completionPrefix ]]; then
        compopt +o default
        local i
        for i in "${!COMPREPLY[@]}"; do COMPREPLY[$i]="${completionPrefix}${COMPREPLY[$i]}"; done
        cur=$completionOriginalCur
    fi
    __%[1]s_handle_special_char "$cur" :
    __%[1]s_handle_special_char "$cur" =`, name)
	return replaceCompletionText(script, old, replacement)
}

func adjustZshCompletion(script string) (string, error) {
	var err error
	script, err = replaceCompletionText(script, `words=("${=words[1,CURRENT]}")`, `words=("${(@Q)words[1,CURRENT]}")`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionSection(script, "    # For zsh, when completing a flag", "    # Extract the directive", `    local -a prefixArgs orderArgs spaceArgs
    out=$(COBRA_ACTIVE_HELP=0 "${words[1]}" __complete "${words[@]:1}" 2>/dev/null)

`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionSection(script, `    local activeHelpMarker=`, "    if [ $((directive & shellCompDirectiveNoSpace))", `    [[ "$out" == `+completionPathPrefixMarker+`* ]] || return 1
    flagPrefix="${out%%$'\n'*}"
    flagPrefix="${flagPrefix#`+completionPathPrefixMarker+`}"
    [[ "$lastParam" == "$flagPrefix"* ]] || return 1
    if [[ $out == *$'\n'* ]]; then out="${out#*$'\n'}"; else out=; fi
    if [[ -n "$flagPrefix" ]]; then
        prefixArgs=(-P "$flagPrefix")
    fi
    local tab="$(printf '\t')"
    while IFS= read -r comp; do
        if [[ -n "$comp" ]]; then
            comp=${comp//:/\\:}
            completions+=("${comp//$tab/:}")
        fi
    done <<<"$out"

`)
	if err != nil {
		return "", err
	}
	return replaceCompletionSection(script, "    if [ $((directive & shellCompDirectiveNoSpace))", "\n}\n", `    if [ $((directive & shellCompDirectiveNoSpace)) -ne 0 ]; then
        spaceArgs=(-S "")
    fi
    if [ $((directive & shellCompDirectiveKeepOrder)) -ne 0 ]; then
        orderArgs=(-V)
    fi

    if [ $((directive & shellCompDirectiveNoFileComp)) -ne 0 ]; then
        _describe "${orderArgs[@]}" completions completions "${prefixArgs[@]}" "${spaceArgs[@]}"
    else
        if [[ -n "$flagPrefix" ]]; then
            compset -P "${(b)flagPrefix}"
        fi
        if [ $((directive & shellCompDirectiveFilterDirs)) -ne 0 ]; then
            _files -/
        else
            _files
        fi
    fi
`)
}

func replaceCompletionSection(script, start, end, replacement string) (string, error) {
	i := strings.Index(script, start)
	if i < 0 {
		return "", fmt.Errorf("unsupported generated completion script")
	}
	j := strings.Index(script[i+len(start):], end)
	if j < 0 {
		return "", fmt.Errorf("unsupported generated completion script")
	}
	return script[:i] + replacement + script[i+len(start)+j:], nil
}

func replaceCompletionText(script, old, replacement string) (string, error) {
	if strings.Count(script, old) != 1 {
		return "", fmt.Errorf("unsupported generated completion script")
	}
	return strings.Replace(script, old, replacement, 1), nil
}

func adjustFishCompletion(script, name string) (string, error) {
	var err error
	cobraName := strings.NewReplacer("-", "_", ":", "_").Replace(name)
	variableName := strings.ReplaceAll(cobraName, ".", "_")
	if cobraName != variableName {
		script = strings.ReplaceAll(script, "__"+cobraName+"_", "__"+variableName+"_")
	}
	script, err = replaceCompletionText(script,
		`set -l lastArg (string escape -- (commandline -ct))`,
		"set -l lastArg (string unescape -- (commandline -ct))\n    set -q lastArg[1]; or set lastArg \"\"")
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionText(script,
		`set -l results (eval $requestComp 2> /dev/null)`,
		`set -lx COBRA_ACTIVE_HELP 0`+"\n    "+`set -l results ($args[1] `+spfcobra.ShellCompRequestCmd+` $args[2..-1] "$lastArg" 2> /dev/null)`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionSection(script, "    # For Fish, when completing a flag", `    printf "%s\n" "$directiveLine"`, `    if not set -q comps[1]; or not string match -q -- '`+completionPathPrefixMarker+`*' "$comps[1]"
        printf ":%s\n" `+fmt.Sprint(spfcobra.ShellCompDirectiveError)+`
        return
    end
    set -l flagPrefix (string replace -- '`+completionPathPrefixMarker+`' '' "$comps[1]")
    set -q flagPrefix[1]; or set flagPrefix ""
    printf "%s\n" "$comps[1]"
    set -e comps[1]
    for comp in $comps
        printf "%s%s\n" "$flagPrefix" "$comp"
    end

`)
	if err != nil {
		return "", err
	}
	variable := "__" + variableName + "_comp_results"
	errorBranch := fmt.Sprintf(`__%[1]s_debug "Received error directive: aborting."
        # Might as well do file completion, in case it helps
        return 1`, variableName)
	errorReplacement := fmt.Sprintf(`__%[1]s_debug "Received error directive: aborting."
        return 0`, variableName)
	script, err = replaceCompletionText(script, errorBranch, errorReplacement)
	if err != nil {
		return "", err
	}
	pathContext := fmt.Sprintf(`set -l dirfilter (math (math --scale 0 $directive / $shellCompDirectiveFilterDirs) %% 2)
    if test (count $%[1]s) -gt 0; and string match -q -- '%[2]s*' "$%[1]s[1]"
        set -l pathPrefix (string replace -- '%[2]s' '' "$%[1]s[1]")
        set -l token (string unescape -- (commandline -ct))
        set -e %[1]s[1]
        set -l nofiles (math (math --scale 0 $directive / $shellCompDirectiveNoFileComp) %% 2)
        if test $nofiles -eq 1
            return 0
        end
        set -l value (string sub -s (math (string length -- "$pathPrefix") + 1) -- "$token")
        set -q value[1]; or set value ""
        set -l escaped (string escape --no-quoted -- "$value" | string replace -a -- '=' '\\=')
        set --global %[1]s
        if test $dirfilter -eq 1
            for entry in (__fish_complete_directories "$escaped")
                set -a %[1]s "$pathPrefix$entry"
            end
        else
            for entry in (complete -C "'' $escaped")
                set -a %[1]s "$pathPrefix$entry"
            end
        end
        __%[3]s_debug "Path completions: $%[1]s"
        return 0
    end`, variable, completionPathPrefixMarker, variableName)
	script, err = replaceCompletionText(script, `set -l dirfilter (math (math --scale 0 $directive / $shellCompDirectiveFilterDirs) % 2)`, pathContext)
	if err != nil {
		return "", err
	}
	return script, nil
}