iris-context/commands/fs/zoxide.go

package fs

import (
	"context"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/versenilvis/fuzzy"
	"github.com/versenilvis/iris/spec"
)

const (
	zoxideCacheTTL     = 2 * time.Second
	zoxideQueryTimeout = 500 * time.Millisecond
	zoxideEmptyLimit   = 20
	zoxideSearchLimit  = 10
)

type zoxideQueryFunc func(context.Context) ([]byte, error)

type zoxideSource struct {
	mu      sync.Mutex
	query   zoxideQueryFunc
	now     func() time.Time
	dirs    []string
	err     error
	fetched time.Time
}

func init() {
	spec.Register(&spec.Spec{
		Name:        "z",
		Description: "jump to directory",
		MaxArgs:     0,
		Generator:   ZoxideGenerator(),
	})
	spec.Register(&spec.Spec{
		Name:        "zi",
		Description: "jump to directory interactively",
		MaxArgs:     0,
		Generator:   ZoxideGenerator(),
	})
}

func newZoxideSource(query zoxideQueryFunc) *zoxideSource {
	return &zoxideSource{query: query, now: time.Now}
}

var defaultZoxideSource = newZoxideSource(func(ctx context.Context) ([]byte, error) {
	return exec.CommandContext(ctx, "zoxide", "query", "-l").Output()
})

// directories caches failures too, preventing a missing zoxide binary from
// spawning a failed process on every keystroke.
func (s *zoxideSource) directories() ([]string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := s.now()
	if !s.fetched.IsZero() && now.Before(s.fetched.Add(zoxideCacheTTL)) {
		return append([]string(nil), s.dirs...), s.err
	}

	ctx, cancel := context.WithTimeout(context.Background(), zoxideQueryTimeout)
	defer cancel()

	out, err := s.query(ctx)
	s.fetched = s.now()
	s.err = err
	s.dirs = nil
	if err != nil {
		return nil, err
	}

	for line := range strings.SplitSeq(string(out), "\n") {
		line = strings.TrimSuffix(line, "\r")
		if line != "" {
			s.dirs = append(s.dirs, line)
		}
	}

	return append([]string(nil), s.dirs...), nil
}

func ZoxideGenerator() spec.GeneratorFunc {
	return newZoxideGenerator(defaultZoxideSource.directories)
}

func newZoxideGenerator(directorySource func() ([]string, error)) spec.GeneratorFunc {
	return func(tokens []string, prefix string, partial string) []spec.Suggestion {
		fullQuery := strings.Join(tokens[1:], " ")
		local := spec.FileGenerator("/")(tokens, prefix, fullQuery)
		if strings.ContainsAny(fullQuery, `/\`) {
			return local
		}

		dirs, err := directorySource()
		if err != nil {
			return local
		}

		external := zoxideSuggestions(dirs, fullQuery)
		return mergeDirectorySuggestions(local, external, spec.GetCWD())
	}
}

func zoxideSuggestions(dirs []string, query string) []spec.Suggestion {
	candidates := dirs
	limit := zoxideEmptyLimit
	if query != "" {
		candidates = matchedZoxideDirectories(dirs, query)
		limit = zoxideSearchLimit
	}

	home, _ := os.UserHomeDir()
	results := make([]spec.Suggestion, 0, min(limit, len(candidates)))
	seen := make(map[string]struct{}, len(candidates))
	for _, candidate := range candidates {
		path := filepath.Clean(candidate)
		if !filepath.IsAbs(path) {
			continue
		}
		if _, ok := seen[path]; ok {
			continue
		}
		seen[path] = struct{}{}

		info, err := os.Stat(path)
		if err != nil || !info.IsDir() {
			continue
		}

		results = append(results, spec.Suggestion{
			Cmd:      path,
			Desc:     displayDirectory(path, home),
			Priority: 40,
		})
		if len(results) == limit {
			break
		}
	}

	return results
}

func matchedZoxideDirectories(dirs []string, query string) []string {
	if len(dirs) == 0 {
		return nil
	}

	results := make([]string, 0, len(dirs))
	seen := make(map[string]struct{}, len(dirs))
	lowerQuery := strings.ToLower(query)

	for _, dir := range dirs {
		if strings.Contains(strings.ToLower(dir), lowerQuery) {
			results = append(results, dir)
			seen[dir] = struct{}{}
		}
	}

	matches := fuzzy.NewSearcher(dirs).SearchWithScores(query, &fuzzy.SearchOptions{Limit: len(dirs)})
	for _, match := range matches {
		if _, ok := seen[match.Str]; ok {
			continue
		}
		results = append(results, match.Str)
		seen[match.Str] = struct{}{}
	}

	return results
}

func mergeDirectorySuggestions(local, external []spec.Suggestion, cwd string) []spec.Suggestion {
	results := make([]spec.Suggestion, 0, len(local)+len(external))
	seen := make(map[string]struct{}, len(local)+len(external))

	add := func(suggestion spec.Suggestion) {
		key := directoryKey(suggestion.Cmd, cwd)
		if _, ok := seen[key]; ok {
			return
		}
		seen[key] = struct{}{}
		results = append(results, suggestion)
	}
	for _, suggestion := range local {
		add(suggestion)
	}
	for _, suggestion := range external {
		add(suggestion)
	}

	return results
}

func directoryKey(path, cwd string) string {
	if path == "" {
		return ""
	}
	if !filepath.IsAbs(path) {
		path = filepath.Join(cwd, path)
	}
	return filepath.Clean(path)
}

func displayDirectory(path, home string) string {
	path = filepath.Clean(path)
	if home == "" {
		return path
	}

	rel, err := filepath.Rel(filepath.Clean(home), path)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return path
	}
	if rel == "." {
		return "~"
	}
	return filepath.Join("~", rel)
}