grokkit/cmd/chat.go

package cmd

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

type ChatHistory struct {
	Messages []map[string]string `json:"messages"`
}

type ToolCall struct {
	Tool        string `json:"tool"`
	File        string `json:"file,omitempty"`
	Path        string `json:"path,omitempty"`
	Instruction string `json:"instruction,omitempty"`
	Description string `json:"description,omitempty"`
	Message     string `json:"message,omitempty"`
}

func loadChatHistory() []map[string]string {
	histFile := getChatHistoryFile()
	data, err := os.ReadFile(histFile)
	if err != nil {
		return nil
	}

	var hist ChatHistory
	if err := json.Unmarshal(data, &hist); err != nil {
		return nil
	}
	return hist.Messages
}

func saveChatHistory(messages []map[string]string) error {
	histFile := getChatHistoryFile()
	hist := ChatHistory{Messages: messages}
	data, err := json.MarshalIndent(hist, "", " ")
	if err != nil {
		return err
	}
	return os.WriteFile(histFile, data, 0644)
}

func getChatHistoryFile() string {
	configFile := viper.GetString("chat.history_file")
	if configFile != "" {
		return configFile
	}
	home, _ := os.UserHomeDir()
	if home == "" {
		home = "."
	}
	histDir := filepath.Join(home, ".config", "grokkit")
	_ = os.MkdirAll(histDir, 0755)
	return filepath.Join(histDir, "chat_history.json")
}

var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Interactive chat with Grok (use --agent for tool-enabled mode)",
	Long: `Start a persistent conversation with Grok.

Normal mode: coherent, reasoning-focused chat (uses full model).
Agent mode (--agent): Grok can call tools with safety previews (uses fast model).`,
	Run: runChat,
}

func init() {
	chatCmd.Flags().Bool("agent", false, "Enable agent mode with tool calling (uses fast non-reasoning model)")
	chatCmd.Flags().String("model", "", "Override model")
	rootCmd.AddCommand(chatCmd)
}

func runChat(cmd *cobra.Command, args []string) {
	agentMode, _ := cmd.Flags().GetBool("agent")
	modelFlag, _ := cmd.Flags().GetString("model")

	var model string
	if modelFlag != "" {
		model = modelFlag
	} else if agentMode {
		model = config.GetModel("chat-agent", "")
	} else {
		model = config.GetModel("chat", "")
	}

	client := grok.NewClient()

	systemPrompt := map[string]string{
		"role":    "system",
		"content": fmt.Sprintf("You are Grok 4, the latest and most powerful model from xAI (2026). You are currently running as `%s`. Be helpful, truthful, and a little irreverent.", model),
	}

	if agentMode {
		systemPrompt["content"] = "You are Grok in Agent Mode.\n" +
			"You can call tools using this exact JSON format inside ```tool blocks:\n\n" +
			"```tool\n" +
			"{\"tool\": \"edit\", \"file\": \"main.go\", \"instruction\": \"Add error handling\"}\n" +
			"```\n\n" +
			"Available tools: edit, scaffold, testgen, lint, commit.\n\n" +
			"Always use tools when the user asks you to change code, generate tests, lint, or commit.\n" +
			"After every tool call, wait for the result before continuing.\n" +
			"Be concise and action-oriented."
	}

	history := loadChatHistory()
	if len(history) == 0 {
		history = []map[string]string{systemPrompt}
	} else if history[0]["role"] != "system" {
		history = append([]map[string]string{systemPrompt}, history...)
	} else {
		history[0] = systemPrompt
	}

	color.Cyan("┌──────────────────────────────────────────────────────────────┐")
	color.Cyan("│ Grokkit Chat %s — Model: %s │", map[bool]string{true: "(Agent Mode)", false: ""}[agentMode], model)
	color.Cyan("│ Type /quit to exit │")
	color.Cyan("└──────────────────────────────────────────────────────────────┘\n")

	scanner := bufio.NewScanner(os.Stdin)
	for {
		color.Yellow("You > ")
		if !scanner.Scan() {
			break
		}

		input := strings.TrimSpace(scanner.Text())
		if input == "" {
			continue
		}
		if input == "/quit" || input == "/q" || input == "exit" {
			color.Cyan("\nGoodbye 👋\n")
			break
		}

		history = append(history, map[string]string{"role": "user", "content": input})

		color.Green("Grok > ")
		reply := client.Stream(history, model)

		// Agent tool call handling
		if agentMode && strings.Contains(reply, "```tool") {
			handleToolCall(reply, &history)
			continue
		}

		history = append(history, map[string]string{"role": "assistant", "content": reply})
		_ = saveChatHistory(history)

		fmt.Println()
	}
}

func handleToolCall(reply string, history *[]map[string]string) {
	start := strings.Index(reply, "```tool")
	if start == -1 {
		return
	}
	end := strings.Index(reply[start+7:], "```")
	if end == -1 {
		return
	}
	block := strings.TrimSpace(reply[start+7 : start+7+end])

	var tc ToolCall
	if err := json.Unmarshal([]byte(block), &tc); err != nil {
		color.Red("Failed to parse tool call: %v", err)
		return
	}

	color.Yellow("\n[Agent] Grok wants to call tool: %s", tc.Tool)

	switch tc.Tool {
	case "edit":
		if tc.File != "" && tc.Instruction != "" {
			// Reuse the existing edit command's logic (preview + confirm)
			editCmd.RunE = func(cmd *cobra.Command, args []string) error {
				RunEditWithInstruction(tc.File, tc.Instruction) // we'll define this helper next
				return nil
			}
			_ = editCmd.Execute()
		}
	case "scaffold":
		if tc.Path != "" && tc.Description != "" {
			scaffoldCmd.RunE = func(cmd *cobra.Command, args []string) error {
				RunScaffoldWithDescription(tc.Path, tc.Description)
				return nil
			}
			_ = scaffoldCmd.Execute()
		}
	case "testgen":
		if tc.File != "" {
			testgenCmd.RunE = func(cmd *cobra.Command, args []string) error {
				RunTestgenWithFile(tc.File)
				return nil
			}
			_ = testgenCmd.Execute()
		}
	case "lint":
		if tc.File != "" {
			lintCmd.RunE = func(cmd *cobra.Command, args []string) error {
				RunLintWithFile(tc.File)
				return nil
			}
			_ = lintCmd.Execute()
		}
	case "commit":
		if tc.Message != "" {
			_, _ = git.Run([]string{"commit", "-m", tc.Message})
		}
	default:
		color.Red("Unknown tool: %s", tc.Tool)
	}

	*history = append(*history, map[string]string{"role": "assistant", "content": reply})
}