grokkit/internal/mcp/server.go

package mcp

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

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"

	"gmgauthier.com/grokkit/internal/analyze"
	"gmgauthier.com/grokkit/internal/docs"
	"gmgauthier.com/grokkit/internal/git"
	"gmgauthier.com/grokkit/internal/grok"
	"gmgauthier.com/grokkit/internal/linter"
	"gmgauthier.com/grokkit/internal/logger"
	"gmgauthier.com/grokkit/internal/prompts"
	"gmgauthier.com/grokkit/internal/recipe"
	"gmgauthier.com/grokkit/internal/testgen"
)

// Server wraps an MCP server and provides Grokkit-specific tools.
type Server struct {
	mcpServer  *server.MCPServer
	grokClient *grok.Client
}

// NewServer creates a new MCP server with Grokkit tools.
func NewServer() (*Server, error) {
	// Create the underlying MCP server
	s := server.NewMCPServer(
		"grokkit",
		"0.1.0",
		server.WithToolCapabilities(false),
		server.WithRecovery(),
	)

	grokClient := grok.NewClient()

	mcpSrv := &Server{
		mcpServer:  s,
		grokClient: grokClient,
	}

	if err := mcpSrv.registerTools(); err != nil {
		return nil, fmt.Errorf("failed to register tools: %w", err)
	}

	mcpSrv.registerResources()

	return mcpSrv, nil
}

// registerTools registers all available Grokkit tools.
func (s *Server) registerTools() error {
	// lint_code tool
	lintTool := mcp.NewTool("lint_code",
		mcp.WithDescription("Run a linter on a source file and return results"),
		mcp.WithString("file_path", mcp.Required(), mcp.Description("Path to the file to lint")),
	)
	s.mcpServer.AddTool(lintTool, s.handleLintCode)

	// analyze_code tool
	analyzeTool := mcp.NewTool("analyze_code",
		mcp.WithDescription("Perform deep project analysis and return an educational Markdown report"),
		mcp.WithString("dir", mcp.Description("Directory to analyze (default: current directory)")),
	)
	s.mcpServer.AddTool(analyzeTool, s.handleAnalyzeCode)

	// generate_docs tool
	docsTool := mcp.NewTool("generate_docs",
		mcp.WithDescription("Generate documentation comments for a source file"),
		mcp.WithString("file_path", mcp.Required(), mcp.Description("Path to the file to document")),
		mcp.WithString("style", mcp.Description("Documentation style (optional)")),
	)
	s.mcpServer.AddTool(docsTool, s.handleGenerateDocs)

	// generate_tests tool
	testsTool := mcp.NewTool("generate_tests",
		mcp.WithDescription("Generate unit tests for a source file"),
		mcp.WithString("file_path", mcp.Required(), mcp.Description("Path to the file to generate tests for")),
	)
	s.mcpServer.AddTool(testsTool, s.handleGenerateTests)

	// generate_commit_msg tool
	commitTool := mcp.NewTool("generate_commit_msg",
		mcp.WithDescription("Generate a conventional commit message from staged changes"),
	)
	s.mcpServer.AddTool(commitTool, s.handleGenerateCommitMsg)

	// run_recipe tool
	recipeTool := mcp.NewTool("run_recipe",
		mcp.WithDescription("Execute a named recipe from .grokkit/recipes"),
		mcp.WithString("name", mcp.Required(), mcp.Description("Name of the recipe to run")),
		mcp.WithObject("params", mcp.Description("Parameters for the recipe (optional)")),
	)
	s.mcpServer.AddTool(recipeTool, s.handleRunRecipe)

	logger.Info("MCP tools registered", "count", 6)
	return nil
}

// handleLintCode implements the lint_code MCP tool.
func (s *Server) handleLintCode(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	filePath, err := req.RequireString("file_path")
	if err != nil {
		return mcp.NewToolResultText("Error: file_path is required"), nil
	}

	logger.Info("MCP lint_code called", "file", filePath)

	result, err := linter.LintFile(filePath)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Linting failed: %v", err)), nil
	}

	output := fmt.Sprintf("✅ Lint completed for %s\n", filePath)
	output += fmt.Sprintf("Language: %s\n", result.Language)
	output += fmt.Sprintf("Linter: %s\n", result.LinterUsed)
	output += fmt.Sprintf("Has issues: %v\n", result.HasIssues)

	if result.Output != "" {
		output += "\n--- Linter Output ---\n" + result.Output
	}

	return mcp.NewToolResultText(output), nil
}

// handleAnalyzeCode implements the analyze_code MCP tool.
func (s *Server) handleAnalyzeCode(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	dir := "."
	if d := req.GetString("dir", ""); d != "" {
		dir = d
	}

	logger.Info("MCP analyze_code called", "dir", dir)

	// Fully functional implementation using logic from cmd/analyze.go
	files, err := analyze.DiscoverSourceFiles(dir)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Failed to discover files: %v", err)), nil
	}

	if len(files) == 0 {
		return mcp.NewToolResultText("No supported source files found in directory."), nil
	}

	lang := linter.DetectPrimaryLanguage(files)
	if lang == "" {
		lang = "unknown"
	}

	_, promptContent, err := prompts.LoadAnalysisPrompt(dir, lang)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Failed to load analysis prompt: %v", err)), nil
	}

	context := analyze.BuildProjectContext(dir, files)

	messages := []map[string]string{
		{"role": "system", "content": promptContent},
		{"role": "user", "content": fmt.Sprintf("Analyze this %s project and generate the full educational Markdown report now:\n\n%s", lang, context)},
	}

	report := s.grokClient.StreamSilent(messages, "grok-4")

	result := fmt.Sprintf("✅ Analysis complete for %s\n", dir)
	result += fmt.Sprintf("Files analyzed: %d\n", len(files))
	result += fmt.Sprintf("Primary language: %s\n\n", lang)
	result += report

	return mcp.NewToolResultText(result), nil
}

// handleGenerateDocs implements the generate_docs MCP tool.
func (s *Server) handleGenerateDocs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	filePath, err := req.RequireString("file_path")
	if err != nil {
		return mcp.NewToolResultText("Error: file_path is required"), nil
	}

	style := req.GetString("style", "")

	logger.Info("MCP generate_docs called", "file", filePath, "style", style)

	// Fully functional implementation using logic from cmd/docs.go
	lang, err := linter.DetectLanguage(filePath)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Unsupported language: %v", err)), nil
	}

	content, err := os.ReadFile(filePath)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Failed to read file: %v", err)), nil
	}

	messages := docs.BuildDocsMessages(lang.Name, string(content))
	response := s.grokClient.StreamSilent(messages, "grok-4")

	documented := grok.CleanCodeResponse(response)

	result := fmt.Sprintf("✅ Documentation generated for %s\n", filePath)
	result += fmt.Sprintf("Language: %s\n", lang.Name)
	result += "\n--- Documented Code ---\n"
	result += documented

	return mcp.NewToolResultText(result), nil
}

// handleGenerateTests implements the generate_tests MCP tool.
func (s *Server) handleGenerateTests(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	filePath, err := req.RequireString("file_path")
	if err != nil {
		return mcp.NewToolResultText("Error: file_path is required"), nil
	}

	logger.Info("MCP generate_tests called", "file", filePath)

	// Fully functional using logic from cmd/testgen.go
	langObj, err := linter.DetectLanguage(filePath)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Unsupported language: %v", err)), nil
	}

	lang := langObj.Name
	if lang == "C/C++" {
		ext := strings.ToLower(filepath.Ext(filePath))
		if ext == ".cpp" || ext == ".cc" || ext == ".cxx" {
			lang = "C++"
		} else {
			lang = "C"
		}
	}

	prompt := testgen.GetTestPrompt(lang)
	if prompt == "" {
		return mcp.NewToolResultText("No test prompt available for this language"), nil
	}

	content, err := os.ReadFile(filePath)
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Failed to read file: %v", err)), nil
	}

	messages := []map[string]string{
		{"role": "system", "content": prompt},
		{"role": "user", "content": fmt.Sprintf("Generate comprehensive tests for the following %s code:\n\n%s", lang, string(content))},
	}

	response := s.grokClient.StreamSilent(messages, "grok-4")
	cleaned := grok.CleanCodeResponse(response)

	result := fmt.Sprintf("✅ Tests generated for %s\n", filePath)
	result += fmt.Sprintf("Language: %s\n\n", lang)
	result += "--- Generated Test Code ---\n"
	result += cleaned

	return mcp.NewToolResultText(result), nil
}

// handleGenerateCommitMsg implements the generate_commit_msg MCP tool.
func (s *Server) handleGenerateCommitMsg(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	logger.Info("MCP generate_commit_msg called")

	// Fully functional implementation using git diff + Grok
	if !git.IsRepo() {
		return mcp.NewToolResultText("Error: Not in a git repository"), nil
	}

	diff, err := git.Run([]string{"diff", "--cached"})
	if err != nil || strings.TrimSpace(diff) == "" {
		return mcp.NewToolResultText("No staged changes found. Stage some changes with `git add` first."), nil
	}

	prompt := `You are an expert at writing conventional commit messages.

Generate a single, concise conventional commit message for the following git diff.

Rules:
- Use format: <type>(<scope>): <description>
- Types: feat, fix, docs, style, refactor, perf, test, chore, ci, build
- Keep subject line under 72 characters
- Do not include body unless absolutely necessary
- Return ONLY the commit message, no explanations, no markdown, no quotes.

Diff:
` + diff

	messages := []map[string]string{
		{"role": "system", "content": prompt},
		{"role": "user", "content": "Generate the commit message now."},
	}

	response := s.grokClient.StreamSilent(messages, "grok-4")
	cleaned := strings.TrimSpace(grok.CleanCodeResponse(response))

	result := "✅ Conventional commit message generated:\n\n"
	result += cleaned

	return mcp.NewToolResultText(result), nil
}

// handleRunRecipe implements the run_recipe MCP tool.
func (s *Server) handleRunRecipe(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	name, err := req.RequireString("name")
	if err != nil {
		return mcp.NewToolResultText("Error: name is required"), nil
	}

	logger.Info("MCP run_recipe called", "recipe", name)

	// Fully functional using internal/recipe package
	r, err := recipe.Load(name, nil) // no custom params for now
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Failed to load recipe '%s': %v", name, err)), nil
	}

	client := grok.NewClient()
	runner := recipe.NewRunner(r, client, "grok-4")

	err = runner.Run()
	if err != nil {
		return mcp.NewToolResultText(fmt.Sprintf("Recipe execution failed: %v", err)), nil
	}

	result := fmt.Sprintf("✅ Recipe '%s' executed successfully\n", name)
	result += "Check the terminal output above for full details.\n"

	return mcp.NewToolResultText(result), nil
}

// registerResources registers MCP resources (recipes and prompts).
func (s *Server) registerResources() {
	// Recipes as resources
	s.mcpServer.AddResource(
		mcp.NewResource("recipes://local", "Local Recipes",
			mcp.WithResourceDescription("Project-local workflow recipes from .grokkit/recipes/"),
			mcp.WithMIMEType("application/json"),
		),
		s.handleRecipesLocal,
	)

	s.mcpServer.AddResource(
		mcp.NewResource("recipes://global", "Global Recipes",
			mcp.WithResourceDescription("Global recipes from ~/.config/grokkit/recipes/"),
			mcp.WithMIMEType("application/json"),
		),
		s.handleRecipesGlobal,
	)

	// Prompts as resources
	s.mcpServer.AddResource(
		mcp.NewResource("prompts://language", "Language Prompts",
			mcp.WithResourceDescription("Language-specific analysis prompts"),
			mcp.WithMIMEType("text/markdown"),
		),
		s.handlePrompts,
	)

	logger.Info("MCP resources registered", "count", 3)
}

// Placeholder handlers for resources (to be implemented)
func (s *Server) handleRecipesLocal(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
	return []mcp.ResourceContents{
		mcp.TextResourceContents{
			URI:      "recipes://local",
			MIMEType: "application/json",
			Text:     `{"status": "placeholder - local recipes not yet implemented"}`,
		},
	}, nil
}

func (s *Server) handleRecipesGlobal(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
	return []mcp.ResourceContents{
		mcp.TextResourceContents{
			URI:      "recipes://global",
			MIMEType: "application/json",
			Text:     `{"status": "placeholder - global recipes not yet implemented"}`,
		},
	}, nil
}

func (s *Server) handlePrompts(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
	return []mcp.ResourceContents{
		mcp.TextResourceContents{
			URI:      "prompts://language",
			MIMEType: "text/markdown",
			Text:     `# Language Prompts\n\nPlaceholder - language prompts not yet implemented.`,
		},
	}, nil
}

// Run starts the MCP server using stdio transport (for Claude Code, etc.).
func (s *Server) Run(ctx context.Context) error {
	logger.Info("Starting Grokkit MCP server (stdio transport)")
	return server.ServeStdio(s.mcpServer)
}