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/agent"
	"gmgauthier.com/grokkit/internal/grok"
)

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

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")

	// Model switching logic
	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 // refresh system prompt
	}

	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 mode: check for tool calls
		if agentMode && strings.Contains(reply, "```tool") {
			agent.HandleToolCall(reply, &history, model)
			continue
		}

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

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

		fmt.Println()
	}
}