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...]", 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" } }