MuyueWorkspace/internal/mcp/registry.go

package mcp

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v3"
)

type RegistryServer struct {
	Name        string            `yaml:"name" json:"name"`
	Description string            `yaml:"description" json:"description"`
	Category    string            `yaml:"category" json:"category"`
	Package     string            `yaml:"package" json:"package"`
	Command     string            `yaml:"command" json:"command"`
	Args        []string          `yaml:"args" json:"args"`
	Env         map[string]string `yaml:"env,omitempty" json:"env,omitempty"`
	RequiredEnv []string          `yaml:"required_env,omitempty" json:"required_env,omitempty"`
	HomePage    string            `yaml:"homepage,omitempty" json:"homepage,omitempty"`
	Tags        []string          `yaml:"tags,omitempty" json:"tags,omitempty"`
	Version     string            `yaml:"version,omitempty" json:"version,omitempty"`
	InstallType string            `yaml:"install_type" json:"install_type"`
}

type Registry struct {
	SchemaVersion string           `yaml:"schema_version"`
	UpdatedAt     time.Time        `yaml:"updated_at"`
	Servers       []RegistryServer `yaml:"servers"`
}

type MCPStatus struct {
	Name      string `json:"name"`
	Installed bool   `json:"installed"`
	Running   bool   `json:"running"`
	Healthy   bool   `json:"healthy"`
	Version   string `json:"version"`
	Error     string `json:"error,omitempty"`
}

type EditorConfig struct {
	Name             string
	ConfigPath       string
	ConfigKey        string
	LocalConfigPath  string
	Format           string
	TransformCommand func(entry mcpEntry) interface{}
}

var (
	registryMu     sync.RWMutex
	registryCache  *Registry
	registryPath   string
)

func init() {
	home, _ := os.UserHomeDir()
	if home != "" {
		registryPath = filepath.Join(home, ".muyue", "mcp-registry.yaml")
	}
}

func SetRegistryPath(p string) {
	registryMu.Lock()
	defer registryMu.Unlock()
	registryPath = p
	registryCache = nil
}

func DefaultRegistry() *Registry {
	return &Registry{
		SchemaVersion: "v1",
		UpdatedAt:     time.Now(),
		Servers: []RegistryServer{
			{
				Name: "filesystem", Description: "File system operations for AI tools",
				Category: "core", Package: "@modelcontextprotocol/server-filesystem",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-filesystem"},
				InstallType: "npm", Tags: []string{"files", "core"},
			},
			{
				Name: "github", Description: "GitHub API integration",
				Category: "vcs", Package: "@modelcontextprotocol/server-github",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-github"},
				Env: map[string]string{"GITHUB_PERSONAL_ACCESS_TOKEN": ""},
				RequiredEnv: []string{"GITHUB_PERSONAL_ACCESS_TOKEN"},
				InstallType: "npm", Tags: []string{"github", "git"},
			},
			{
				Name: "git", Description: "Git repository operations",
				Category: "vcs", Package: "@modelcontextprotocol/server-git",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-git"},
				InstallType: "npm", Tags: []string{"git"},
			},
			{
				Name: "fetch", Description: "Web fetching and HTTP requests",
				Category: "web", Package: "@modelcontextprotocol/server-fetch",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-fetch"},
				InstallType: "npm", Tags: []string{"web", "http"},
			},
			{
				Name: "memory", Description: "Persistent memory/knowledge graph",
				Category: "core", Package: "@modelcontextprotocol/server-memory",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-memory"},
				InstallType: "npm", Tags: []string{"memory", "core"},
			},
			{
				Name: "sequential-thinking", Description: "Structured reasoning and chain-of-thought",
				Category: "ai", Package: "@modelcontextprotocol/server-sequential-thinking",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-sequential-thinking"},
				InstallType: "npm", Tags: []string{"ai", "reasoning"},
			},
			{
				Name: "brave-search", Description: "Web search via Brave Search API",
				Category: "web", Package: "@modelcontextprotocol/server-brave-search",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-brave-search"},
				Env: map[string]string{"BRAVE_API_KEY": ""},
				RequiredEnv: []string{"BRAVE_API_KEY"},
				InstallType: "npm", Tags: []string{"search", "web"},
			},
			{
				Name: "sqlite", Description: "SQLite database operations",
				Category: "database", Package: "@modelcontextprotocol/server-sqlite",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-sqlite"},
				InstallType: "npm", Tags: []string{"database", "sqlite"},
			},
			{
				Name: "postgres", Description: "PostgreSQL database operations",
				Category: "database", Package: "@modelcontextprotocol/server-postgres",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-postgres"},
				InstallType: "npm", Tags: []string{"database", "postgres"},
			},
			{
				Name: "docker", Description: "Docker container management",
				Category: "devops", Package: "@modelcontextprotocol/server-docker",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-docker"},
				InstallType: "npm", Tags: []string{"docker", "devops"},
			},
			{
				Name: "minimax-web-search", Description: "Web search via MiniMax API",
				Category: "ai", Package: "@minimax/mcp-web-search",
				Command: "npx", Args: []string{"-y", "@minimax/mcp-web-search"},
				Env: map[string]string{"MINIMAX_API_KEY": ""},
				RequiredEnv: []string{"MINIMAX_API_KEY"},
				InstallType: "npm", Tags: []string{"ai", "search"},
			},
			{
				Name: "minimax-image", Description: "Image understanding via MiniMax API",
				Category: "ai", Package: "@minimax/mcp-image-understanding",
				Command: "npx", Args: []string{"-y", "@minimax/mcp-image-understanding"},
				Env: map[string]string{"MINIMAX_API_KEY": ""},
				RequiredEnv: []string{"MINIMAX_API_KEY"},
				InstallType: "npm", Tags: []string{"ai", "image"},
			},
			{
				Name: "puppeteer", Description: "Browser automation with Puppeteer",
				Category: "web", Package: "@modelcontextprotocol/server-puppeteer",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-puppeteer"},
				InstallType: "npm", Tags: []string{"browser", "automation"},
			},
			{
				Name: "everything", Description: "Test/debug MCP server with all features",
				Category: "testing", Package: "@modelcontextprotocol/server-everything",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-everything"},
				InstallType: "npm", Tags: []string{"testing", "debug"},
			},
			{
				Name: "slack", Description: "Slack workspace integration",
				Category: "communication", Package: "@modelcontextprotocol/server-slack",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-slack"},
				Env: map[string]string{"SLACK_BOT_TOKEN": ""},
				RequiredEnv: []string{"SLACK_BOT_TOKEN"},
				InstallType: "npm", Tags: []string{"slack", "communication"},
			},
			{
				Name: "google-maps", Description: "Google Maps integration",
				Category: "web", Package: "@modelcontextprotocol/server-google-maps",
				Command: "npx", Args: []string{"-y", "@modelcontextprotocol/server-google-maps"},
				Env: map[string]string{"GOOGLE_MAPS_API_KEY": ""},
				RequiredEnv: []string{"GOOGLE_MAPS_API_KEY"},
				InstallType: "npm", Tags: []string{"maps", "location"},
			},
		},
	}
}

func LoadRegistry() (*Registry, error) {
	registryMu.RLock()
	if registryCache != nil {
		defer registryMu.RUnlock()
		return registryCache, nil
	}
	registryMu.RUnlock()

	reg, err := loadRegistryFromDisk()
	if err != nil {
		defaultReg := DefaultRegistry()
		registryMu.Lock()
		registryCache = defaultReg
		registryMu.Unlock()
		return defaultReg, nil
	}

	registryMu.Lock()
	registryCache = reg
	registryMu.Unlock()
	return reg, nil
}

func loadRegistryFromDisk() (*Registry, error) {
	if registryPath == "" {
		return nil, fmt.Errorf("registry path not set")
	}

	data, err := os.ReadFile(registryPath)
	if err != nil {
		return nil, err
	}

	var reg Registry
	if err := yaml.Unmarshal(data, &reg); err != nil {
		return nil, fmt.Errorf("parse registry: %w", err)
	}

	return &reg, nil
}

func SaveRegistry(reg *Registry) error {
	if registryPath == "" {
		return fmt.Errorf("registry path not set")
	}

	reg.UpdatedAt = time.Now()
	data, err := yaml.Marshal(reg)
	if err != nil {
		return fmt.Errorf("marshal registry: %w", err)
	}

	if err := os.MkdirAll(filepath.Dir(registryPath), 0755); err != nil {
		return err
	}

	if err := os.WriteFile(registryPath, data, 0644); err != nil {
		return err
	}

	registryMu.Lock()
	registryCache = reg
	registryMu.Unlock()
	return nil
}

func AddToRegistry(server RegistryServer) error {
	reg, err := LoadRegistry()
	if err != nil {
		return err
	}

	for _, s := range reg.Servers {
		if s.Name == server.Name {
			return fmt.Errorf("server %q already exists in registry", server.Name)
		}
	}

	reg.Servers = append(reg.Servers, server)
	return SaveRegistry(reg)
}

func RemoveFromRegistry(name string) error {
	reg, err := LoadRegistry()
	if err != nil {
		return err
	}

	for i, s := range reg.Servers {
		if s.Name == name {
			reg.Servers = append(reg.Servers[:i], reg.Servers[i+1:]...)
			return SaveRegistry(reg)
		}
	}

	return fmt.Errorf("server %q not found in registry", name)
}

func InitRegistry() error {
	if _, err := os.Stat(registryPath); err == nil {
		return nil
	}
	return SaveRegistry(DefaultRegistry())
}

func ResolveEnv(env map[string]string, providerKeys map[string]string) map[string]string {
	resolved := make(map[string]string)
	for k, v := range env {
		if v != "" {
			resolved[k] = v
			continue
		}

		if providerKeys != nil {
			for providerKey, apiKey := range providerKeys {
				if strings.EqualFold(k, providerKey) || strings.Contains(strings.ToUpper(k), strings.ToUpper(providerKey)) {
					if apiKey != "" {
						resolved[k] = apiKey
					}
				}
			}
		}

		if resolved[k] == "" {
			if envVal := os.Getenv(k); envVal != "" {
				resolved[k] = envVal
			}
		}
	}
	return resolved
}

func ValidateConfig(configPath string) error {
	data, err := os.ReadFile(configPath)
	if err != nil {
		return fmt.Errorf("read config: %w", err)
	}

	var cfg map[string]interface{}
	if err := json.Unmarshal(data, &cfg); err != nil {
		return fmt.Errorf("parse config: %w", err)
	}

	return nil
}

func DiscoverNpmServers() ([]RegistryServer, error) {
	var servers []RegistryServer

	packages := []struct {
		pkg   string
		name  string
		desc  string
		cat   string
		args  []string
	}{
		{"@modelcontextprotocol/server-filesystem", "filesystem", "File system operations", "core", []string{"-y", "@modelcontextprotocol/server-filesystem"}},
		{"@modelcontextprotocol/server-github", "github", "GitHub API integration", "vcs", []string{"-y", "@modelcontextprotocol/server-github"}},
		{"@modelcontextprotocol/server-fetch", "fetch", "Web fetching", "web", []string{"-y", "@modelcontextprotocol/server-fetch"}},
		{"@modelcontextprotocol/server-memory", "memory", "Persistent memory", "core", []string{"-y", "@modelcontextprotocol/server-memory"}},
	}

	for _, p := range packages {
		servers = append(servers, RegistryServer{
			Name:        p.name,
			Description: p.desc,
			Category:    p.cat,
			Package:     p.pkg,
			Command:     "npx",
			Args:        p.args,
			InstallType: "npm",
		})
	}

	return servers, nil
}

func GetInstalledVersion(name string) string {
	home, _ := os.UserHomeDir()
	if home == "" {
		return ""
	}
	receiptPath := filepath.Join(home, ".muyue", "receipts", "mcp", name+".json")
	data, err := os.ReadFile(receiptPath)
	if err != nil {
		return ""
	}
	var receipt struct {
		Version string `json:"version"`
	}
	if json.Unmarshal(data, &receipt) == nil {
		return receipt.Version
	}
	return ""
}

func SaveReceipt(name, version string) error {
	home, _ := os.UserHomeDir()
	if home == "" {
		return nil
	}
	receiptDir := filepath.Join(home, ".muyue", "receipts", "mcp")
	os.MkdirAll(receiptDir, 0755)

	receipt := struct {
		Name      string `json:"name"`
		Version   string `json:"version"`
		UpdatedAt string `json:"updated_at"`
	}{
		Name:      name,
		Version:   version,
		UpdatedAt: time.Now().Format(time.RFC3339),
	}

	data, _ := json.MarshalIndent(receipt, "", "  ")
	return os.WriteFile(filepath.Join(receiptDir, name+".json"), data, 0644)
}

func BuildProviderKeyMap(cfg interface{ GetAPIKeys() map[string]string }) map[string]string {
	if cfg == nil {
		return nil
	}
	return cfg.GetAPIKeys()
}

func EditorConfigs(homeDir string) []EditorConfig {
	if homeDir == "" {
		home, _ := os.UserHomeDir()
		homeDir = home
	}

	transformStdio := func(e mcpEntry) interface{} {
		m := map[string]interface{}{
			"command": e.cmd,
			"args":    e.args,
		}
		if len(e.env) > 0 {
			m["env"] = e.env
		}
		return m
	}

	transformCursor := func(e mcpEntry) interface{} {
		m := map[string]interface{}{
			"type":    "stdio",
			"command": e.cmd,
			"args":    e.args,
		}
		if len(e.env) > 0 {
			m["env"] = e.env
		}
		return m
	}

	return []EditorConfig{
		{
			Name: "crush", ConfigPath: filepath.Join(homeDir, ".config", "crush", "crush.json"),
			ConfigKey: "mcps", Format: "json", TransformCommand: transformStdio,
		},
		{
			Name: "claude-code", ConfigPath: filepath.Join(homeDir, ".claude.json"),
			ConfigKey: "mcpServers", Format: "json", TransformCommand: transformStdio,
		},
		{
			Name: "cursor", ConfigPath: filepath.Join(homeDir, ".cursor", "mcp.json"),
			LocalConfigPath: ".cursor/mcp.json", ConfigKey: "mcpServers",
			Format: "json", TransformCommand: transformCursor,
		},
		{
			Name: "vscode", ConfigPath: filepath.Join(homeDir, ".vscode", "mcp.json"),
			LocalConfigPath: ".vscode/mcp.json", ConfigKey: "servers",
			Format: "json", TransformCommand: transformStdio,
		},
		{
			Name: "windsurf", ConfigPath: filepath.Join(homeDir, ".windsurf", "mcp.json"),
			ConfigKey: "mcpServers", Format: "json", TransformCommand: transformStdio,
		},
	}
}

func CheckServerStatus(name string) MCPStatus {
	status := MCPStatus{Name: name}

	reg, err := LoadRegistry()
	if err != nil {
		status.Error = "registry unavailable"
		return status
	}

	var server *RegistryServer
	for i := range reg.Servers {
		if reg.Servers[i].Name == name {
			server = &reg.Servers[i]
			break
		}
	}
	if server == nil {
		status.Error = "not in registry"
		return status
	}

	_, err = exec.LookPath(server.Command)
	if err != nil {
		status.Error = fmt.Sprintf("command %q not found", server.Command)
		return status
	}
	status.Installed = true

	status.Version = GetInstalledVersion(name)

	home, _ := os.UserHomeDir()
	if home != "" {
		crushingPath := filepath.Join(home, ".config", "crush", "crush.json")
		data, err := os.ReadFile(crushingPath)
		if err == nil {
			var cfg map[string]interface{}
			if json.Unmarshal(data, &cfg) == nil {
				if mcps, ok := cfg["mcps"].(map[string]interface{}); ok {
					if _, exists := mcps[name]; exists {
						status.Running = true
						status.Healthy = true
					}
				}
			}
		}
	}

	return status
}