From a870ad858c571a201be7a15f49401ac89b75347b Mon Sep 17 00:00:00 2001
From: Jay Kumar <70096901+35C4n0r@users.noreply.github.com>
Date: Mon, 11 Mar 2024 14:05:56 +0530
Subject: [PATCH] feat: auto inject completions (#139)

* feat: auto inject completions

* fix: remove duplicate source lines

* chore: simplify string matching

* ref: change command name
---
 pkg/cmd/autocomplete.go | 104 ++++++++++++++++++++++++++++++++++++++++
 pkg/cmd/cmd.go          |   2 +
 2 files changed, 106 insertions(+)
 create mode 100644 pkg/cmd/autocomplete.go

diff --git a/pkg/cmd/autocomplete.go b/pkg/cmd/autocomplete.go
new file mode 100644
index 000000000..50d22fd45
--- /dev/null
+++ b/pkg/cmd/autocomplete.go
@@ -0,0 +1,104 @@
+// Copyright 2024 Daytona Platforms Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+	"github.com/spf13/cobra"
+)
+
+var AutoCompleteCmd = &cobra.Command{
+	Use:     "autocomplete [bash|zsh|fish|powershell]",
+	Short:   "Adds completion script for your shell enviornment",
+	Args:    cobra.ExactArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		shell := args[0]
+		homeDir, err := os.UserHomeDir()
+		if err != nil {
+			fmt.Printf("Error finding user home directory: %s\n", err)
+			return
+		}
+
+		var filePath, profilePath string
+		switch shell {
+		case "bash":
+			filePath = filepath.Join(homeDir, ".daytona.completion_script.bash")
+			profilePath = filepath.Join(homeDir, ".bashrc")
+		case "zsh":
+			filePath = filepath.Join(homeDir, ".daytona.completion_script.zsh")
+			profilePath = filepath.Join(homeDir, ".zshrc")
+		case "fish":
+			filePath = filepath.Join(homeDir, ".config", "fish", "daytona.completion_script.fish")
+			profilePath = filepath.Join(homeDir, ".config", "fish", "config.fish")
+		case "powershell":
+			filePath = filepath.Join(homeDir, "daytona.completion_script.ps1")
+			profilePath = filepath.Join(homeDir, "Documents", "WindowsPowerShell", "Microsoft.PowerShell_profile.ps1")
+		default:
+			fmt.Println("Unsupported shell type. Please use bash, zsh, fish, or powershell.")
+			return
+		}
+
+		file, err := os.Create(filePath)
+		if err != nil {
+			fmt.Printf("Error creating completion script file: %s\n", err)
+			return
+		}
+		defer file.Close()
+
+		switch shell {
+		case "bash":
+			err = cmd.Root().GenBashCompletion(file)
+		case "zsh":
+			err = cmd.Root().GenZshCompletion(file)
+		case "fish":
+			err = cmd.Root().GenFishCompletion(file, true)
+		case "powershell":
+			err = cmd.Root().GenPowerShellCompletionWithDesc(file)
+		}
+
+		if err != nil {
+			fmt.Printf("Error generating completion script: %s\n", err)
+			return
+		}
+
+		sourceCommand := fmt.Sprintf("\nsource %s\n", filePath)
+		if shell == "powershell" {
+			sourceCommand = fmt.Sprintf(". %s\n", filePath)
+		}
+
+		alreadyPresent := false
+		// Read existing content from the file
+		profile, err := os.ReadFile(profilePath)
+
+		if err != nil && !os.IsNotExist(err) {
+			fmt.Printf("Error while reading profile (%s): %s\n", profilePath, err)
+		}
+	
+		if strings.Contains(string(profile), strings.TrimSpace(sourceCommand)) {
+			alreadyPresent = true
+		}
+
+		if !alreadyPresent {
+			// Append the source command to the shell's profile file if not present
+			profile, err := os.OpenFile(profilePath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
+			if err != nil {
+				fmt.Printf("Error opening profile file (%s): %s\n", profilePath, err)
+				return
+			}
+			defer profile.Close()
+
+			if _, err := profile.WriteString(sourceCommand); err != nil {
+				fmt.Printf("Error writing to profile file (%s): %s\n", profilePath, err)
+				return
+			}
+		}
+
+		fmt.Println("Autocomplete script generated and injected successfully.")
+		fmt.Printf("Please source your %s profile to apply the changes or restart your terminal.\n", shell)
+		fmt.Printf("For manual sourcing, use: source %s\n", profilePath)
+	},
+}
\ No newline at end of file
diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go
index f1bd0a3f2..e0ca12613 100644
--- a/pkg/cmd/cmd.go
+++ b/pkg/cmd/cmd.go
@@ -36,6 +36,7 @@ var rootCmd = &cobra.Command{
 var originalStdout *os.File
 
 func Execute() {
+	rootCmd.AddCommand(AutoCompleteCmd)
 	rootCmd.AddCommand(InfoCmd)
 	rootCmd.AddCommand(StartCmd)
 	rootCmd.AddCommand(StopCmd)
@@ -62,6 +63,7 @@ func Execute() {
 		rootCmd.AddCommand(ProviderCmd)
 	}
 
+	rootCmd.CompletionOptions.HiddenDefaultCmd = true
 	rootCmd.PersistentFlags().BoolP("help", "", false, "help for daytona")
 	rootCmd.PersistentFlags().StringVarP(&output.FormatFlag, "output", "o", output.FormatFlag, `Output format. Must be one of (yaml, json)`)