grokkit/cmd/agent.go

package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"gmgauthier.com/grokkit/config"
	"gmgauthier.com/grokkit/internal/grok"
)

var agentCmd = &cobra.Command{
	Use:   "agent INSTRUCTION",
	Short: "Multi-file agent — Grok intelligently edits multiple files with preview",
	Args:  cobra.ExactArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		instruction := args[0]
		modelFlag, _ := cmd.Flags().GetString("model")
		model := config.GetModel(modelFlag)

		client := grok.NewClient()

		color.Yellow("🔍 Agent mode activated. Scanning project...")

		var files []string
		err := filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
			if err != nil || info.IsDir() {
				return err
			}
			if strings.HasSuffix(path, ".go") {
				files = append(files, path)
			}
			return nil
		})
		if err != nil {
			color.Red("Failed to scan project: %v", err)
			os.Exit(1)
		}

		if len(files) == 0 {
			color.Yellow("No .go files found.")
			return
		}

		color.Yellow("📄 Found %d files. Asking Grok for a plan...", len(files))

		planMessages := []map[string]string{
			{"role": "system", "content": "You are an expert software engineer. Given an instruction and list of files, return a clear plan: which files to change and a brief description of what to do in each. List files one per line."},
			{"role": "user", "content": fmt.Sprintf("Instruction: %s\n\nFiles:\n%s", instruction, strings.Join(files, "\n"))},
		}

		plan := client.Stream(planMessages, model)
		color.Cyan("\nGrok's Plan:\n%s", plan)

		fmt.Print("\nProceed with changes? (y/n): ")
		var confirm string
		if _, err := fmt.Scanln(&confirm); err != nil {
			color.Red("Failed to read input: %v", err)
			return
		}
		if confirm != "y" && confirm != "Y" {
			color.Yellow("Aborted.")
			return
		}

		color.Yellow("Applying changes file by file...\n")

		applyAll := false
		for i, file := range files {
			color.Yellow("[%d/%d] → %s", i+1, len(files), file)

			original, err := os.ReadFile(file)
			if err != nil {
				color.Red("Could not read %s", file)
				continue
			}

			backupPath := file + ".bak"
			_ = os.WriteFile(backupPath, original, 0644)

			messages := []map[string]string{
				{"role": "system", "content": "You are an expert programmer. Remove all unnecessary comments including last modified timestamps and ownership comments. Return clean, complete file content with no explanations, markdown, diffs, or extra text."},
				{"role": "user", "content": fmt.Sprintf("File: %s\n\nOriginal content:\n%s\n\nTask: %s", filepath.Base(file), original, instruction)},
			}

			raw := client.StreamSilent(messages, "grok-4-1-fast-non-reasoning")
			newContent := grok.CleanCodeResponse(raw)

			color.Cyan("Proposed changes for %s:", filepath.Base(file))
			fmt.Println("--- a/" + filepath.Base(file))
			fmt.Println("+++ b/" + filepath.Base(file))
			fmt.Print(newContent)

			if !applyAll {
				fmt.Print("\nApply this file? (y/n/a = all remaining): ")
				var answer string
				if _, err := fmt.Scanln(&answer); err != nil {
					color.Yellow("Skipped %s (failed to read input)", file)
					continue
				}
				answer = strings.ToLower(strings.TrimSpace(answer))

				if answer == "a" {
					applyAll = true
				} else if answer != "y" {
					color.Yellow("Skipped %s (backup kept)", file)
					continue
				}
			}

			_ = os.WriteFile(file, []byte(newContent), 0644)
			color.Green("✅ Applied %s", file)
		}

		color.Green("\n🎉 Agent mode complete! All changes applied.")
	},
}

// CleanCodeResponse removes markdown fences and returns pure code content
func CleanCodeResponse(text string) string {
	fence := "```"

	// Remove any line that starts with the fence (opening fence, possibly with language tag)
	text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`.*$`).ReplaceAllString(text, "")

	// Remove any line that is just the fence (closing fence)
	text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`\s*$`).ReplaceAllString(text, "")

	// Trim only leading and trailing whitespace.
	// Do NOT collapse internal blank lines — they are intentional in code.
	text = strings.TrimSpace(text)

	return text
}