<commit_before>package shell

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"

	prompt "github.com/c-bata/go-prompt"
	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/uuid"
	"github.com/spf13/cobra"
)

const (
	completionAnnotation string = "completion"
	ldThreshold          int    = 2
	pachChar             rune   = '🐘'
)

type CompletionFunc func(flag, arg string) []prompt.Suggest

var completions map[string]CompletionFunc = make(map[string]CompletionFunc)

func RegisterCompletionFunc(cmd *cobra.Command, completionFunc CompletionFunc) {
	id := uuid.NewWithoutDashes()

	if cmd.Annotations == nil {
		cmd.Annotations = make(map[string]string)
	} else if _, ok := cmd.Annotations[completionAnnotation]; ok {
		panic("duplicate completion func registration")
	}
	cmd.Annotations[completionAnnotation] = id
	completions[id] = completionFunc
}

type shell struct {
	rootCmd *cobra.Command
}

func newShell(rootCmd *cobra.Command) *shell {
	return &shell{rootCmd: rootCmd}
}

func (s *shell) executor(in string) {
	cmd := exec.Command("bash")
	cmd.Stdin = strings.NewReader("pachctl " + in)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
	return
}

func (s *shell) suggestor(in prompt.Document) []prompt.Suggest {
	args := strings.Fields(in.Text)
	if len(strings.TrimSuffix(in.Text, " ")) < len(in.Text) {
		args = append(args, "")
	}
	cmd := s.rootCmd
	text := ""
	if len(args) > 0 {
		var err error
		cmd, _, err = s.rootCmd.Traverse(args[:len(args)-1])
		if err != nil {
			log.Fatal(err)
		}
		text = args[len(args)-1]
	}
	flag := ""
	if len(args) > 1 {
		if args[len(args)-2][0] == '-' {
			flag = args[len(args)-2]
		}
	}
	suggestions := cmd.SuggestionsFor(text)
	if len(suggestions) > 0 {
		var result []prompt.Suggest
		for _, suggestion := range suggestions {
			cmd, _, err := cmd.Traverse([]string{suggestion})
			if err != nil {
				log.Fatal(err)
			}
			result = append(result, prompt.Suggest{
				Text:        suggestion,
				Description: cmd.Short,
			})
		}
		return result
	}
	if id, ok := cmd.Annotations[completionAnnotation]; ok {
		completionFunc := completions[id]
		suggests := completionFunc(flag, text)
		var result []prompt.Suggest
		for _, s := range suggests {
			sText := s.Text
			if len(text) < len(sText) {
				sText = sText[:len(text)]
			}
			if ld(sText, text, true) < ldThreshold {
				result = append(result, s)
			}
		}
		return result
	}
	return nil
}

func (s *shell) run() {
	prompt.New(
		s.executor,
		s.suggestor,
		prompt.OptionPrefix(fmt.Sprintf("%c > ", pachChar)),
		prompt.OptionLivePrefix(func() (string, bool) {
			cfg, err := config.Read()
			if err != nil {
				return "", false
			}
			activeContext, _, err := cfg.ActiveContext()
			if err != nil {
				return "", false
			}
			return fmt.Sprintf("%c  context:(%s) >", pachChar, activeContext), true
		}),
	).Run()
}

// Run runs a prompt, it does not return.
func Run(rootCmd *cobra.Command) {
	newShell(rootCmd).run()
}

// ld computes the Levenshtein Distance for two strings.
func ld(s, t string, ignoreCase bool) int {
	if ignoreCase {
		s = strings.ToLower(s)
		t = strings.ToLower(t)
	}
	d := make([][]int, len(s)+1)
	for i := range d {
		d[i] = make([]int, len(t)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {<commit_msg>Remove the 🐘 from the prompt.<commit_after>package shell

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"

	prompt "github.com/c-bata/go-prompt"
	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/uuid"
	"github.com/spf13/cobra"
)

const (
	completionAnnotation string = "completion"
	ldThreshold          int    = 2
)

type CompletionFunc func(flag, arg string) []prompt.Suggest

var completions map[string]CompletionFunc = make(map[string]CompletionFunc)

func RegisterCompletionFunc(cmd *cobra.Command, completionFunc CompletionFunc) {
	id := uuid.NewWithoutDashes()

	if cmd.Annotations == nil {
		cmd.Annotations = make(map[string]string)
	} else if _, ok := cmd.Annotations[completionAnnotation]; ok {
		panic("duplicate completion func registration")
	}
	cmd.Annotations[completionAnnotation] = id
	completions[id] = completionFunc
}

type shell struct {
	rootCmd *cobra.Command
}

func newShell(rootCmd *cobra.Command) *shell {
	return &shell{rootCmd: rootCmd}
}

func (s *shell) executor(in string) {
	cmd := exec.Command("bash")
	cmd.Stdin = strings.NewReader("pachctl " + in)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
	return
}

func (s *shell) suggestor(in prompt.Document) []prompt.Suggest {
	args := strings.Fields(in.Text)
	if len(strings.TrimSuffix(in.Text, " ")) < len(in.Text) {
		args = append(args, "")
	}
	cmd := s.rootCmd
	text := ""
	if len(args) > 0 {
		var err error
		cmd, _, err = s.rootCmd.Traverse(args[:len(args)-1])
		if err != nil {
			log.Fatal(err)
		}
		text = args[len(args)-1]
	}
	flag := ""
	if len(args) > 1 {
		if args[len(args)-2][0] == '-' {
			flag = args[len(args)-2]
		}
	}
	suggestions := cmd.SuggestionsFor(text)
	if len(suggestions) > 0 {
		var result []prompt.Suggest
		for _, suggestion := range suggestions {
			cmd, _, err := cmd.Traverse([]string{suggestion})
			if err != nil {
				log.Fatal(err)
			}
			result = append(result, prompt.Suggest{
				Text:        suggestion,
				Description: cmd.Short,
			})
		}
		return result
	}
	if id, ok := cmd.Annotations[completionAnnotation]; ok {
		completionFunc := completions[id]
		suggests := completionFunc(flag, text)
		var result []prompt.Suggest
		for _, s := range suggests {
			sText := s.Text
			if len(text) < len(sText) {
				sText = sText[:len(text)]
			}
			if ld(sText, text, true) < ldThreshold {
				result = append(result, s)
			}
		}
		return result
	}
	return nil
}

func (s *shell) run() {
	prompt.New(
		s.executor,
		s.suggestor,
		prompt.OptionPrefix(">>> "),
		prompt.OptionLivePrefix(func() (string, bool) {
			cfg, err := config.Read()
			if err != nil {
				return "", false
			}
			activeContext, _, err := cfg.ActiveContext()
			if err != nil {
				return "", false
			}
			return fmt.Sprintf("context:(%s) >>> ", activeContext), true
		}),
	).Run()
}

// Run runs a prompt, it does not return.
func Run(rootCmd *cobra.Command) {
	newShell(rootCmd).run()
}

// ld computes the Levenshtein Distance for two strings.
func ld(s, t string, ignoreCase bool) int {
	if ignoreCase {
		s = strings.ToLower(s)
		t = strings.ToLower(t)
	}
	d := make([][]int, len(s)+1)
	for i := range d {
		d[i] = make([]int, len(t)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {
{'fraction_non_alphanumeric': 0.1566100667070952, 'fraction_numerical': 0.004548211036992117, 'mean_word_length': 3.503071672354949, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '915442', 'n_tokens_mistral': 2712, 'n_tokens_neox': 2439, 'n_words': 553}