Added support for colored console output and branch creation function…

go.mod

@@ -9,10 +9,13 @@ require (
 )
 
 require (
+	github.com/fatih/color v1.16.0 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.1.0 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect

go.sum

@@ -2,6 +2,8 @@ github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
+github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM=
+github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE=
 github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
 github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
 github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM=
@@ -19,6 +21,11 @@ github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
 github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4=
@@ -59,6 +66,8 @@ go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI=
 go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ=
 golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjsbSXD66ic0XW0js0R9g=
 golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
 golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=

main.go

@@ -4,9 +4,9 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"log"
 	"net/http"
 	"os"
 	"os/exec"
@@ -39,17 +39,19 @@ func main() {
 	// Initialize Cobra
 	var rootCmd = &cobra.Command{Use: "gitpilotai"}
 	rootCmd.AddCommand(generateCmd)
+	rootCmd.AddCommand(branchCmd)
 	cobra.OnInitialize(initConfig)
 	if err := rootCmd.Execute(); err != nil {
-		fmt.Println(err)
+		color.Red("%v", err)
 		os.Exit(1)
 	}
 }
 
 func initConfig() {
 	viper.AutomaticEnv()
 	if viper.GetString("OPENAI_API_KEY") == "" {
-		log.Fatal("OPENAI_API_KEY environment variable not set")
+		color.Red("OPENAI_API_KEY environment variable not set")
+		return
 	}
 }
 
@@ -58,38 +60,43 @@ var generateCmd = &cobra.Command{
 	Short: "Generate commit messages based on git diff",
 	Run: func(cmd *cobra.Command, args []string) {
 		if !hasGitChanges() {
-			fmt.Println("No changes detected in the Git repository.")
+			color.Yellow("Please make some changes to the repository and try again.")
 			return
 		}
 
 		if err := stageFiles(); err != nil {
-			log.Fatalf("Error staging files: %v", err)
+			color.Red("Error staging files: %v", err)
+			return
 		}
 
 		diff := getGitDiff()
 		if diff == "" {
-			fmt.Println("No diff found.")
+			color.Yellow("No diff found.")
 			return
 		}
 
 		commitMessage, err := GenerateDiff(diff)
 		if err != nil {
-			log.Fatalf("Error generating commit message: %v", err)
+			color.Red("Error generating commit message: %v", err)
+			return
 		}
 
-		fmt.Printf("Generated commit message: %s\n", commitMessage)
+		color.Green("Generated commit message: %s\n", commitMessage)
 
 		if err := commitChanges(commitMessage); err != nil {
-			log.Fatalf("Error committing changes: %v", err)
+			color.Red("Error committing changes: %v", err)
+			return
 		}
 
 		currentBranch, err := detectCurrentBranch()
 		if err != nil {
-			log.Fatalf("Error detecting current branch: %v", err)
+			color.Red("Error detecting current branch: %v", err)
+			return
 		}
 
 		if err := pushChanges(currentBranch); err != nil {
-			log.Fatalf("Error pushing changes: %v", err)
+			color.Red("Error pushing changes: %v", err)
+			return
 		}
 	},
 }
@@ -100,53 +107,61 @@ var branchCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		currentBranch, err := detectCurrentBranch()
 		if err != nil {
-			log.Fatalf("Error detecting current branch: %v", err)
+			color.Red("Error detecting current branch: %v", err)
+			return
 		}
 		if currentBranch != "main" && currentBranch != "master" {
-			log.Fatalf("You must be on the main branch to create a new branch.")
+			color.Red("You must be on the main branch to create a new branch.")
+			return
 		}
 		// Add files
 		if !hasGitChanges() {
-			fmt.Println("No changes detected in the Git repository.")
+			color.Yellow("No changes detected in the Git repository.")
 			return
 		}
 
 		if err := stageFiles(); err != nil {
-			log.Fatalf("Error staging files: %v", err)
+			color.Red("Error staging files: %v", err)
+			return
 		}
 
 		diff := getGitDiff()
 		if diff == "" {
-			fmt.Println("No diff found.")
+			color.Yellow("No diff found.")
 			return
 		}
 		// Generate commit message
 		commitMessage, err := GenerateDiff(diff)
 		if err != nil {
-			log.Fatalf("Error generating commit message: %v", err)
+			color.Red("Error generating commit message: %v", err)
+			return
 		}
 		// Create branch
 		branchName := generateBranchNameFromCommitMessage(commitMessage)
 		// Switch to new branch
 		err = checkoutNewBranchLocally(branchName)
 		if err != nil {
-			log.Fatalf("Error checking out new branch: %v", err)
+			color.Red("Error checking out new branch: %v", err)
+			return
 		}
 		// Commit changes
 		if err := commitChanges(commitMessage); err != nil {
-			log.Fatalf("Error committing changes: %v", err)
+			color.Red("Error committing changes: %v", err)
+			return
 		}
 
 		if err := pushChanges(branchName); err != nil {
-			log.Fatalf("Error pushing changes: %v", err)
+			color.Red("Error pushing changes: %v", err)
+			return
 		}
 	},
 }
 
 func hasGitChanges() bool {
 	out, err := exec.Command("git", "status", "--porcelain").Output()
 	if err != nil {
-		log.Fatalf("Error checking git status: %v", err)
+		color.Red("Error checking git status: %v", err)
+		return false
 	}
 	return len(out) > 0
 }
@@ -155,13 +170,15 @@ func getGitDiff() string {
 	stagedDiffCmd := exec.Command("git", "diff", "--staged")
 	stagedDiff, err := stagedDiffCmd.Output()
 	if err != nil {
-		log.Printf("Error getting staged git diff: %v", err)
+		color.Red("Error getting staged git diff: %v", err)
+		return ""
 	}
 
 	unstagedDiffCmd := exec.Command("git", "diff")
 	unstagedDiff, err := unstagedDiffCmd.Output()
 	if err != nil {
-		log.Printf("Error getting unstaged git diff: %v", err)
+		color.Red("Error getting unstaged git diff: %v", err)
+		return ""
 	}
 
 	totalDiff := strings.TrimSpace(string(stagedDiff)) + "\n" + strings.TrimSpace(string(unstagedDiff))
@@ -218,7 +235,7 @@ func commitChanges(commitMessage string) error {
 	if out, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("git commit failed: %s, %v", out, err)
 	}
-	fmt.Println("Changes committed successfully.")
+	color.Green("Commit message: %s\n", commitMessage)
 	return nil
 }
 
@@ -243,7 +260,7 @@ func pushChanges(branchName string) error {
 		return fmt.Errorf("git push failed: %s, %v", out, err)
 	}
 
-	fmt.Printf("Changes pushed successfully to branch %s.\n", branchName)
+	color.Green("Changes pushed successfully to branch %s.\n", branchName)
 	return nil
 }
 
@@ -255,7 +272,7 @@ func stageFiles() error {
 		return fmt.Errorf("error staging files: %s, %v", out, err)
 	}
 
-	fmt.Println("Files staged successfully.")
+	color.Green("Files staged successfully.")
 	return nil
 }
 
@@ -272,7 +289,8 @@ func generateBranchNameFromCommitMessage(commitMessage string) string {
 	prompt := fmt.Sprintf("Generate a branch name from a commit message. The branch name should be in a valid format, e.g. branch-name-of-feature. Here is the commit message:%s", commitMessage)
 	branch, err := makeOpenAPIRequestFromPrompt(prompt)
 	if err != nil {
-		log.Fatalf("Error generating branch name: %v", err)
+		color.Red("Error generating branch name: %v", err)
+		return ""
 	}
 	return branch
 }
@@ -300,6 +318,6 @@ func checkoutNewBranchLocally(branchName string) error {
 	if out, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("error checking out new branch: %s, %v", out, err)
 	}
-	fmt.Printf("Switched to new branch %s\n", branchName)
+	color.Green("Switched to new branch %s\n", branchName)
 	return nil
 }