grokkit/cmd/docs.go

package cmd

import (
	"fmt"
	"os"
	"strings"

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

var autoApply bool

var docsCmd = &cobra.Command{
	Use:   "docs <file> [file...]",
	Short: "Generate documentation comments for source files",
	Long: `Detects the programming language of each file and uses Grok AI to generate
language-appropriate documentation comments.

Supported doc styles:
  Go          godoc (// FuncName does...)
  Python      PEP 257 docstrings
  C/C++       Doxygen (/** @brief ... */)
  JS/TS       JSDoc (/** @param ... */)
  Rust        rustdoc (/// Summary)
  Ruby        YARD (# @param ...)
  Java        Javadoc (/** @param ... */)
  Shell       Shell comments (# function: ...)

Safety features:
- Creates .bak backup before modifying files
- Shows preview of changes before applying
- Requires confirmation unless --auto-apply is used`,
	Args: cobra.MinimumNArgs(1),
	Run:  runDocs,
}

func init() {
	docsCmd.Flags().BoolVar(&autoApply, "auto-apply", false, "Apply documentation without confirmation")
}

func runDocs(cmd *cobra.Command, args []string) {
	modelFlag, _ := cmd.Flags().GetString("model")
	model := config.GetModel("docs", modelFlag)

	client := grok.NewClient()
	for _, filePath := range args {
		processDocsFile(client, model, filePath)
	}
}

func processDocsFile(client *grok.Client, model, filePath string) {
	logger.Info("starting docs operation", "file", filePath)

	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		logger.Error("file not found", "file", filePath)
		color.Red("❌ File not found: %s", filePath)
		return
	}

	lang, err := linter.DetectLanguage(filePath)
	if err != nil {
		logger.Warn("unsupported language", "file", filePath, "error", err)
		color.Yellow("⚠️  Skipping %s: %v", filePath, err)
		return
	}

	originalContent, err := os.ReadFile(filePath)
	if err != nil {
		logger.Error("failed to read file", "file", filePath, "error", err)
		color.Red("❌ Failed to read file: %v", err)
		return
	}

	color.Cyan("📝 Generating %s docs for: %s", lang.Name, filePath)
	logger.Info("requesting AI documentation", "file", filePath, "language", lang.Name)

	messages := buildDocsMessages(lang.Name, string(originalContent))
	response := client.StreamSilent(messages, model)

	if response == "" {
		logger.Error("AI returned empty response", "file", filePath)
		color.Red("❌ Failed to get AI response")
		return
	}

	documented := grok.CleanCodeResponse(response)
	logger.Info("received AI documentation", "file", filePath, "size", len(documented))

	// Show preview
	fmt.Println()
	color.Cyan("📋 Preview of documented code:")
	fmt.Println(strings.Repeat("-", 80))
	lines := strings.Split(documented, "\n")
	previewLines := 50
	if len(lines) < previewLines {
		previewLines = len(lines)
	}
	for i := 0; i < previewLines; i++ {
		fmt.Println(lines[i])
	}
	if len(lines) > previewLines {
		fmt.Printf("... (%d more lines)\n", len(lines)-previewLines)
	}
	fmt.Println(strings.Repeat("-", 80))
	fmt.Println()

	if !autoApply {
		color.Yellow("Apply documentation to %s? (y/N): ", filePath)
		var confirm string
		if _, err := fmt.Scanln(&confirm); err != nil {
			color.Yellow("❌ Cancelled. No changes made to: %s", filePath)
			return
		}
		confirm = strings.ToLower(strings.TrimSpace(confirm))
		if confirm != "y" && confirm != "yes" {
			logger.Info("user cancelled docs application", "file", filePath)
			color.Yellow("❌ Cancelled. No changes made to: %s", filePath)
			return
		}
	}

	// Create backup
	backupPath := filePath + ".bak"
	if err := os.WriteFile(backupPath, originalContent, 0644); err != nil {
		logger.Error("failed to create backup", "file", filePath, "error", err)
		color.Red("❌ Failed to create backup: %v", err)
		return
	}
	logger.Info("backup created", "backup", backupPath)

	if err := os.WriteFile(filePath, []byte(documented), 0644); err != nil {
		logger.Error("failed to write documented file", "file", filePath, "error", err)
		color.Red("❌ Failed to write file: %v", err)
		return
	}

	logger.Info("documentation applied successfully", "file", filePath)
	color.Green("✅ Documentation applied: %s", filePath)
	color.Cyan("💾 Original saved to: %s", backupPath)
}

func buildDocsMessages(language, code string) []map[string]string {
	style := docStyle(language)
	systemPrompt := fmt.Sprintf(
		"You are a documentation expert. Add %s documentation comments to the provided code. "+
			"Return ONLY the documented code with no explanations, markdown, or extra text. "+
			"Do NOT include markdown code fences. Document all public functions, methods, types, and constants.",
		style,
	)

	userPrompt := fmt.Sprintf("Add documentation comments to the following %s code:\n\n%s", language, code)

	return []map[string]string{
		{"role": "system", "content": systemPrompt},
		{"role": "user", "content": userPrompt},
	}
}

func docStyle(language string) string {
	switch strings.ToLower(language) {
	case "go":
		return "godoc"
	case "python":
		return "PEP 257 docstring"
	case "c", "c++":
		return "Doxygen"
	case "javascript", "typescript":
		return "JSDoc"
	case "rust":
		return "rustdoc"
	case "ruby":
		return "YARD"
	case "java":
		return "Javadoc"
	case "shell", "bash":
		return "shell comment"
	default:
		return "standard documentation comment"
	}
}