grokkit/cmd/scaffold.go

package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

var scaffoldCmd = &cobra.Command{
	Use:   "scaffold FILE DESCRIPTION",
	Short: "Scaffold a new file with Grok (safe preview + confirmation)",
	Args:  cobra.ExactArgs(2),
	Run: func(cmd *cobra.Command, args []string) {
		filePath := args[0]
		description := args[1]

		withTests, _ := cmd.Flags().GetBool("with-tests")
		dryRun, _ := cmd.Flags().GetBool("dry-run")
		force, _ := cmd.Flags().GetBool("force")
		yes, _ := cmd.Flags().GetBool("yes")
		langOverride, _ := cmd.Flags().GetString("lang")

		modelFlag, _ := cmd.Flags().GetString("model")
		model := config.GetModel("scaffold", modelFlag)

		logger.Info("scaffold command started",
			"file", filePath,
			"description", description,
			"with_tests", withTests,
			"model", model)

		// Safety: don't overwrite existing file unless --force
		if _, err := os.Stat(filePath); err == nil && !force {
			color.Red("File already exists: %s (use --force to overwrite)", filePath)
			os.Exit(1)
		}

		dir := filepath.Dir(filePath)
		if dir != "." && dir != "" {
			if err := os.MkdirAll(dir, 0755); err != nil {
				logger.Error("failed to create directory", "dir", dir, "error", err)
				color.Red("Failed to create directory: %v", err)
				os.Exit(1)
			}
		}

		// Detect language and harvest context
		lang := detectLanguage(filePath, langOverride)
		context := harvestContext(filePath, lang)

		// Build system prompt with style enforcement
		systemPrompt := fmt.Sprintf(`You are an expert %s programmer.
Match the exact style, naming conventions, error handling, logging, and package structure of this project.
Return ONLY the complete code file. No explanations, no markdown, no backticks.`, lang)

		if withTests {
			systemPrompt += "\nAlso generate a basic _test.go file with at least one test when --with-tests is used."
		}

		messages := []map[string]string{
			{"role": "system", "content": systemPrompt},
			{"role": "user", "content": fmt.Sprintf(
				"File path: %s\n\nProject context:\n%s\n\nCreate this new file:\n%s",
				filePath, context, description)},
		}

		color.Yellow("Asking Grok to scaffold %s...\n", filepath.Base(filePath))
		client := grok.NewClient()
		raw := client.StreamSilent(messages, model)
		newContent := grok.CleanCodeResponse(raw)
		color.Green("Response received")

		// Preview
		color.Cyan("\nProposed new file:")
		fmt.Printf("--- /dev/null\n")
		fmt.Printf("+++ b/%s\n", filepath.Base(filePath))
		fmt.Println(newContent)

		if dryRun {
			color.Yellow("\n--dry-run: file not written")
			return
		}

		if !yes {
			fmt.Print("\n\nCreate this file? (y/n): ")
			var confirm string
			if _, err := fmt.Scanln(&confirm); err != nil || (confirm != "y" && confirm != "Y") {
				color.Yellow("Scaffold cancelled.")
				return
			}
		}

		// Write main file
		if err := os.WriteFile(filePath, []byte(newContent), 0644); err != nil {
			logger.Error("failed to write file", "file", filePath, "error", err)
			color.Red("Failed to write file: %v", err)
			os.Exit(1)
		}
		color.Green("✓ Created: %s", filePath)

		// Optional test file
		if withTests {
			testPath := strings.TrimSuffix(filePath, filepath.Ext(filePath)) + "_test.go"
			testMessages := []map[string]string{
				{"role": "system", "content": systemPrompt},
				{"role": "user", "content": fmt.Sprintf("Generate a basic test file for %s using the same style.", filepath.Base(filePath))},
			}
			testRaw := client.StreamSilent(testMessages, model)
			testContent := grok.CleanCodeResponse(testRaw)

			if err := os.WriteFile(testPath, []byte(testContent), 0644); err == nil {
				color.Green("✓ Created test: %s", filepath.Base(testPath))
			}
		}

		logger.Info("scaffold completed successfully", "file", filePath, "with_tests", withTests)
	},
}

// Simple language detector (can be moved to internal/linter later)
func detectLanguage(path, override string) string {
	if override != "" {
		return override
	}
	ext := strings.ToLower(filepath.Ext(path))
	switch ext {
	case ".go":
		return "Go"
	case ".py":
		return "Python"
	case ".js", ".ts":
		return "TypeScript"
	case ".c":
		return "C"
	case ".cpp":
		return "C++"
	case ".java":
		return "Java"
	// add more as needed
	default:
		return "code"
	}
}

// Basic context harvester (~4000 token cap)
func harvestContext(filePath, lang string) string {
	var sb strings.Builder
	dir := filepath.Dir(filePath)

	// Siblings
	files, _ := os.ReadDir(dir)
	for _, f := range files {
		if f.IsDir() || strings.HasPrefix(f.Name(), ".") {
			continue
		}
		if filepath.Ext(f.Name()) == filepath.Ext(filePath) && f.Name() != filepath.Base(filePath) {
			content, _ := os.ReadFile(filepath.Join(dir, f.Name()))
			if len(content) > 2000 {
				content = content[:2000]
			}
			sb.WriteString(fmt.Sprintf("=== %s ===\n%s\n\n", f.Name(), string(content)))
		}
	}

	// Rough token cap
	if sb.Len() > 4000 {
		return sb.String()[:4000] + "\n... (truncated)"
	}
	return sb.String()
}

func init() {
	scaffoldCmd.Flags().Bool("with-tests", false, "Also scaffold a basic _test.go file")
	scaffoldCmd.Flags().Bool("dry-run", false, "Preview only, do not write")
	scaffoldCmd.Flags().BoolP("yes", "y", false, "Skip confirmation prompt")
	scaffoldCmd.Flags().Bool("force", false, "Overwrite existing file")
	scaffoldCmd.Flags().String("lang", "", "Force language for prompt (Go, Python, etc.)")
}