grokkit/cmd/edit.go

130 lines
3.1 KiB
Go

package cmd
import (
"fmt"
"os"
"path/filepath"
"strings"
"github.com/spf13/viper"
"gmgauthier.com/grokkit/internal/git"
"github.com/fatih/color"
"github.com/spf13/cobra"
"gmgauthier.com/grokkit/internal/grok"
)
var editCmd = &cobra.Command{
Use: "edit FILE INSTRUCTION",
Short: "Edit a file in-place with Grok",
Run: func(cmd *cobra.Command, args []string) {
if len(args) < 2 {
color.Red("Usage: grokkit edit <file> <instruction>")
os.Exit(1)
}
filePath := args[0]
instruction := strings.Join(args[1:], " ")
content, err := os.ReadFile(filePath)
if err != nil {
color.Red("Error reading file %s: %v", filePath, err)
os.Exit(1)
}
backupPath := filePath + ".bak"
if backupErr := os.WriteFile(backupPath, content, 0o644); backupErr != nil {
color.Yellow("Warning: backup failed %s: %v", backupPath, backupErr)
}
model := viper.GetString("model")
client := grok.NewClient()
systemPrompt := `You are an expert software engineer.
Your task is to edit code files based on instructions.
IMPORTANT: Respond ONLY with the COMPLETE new file content.
No explanations, no markdown, no comments outside the code.`
userPrompt := fmt.Sprintf(`File: %s
Current content:
%s
Instruction: %s
Output ONLY the full updated file content:`, filePath, content, instruction)
messages := []map[string]string{
{"role": "system", "content": systemPrompt},
{"role": "user", "content": userPrompt},
}
cleanCodeResponse := func(s string) string {
lines := strings.Split(s, "\n")
if len(lines) > 0 {
first := strings.TrimSpace(lines[0])
if strings.HasPrefix(first, "```") {
lines = lines[1:]
}
}
if len(lines) > 0 {
last := strings.TrimSpace(lines[len(lines)-1])
if strings.HasSuffix(last, "```") {
lines = lines[:len(lines)-1]
}
}
return strings.TrimSpace(strings.Join(lines, "\n"))
}
newContent := cleanCodeResponse(client.Stream(messages, model))
dir := filepath.Dir(filePath)
base := filepath.Base(filePath)
tmpA, err := os.CreateTemp(dir, "grokkit."+base+".old~")
if err != nil {
color.Red("Temp error: %v", err)
return
}
defer os.Remove(tmpA.Name())
tmpA.Write(content)
tmpA.Close()
tmpB, err := os.CreateTemp(dir, "grokkit."+base+".new~")
if err != nil {
color.Red("Temp error: %v", err)
return
}
defer os.Remove(tmpB.Name())
tmpB.Write([]byte(newContent))
tmpB.Close()
diffOut := git.Run([]string{"diff", "--no-index", "--no-color", tmpA.Name(), tmpB.Name()})
fmt.Println(color.YellowString("📄 Backup: %s", backupPath))
fmt.Println()
if diffOut != "" {
fmt.Println(color.YellowString("🔄 Proposed changes:"))
fmt.Print(color.CyanString(diffOut))
} else {
color.Yellow("No changes.")
}
fmt.Println()
fmt.Print(color.YellowString("Apply? (y/N): "))
var confirm string
fmt.Scanln(&confirm)
if strings.ToLower(strings.TrimSpace(confirm)) != "y" {
color.Yellow("Aborted.")
return
}
if err := os.WriteFile(filePath, []byte(newContent), 0644); err != nil {
color.Red("Error writing file: %v", err)
os.Exit(1)
}
color.Green("File %s updated successfully!", filePath)
},
}