From 3e93a3e9b2c4d209b0a14f05e8ee2a13d3e99033 Mon Sep 17 00:00:00 2001
From: Jeff Wong <awole20@gmail.com>
Date: Mon, 2 Sep 2024 18:54:14 -0700
Subject: [PATCH] FEATURE: add autocomplete (#857)

---
 launcher_go/v2/main.go                   | 14 +++++++
 launcher_go/v2/utils/find_config.go      | 51 ++++++++++++++++++++++++
 launcher_go/v2/utils/find_config_test.go | 39 ++++++++++++++++++
 3 files changed, 104 insertions(+)
 create mode 100644 launcher_go/v2/utils/find_config.go
 create mode 100644 launcher_go/v2/utils/find_config_test.go

diff --git a/launcher_go/v2/main.go b/launcher_go/v2/main.go
index 0c8b669..55cd241 100644
--- a/launcher_go/v2/main.go
+++ b/launcher_go/v2/main.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"github.com/alecthomas/kong"
 	"github.com/discourse/discourse_docker/launcher_go/v2/utils"
+	"github.com/posener/complete"
+	"github.com/willabides/kongplete"
 	"golang.org/x/sys/unix"
 	"os"
 	"os/exec"
@@ -30,14 +32,26 @@ type Cli struct {
 	StopCmd    StopCmd    `cmd:"" name:"stop" help:"Stops container."`
 	RestartCmd RestartCmd `cmd:"" name:"restart" help:"Stops then starts container."`
 	RebuildCmd RebuildCmd `cmd:"" name:"rebuild" help:"Builds new image, then destroys old container, and starts new container."`
+
+	InstallCompletions kongplete.InstallCompletions `cmd:"" aliases:"sh" help:"Print shell autocompletions. Add output to dotfiles, or 'source <(./launcher2 sh)'."`
 }
 
 func main() {
 	cli := Cli{}
 	runCtx, cancel := context.WithCancel(context.Background())
 
+	// pre parse to get config dir for prediction of conf dir
+	confFiles := utils.FindConfigNames()
+
 	parser := kong.Must(&cli, kong.UsageOnError(), kong.Bind(&runCtx), kong.Vars{"version": utils.Version})
 
+	// Run kongplete.Complete to handle completion requests
+	kongplete.Complete(parser,
+		kongplete.WithPredictor("config", complete.PredictSet(confFiles...)),
+		kongplete.WithPredictor("file", complete.PredictFiles("*")),
+		kongplete.WithPredictor("dir", complete.PredictDirs("*")),
+	)
+
 	ctx, err := parser.Parse(os.Args[1:])
 	parser.FatalIfErrorf(err)
 
diff --git a/launcher_go/v2/utils/find_config.go b/launcher_go/v2/utils/find_config.go
new file mode 100644
index 0000000..85a1cae
--- /dev/null
+++ b/launcher_go/v2/utils/find_config.go
@@ -0,0 +1,51 @@
+package utils
+
+import (
+	"bytes"
+	"flag"
+	"io/ioutil"
+	"os"
+	"strings"
+)
+
+// Find config names for autocomplete, given the current --conf-dir argument.
+func FindConfigNames() []string {
+	compLine := os.Getenv("COMP_LINE")
+	flagLine := []string{}
+	found := false
+	// the flag package wants a valid flag first
+	// drop all COMP_LINE args until we find something that starts with --conf-dir
+	for _, s := range strings.Fields(compLine) {
+		if found {
+			flagLine = append(flagLine, s)
+		}
+		if strings.HasPrefix(s, "--conf-dir") {
+			flagLine = append(flagLine, s)
+			found = true
+		}
+	}
+	flags := flag.NewFlagSet("f", flag.ContinueOnError)
+	//squelch helptext
+	flags.SetOutput(&bytes.Buffer{})
+	confDirArg := flags.String("conf-dir", "./containers", "conf dir")
+	flags.Parse(flagLine)
+
+	// search in the current conf dir for any files
+	confDir := strings.TrimRight(*confDirArg, "/") + "/"
+	confFiles := []string{}
+	files, err := ioutil.ReadDir(confDir)
+	if err == nil {
+		for _, file := range files {
+			if !file.IsDir() {
+				if strings.HasSuffix(file.Name(), ".yml") {
+					confName, _ := strings.CutSuffix(file.Name(), ".yml")
+					confFiles = append(confFiles, confName)
+				} else if strings.HasSuffix(file.Name(), ".yaml") {
+					confName, _ := strings.CutSuffix(file.Name(), ".yaml")
+					confFiles = append(confFiles, confName)
+				}
+			}
+		}
+	}
+	return confFiles
+}
diff --git a/launcher_go/v2/utils/find_config_test.go b/launcher_go/v2/utils/find_config_test.go
new file mode 100644
index 0000000..1d6340d
--- /dev/null
+++ b/launcher_go/v2/utils/find_config_test.go
@@ -0,0 +1,39 @@
+package utils_test
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/discourse/discourse_docker/launcher_go/v2/utils"
+	"os"
+)
+
+var _ = Describe("FindConfig", func() {
+
+	It("Parses and returns yml or yaml files", func() {
+		os.Setenv("COMP_LINE", "launcher2 build --conf-dir ../test/containers")
+		Expect(utils.FindConfigNames()).To(ContainElements("test", "test2"))
+	})
+
+	It("Parses and returns yml or yaml files with trailing slash", func() {
+		os.Setenv("COMP_LINE", "launcher2 build --conf-dir ../test/containers/")
+		Expect(utils.FindConfigNames()).To(ContainElements("test", "test2"))
+	})
+
+	It("Parses and returns yml or yaml files on equals", func() {
+		os.Setenv("COMP_LINE", "launcher2 --conf-dir=../test/containers other args")
+		Expect(utils.FindConfigNames()).To(ContainElements("test", "test2"))
+	})
+
+	It("doesn't error when dir does not exist when set", func() {
+		os.Setenv("COMP_LINE", "launcher2 --conf-dir=./does-not-exist")
+		Expect(utils.FindConfigNames()).To(BeEmpty())
+	})
+
+	It("doesn't error when dir does not exist", func() {
+		//by default it look is in ./containers directory, which does not exist
+		// in this directory
+		os.Setenv("COMP_LINE", "launcher2")
+		Expect(utils.FindConfigNames()).To(BeEmpty())
+	})
+})
-- 
2.25.1