iris-context/spec/lookup.go

package spec

import (
	"maps"
	"slices"
	"strings"
	"sync"

	"github.com/versenilvis/iris/integration/shell"
	"github.com/versenilvis/iris/internal/logger"
	"github.com/versenilvis/iris/spec/alias"
)

var (
	ShellAliases   = map[string]string{}
	shellAliasesMu sync.RWMutex
)

func GetAlias(name string) (string, bool) {
	shellAliasesMu.RLock()
	defer shellAliasesMu.RUnlock()
	val, ok := ShellAliases[name]
	return val, ok
}

func GetAliasesCopy() map[string]string {
	shellAliasesMu.RLock()
	defer shellAliasesMu.RUnlock()
	cp := make(map[string]string, len(ShellAliases))
	maps.Copy(cp, ShellAliases)
	return cp
}

// Lookup finds matching suggestions for your input by looking at how many words you typed
// it changes aliases to real commands, finds subcommands inside others, and runs generators for more suggestions
// e.g. Lookup("git che") -> suggests "git checkout"
// e.g. Lookup("git checkout ") -> suggests branch names via generator
func Lookup(input string) []Suggestion {
	if shell.Current != nil {
		aliases := shell.Current.ScanAliases()
		shellAliasesMu.Lock()
		ShellAliases = aliases
		shellAliasesMu.Unlock()
	}

	shellAliasesMu.RLock()
	aliases := ShellAliases
	shellAliasesMu.RUnlock()

	if input == "" {
		return topLevelSuggestions("", aliases)
	}

	tokens := Tokenize(input)
	originalTokens := append([]string(nil), tokens...)
	originalPrefix := ""
	expandedPrefix := ""

	if len(tokens) == 1 && tokens[0] == "" {
		return topLevelSuggestions("", aliases)
	}
	// NOTE: remember that wrapper already check if we type nothing, but I just want to make sure
	// in the future, maybe we will write unit test or using Lookup in another module
	// it will be safer because we maybe forget to check it in that module

	scanExternalCommands()

	// if you have an alias in your shell config like: alias nv="nvim"
	// expand it even when there's only one token (e.g. "nv ") so that the
	// target spec's Generator (FileGenerator etc.) can still fire.
	// only expand when there's a trailing space — i.e. the user has committed
	// to the alias name and is now typing arguments (tokens last elem == "")
	hasTrailingSpace := len(tokens) > 0 && tokens[len(tokens)-1] == ""
	if hasTrailingSpace || len(tokens) > 1 {
		if target, ok := aliases[tokens[0]]; ok {
			aliasTokens := Tokenize(target)
			if len(aliasTokens) > 0 && aliasTokens[len(aliasTokens)-1] == "" {
				aliasTokens = aliasTokens[:len(aliasTokens)-1]
			}
			tokens = append(aliasTokens, tokens[1:]...)
		}
	}

	// Expand tool aliases
	isAliasExpanded := false
	aliasExpansionEnd := 0
	if len(tokens) > 2 {
		rootCmdName := tokens[0]
		if provider := alias.GetProvider(rootCmdName); provider != nil {
			cwd := GetCWD()
			toolAliases := provider.GetAliases(cwd)
			subCmd := tokens[1]
			for _, a := range toolAliases {
				if a.Name == subCmd {
					if strings.HasPrefix(a.Expansion, "!") {
						return nil
					}
					aliasTokens := Tokenize(a.Expansion)
					if len(aliasTokens) > 0 && aliasTokens[len(aliasTokens)-1] == "" {
						aliasTokens = aliasTokens[:len(aliasTokens)-1]
					}
					if len(aliasTokens) > 0 && aliasTokens[0] == rootCmdName {
						aliasTokens = aliasTokens[1:]
					}
					newTokens := make([]string, 0, len(tokens)-1+len(aliasTokens))
					newTokens = append(newTokens, tokens[0])
					newTokens = append(newTokens, aliasTokens...)
					newTokens = append(newTokens, tokens[2:]...)
					tokens = newTokens
					isAliasExpanded = true
					aliasExpansionEnd = len(aliasTokens)
					expandedPrefix = strings.Join(tokens[:1+aliasExpansionEnd], " ")
					if len(originalTokens) >= 2 {
						originalPrefix = strings.Join(originalTokens[:2], " ")
					}
					break
				}
			}
		}
	}

	if len(tokens) == 1 {
		query := tokens[0]
		results := topLevelSuggestions(query, aliases)

		if spec, exists := Registry[query]; exists {
			hasTrailingSpace := query != "" && query[len(query)-1] == ' '

			if hasTrailingSpace {
				partial := ""
				prefix := query

				for _, sub := range spec.Subcommands {
					results = append(results, Suggestion{
						Cmd: strings.TrimSpace(query) + " " + sub.Name, Desc: sub.Description, Icon: query, Priority: sub.Priority,
					})
				}
				if spec.Generator != nil {
					genResults := spec.Generator(tokens, prefix, partial)
					for _, g := range genResults {
						results = append(results, Suggestion{
							Cmd: strings.TrimSpace(query) + " " + g.Cmd, Desc: g.Desc, Icon: query, Priority: g.Priority,
						})
					}
				}
			}
		}
		return results
	}

	rootCmdName := tokens[0]
	spec, exists := Registry[rootCmdName]
	logger.Debugf("core lookup tokens: %v, registry exists: %v", tokens, exists)
	if !exists {
		partial := tokens[len(tokens)-1]
		args := tokens[1:]
		if len(args) > 0 && args[len(args)-1] == partial {
			args = args[:len(args)-1]
		}
		if cobSugg := QueryCobraComplete(rootCmdName, args, partial); len(cobSugg) > 0 {
			return cobSugg
		}
		return nil
	}

	currentSubs, currentOpts, currentGen := spec.Subcommands, spec.Options, spec.Generator
	depth := 1

	for depth < len(tokens)-1 {
		tok := tokens[depth]
		if tok == "" || strings.HasPrefix(tok, "-") || strings.Contains(tok, "=") {
			depth++
			continue
		}

		found := false
		for _, sub := range currentSubs {
			match := sub.Name == tok
			if !match {
				if slices.Contains(sub.Aliases, tok) {
					match = true
				}
			}

			if match {
				currentSubs, currentOpts, currentGen = sub.Subcommands, sub.Options, sub.Generator
				found = true
				break
			}
		}

		if found {
			depth++
			continue
		}
		// invalid subcommand word typed
		if len(currentSubs) > 0 && (!isAliasExpanded || depth > aliasExpansionEnd) {
			return nil
		}
		break
	}

	results := []Suggestion{}

	currentLimit := spec.MaxArgs
	tempSubs := spec.Subcommands
	for i := 1; i < depth; i++ {
		tok := tokens[i]
		for _, sub := range tempSubs {
			if sub.Name == tok {
				currentLimit = sub.MaxArgs
				tempSubs = sub.Subcommands
				break
			}
		}
	}

	argCount := 0
	for i := depth; i < len(tokens)-1; i++ {
		t := tokens[i]
		if t != "" && !strings.HasPrefix(t, "-") && !strings.Contains(t, "=") {
			argCount++
		}
	}

	partial := tokens[len(tokens)-1]
	allowMoreArgs := currentLimit <= 0 || argCount < currentLimit

	logger.Debugf("core query tokens: %v (partial: '%s')", tokens, partial)
	logger.Debugf("core depth: %d, argCount: %d, limit: %d, allowMore: %v", depth, argCount, currentLimit, allowMoreArgs)

	prefixBuilder := strings.Builder{}
	for i := 0; i < depth; i++ {
		if i > 0 {
			prefixBuilder.WriteByte(' ')
		}
		prefixBuilder.WriteString(tokens[i])
	}
	prefix := prefixBuilder.String()

	linePrefixBuilder := strings.Builder{}
	for i := 0; i < len(tokens)-1; i++ {
		if i > 0 {
			linePrefixBuilder.WriteByte(' ')
		}
		linePrefixBuilder.WriteString(tokens[i])
	}
	linePrefix := linePrefixBuilder.String()

	if currentGen != nil && allowMoreArgs {
		genPartial := partial
		genLinePrefix := linePrefix

		if partial == "" && len(tokens) >= 3 {
			prevToken := tokens[len(tokens)-2]
			if strings.HasSuffix(prevToken, "/") || strings.HasSuffix(prevToken, "\\") {
				genPartial = prevToken
				lpBuilder := strings.Builder{}
				for i := 0; i < len(tokens)-2; i++ {
					if i > 0 {
						lpBuilder.WriteByte(' ')
					}
					lpBuilder.WriteString(tokens[i])
				}
				genLinePrefix = lpBuilder.String()
			}
		}

		genResults := currentGen(tokens, prefix, genPartial)

		for _, g := range genResults {
			if genPartial != "" && !HasPrefix(g.Cmd, genPartial) && !strings.Contains(g.Cmd, genPartial) {
				continue
			}

			suggested := g.Cmd
			if strings.Contains(suggested, " ") && !strings.HasPrefix(suggested, "\"") {
				suggested = "\"" + suggested + "\""
			}

			// if the suggestion is a full path that includes
			// words already in the command line (multi-word support), we replace
			// the entire argument part by using prefix only
			finalCmd := ""
			if len(tokens) > depth+1 && strings.HasPrefix(g.Cmd, tokens[depth]) {
				finalCmd = prefix + " " + suggested
			} else if genLinePrefix != "" {
				finalCmd = genLinePrefix + " " + suggested
			} else {
				finalCmd = suggested
			}

			newTokens := Tokenize(finalCmd)
			if len(newTokens) > 0 {
				lastToken := newTokens[len(newTokens)-1]
				isDuplicate := false
				for i := 0; i < len(tokens)-1; i++ {
					if tokens[i] == lastToken {
						isDuplicate = true
						break
					}
				}
				if isDuplicate {
					continue
				}
			}

			results = append(results, Suggestion{
				Cmd:      finalCmd,
				Desc:     g.Desc,
				Icon:     rootCmdName,
				Priority: g.Priority,
			})
		}
	}

	if allowMoreArgs {
		for _, sub := range currentSubs {
			if partial == "" || HasPrefix(sub.Name, partial) {
				results = append(results, Suggestion{
					Cmd: prefix + " " + sub.Name, Desc: sub.Description, Icon: rootCmdName, Priority: sub.Priority,
				})
			}
		}

		if depth == 1 {
			if provider := alias.GetProvider(rootCmdName); provider != nil {
				cwd := GetCWD()
				for _, a := range provider.GetAliases(cwd) {
					if partial == "" || HasPrefix(a.Name, partial) {
						priority := 70
						switch a.Scope {
						case "local", "worktree", "command":
							priority = 85
						case "system":
							priority = 65
						}
						results = append(results, Suggestion{
							Cmd:      prefix + " " + a.Name,
							Desc:     "Alias for: " + a.Expansion,
							Icon:     rootCmdName,
							Priority: priority,
							Source:   "alias",
						})
					}
				}
			}
		}
	}

	usedOpts := make(map[string]bool)
	for _, t := range tokens {
		if strings.HasPrefix(t, "-") {
			usedOpts[t] = true
		}
	}
	for _, opt := range currentOpts {
		trimmedOpt := strings.TrimLeft(opt.Name, "-")
		if !usedOpts[opt.Name] && (partial == "" || HasPrefix(opt.Name, partial) || HasPrefix(trimmedOpt, partial)) {
			optPriority := opt.Priority
			if optPriority == 0 {
				if strings.HasPrefix(partial, "-") {
					optPriority = 80
				} else {
					optPriority = 10
				}
			}
			results = append(results, Suggestion{
				Cmd: linePrefix + " " + opt.Name, Desc: opt.Description, Icon: rootCmdName, Priority: optPriority,
			})
		}
	}

	if isAliasExpanded && expandedPrefix != "" && originalPrefix != "" {
		for i := range results {
			if strings.HasPrefix(results[i].Cmd, expandedPrefix) {
				results[i].Cmd = originalPrefix + strings.TrimPrefix(results[i].Cmd, expandedPrefix)
			}
		}
	}

	return results
}

func topLevelSuggestions(query string, aliases map[string]string) []Suggestion {
	scanExternalCommands()
	results, seen := []Suggestion{}, make(map[string]bool)

	for name, target := range aliases {
		if !seen[name] && (query == "" || HasPrefix(name, query)) {
			results = append(results, Suggestion{
				Cmd: name, Desc: target, Icon: "alias",
			})
			seen[name] = true
		}
	}

	for name, spec := range Registry {
		if seen[name] {
			continue
		}
		match := false
		if query == "" || HasPrefix(name, query) {
			match = true
		} else {
			for _, a := range spec.Aliases {
				if HasPrefix(a, query) {
					match = true
					break
				}
			}
		}
		if match {
			icon := spec.Icon
			if icon == "" {
				icon = name
			}
			results = append(results, Suggestion{Cmd: name, Desc: spec.Description, Icon: icon})
			seen[name] = true
		}
	}

	for name := range pathCmds {
		if !seen[name] && (query == "" || HasPrefix(name, query)) {
			results = append(results, Suggestion{
				Cmd: name, Desc: "system command", Icon: "system",
			})
		}
	}

	return results
}